129 lines
4.6 KiB
Python
129 lines
4.6 KiB
Python
import os
|
|
import time
|
|
from typing import TypeVar, Generic, TYPE_CHECKING
|
|
|
|
from modules import shared
|
|
from modules.processing import StableDiffusionProcessing, Processed
|
|
|
|
from scripts.lib.extractor import ExtractorBase
|
|
from scripts.lib.features.featureinfo import MultiImageFeatures
|
|
from scripts.lib.ui import retrieve_layers, retrieve_steps
|
|
from scripts.lib.putils import ProcessedBuilder
|
|
from scripts.lib.report import message as E
|
|
from scripts.lib.utils import sorted_values, sorted_items
|
|
from scripts.lib.colorizer import Colorizer
|
|
|
|
if TYPE_CHECKING:
|
|
from scripts.dumpunet import Script
|
|
|
|
TInfo = TypeVar("TInfo")
|
|
|
|
class FeatureExtractorBase(Generic[TInfo], ExtractorBase):
|
|
|
|
# image_index -> step -> Features
|
|
extracted_features: MultiImageFeatures[TInfo]
|
|
|
|
# steps to process
|
|
steps: list[int]
|
|
|
|
# layers to process
|
|
layers: list[str]
|
|
|
|
# dump path
|
|
path: str|None
|
|
|
|
def __init__(
|
|
self,
|
|
runner: "Script",
|
|
enabled: bool,
|
|
total_steps: int,
|
|
layer_input: str,
|
|
step_input: str,
|
|
path: str|None
|
|
):
|
|
super().__init__(runner, enabled)
|
|
|
|
self.extracted_features = MultiImageFeatures()
|
|
self.steps = []
|
|
self.layers = []
|
|
self.path = None
|
|
|
|
if not self.enabled:
|
|
return
|
|
|
|
assert layer_input is not None and layer_input != "", E("<Layers> must not be empty.")
|
|
|
|
self.layers = retrieve_layers(layer_input)
|
|
self.steps = (
|
|
retrieve_steps(step_input)
|
|
or list(range(1, total_steps+1))
|
|
)
|
|
|
|
if path is not None:
|
|
assert path != "", E("<Output path> must not be empty.")
|
|
# mkdir -p path
|
|
if os.path.exists(path):
|
|
assert os.path.isdir(path), E("<Output path> already exists and is not a directory.")
|
|
else:
|
|
os.makedirs(path, exist_ok=True)
|
|
|
|
self.path = path
|
|
|
|
def on_setup(self):
|
|
self.extracted_features = MultiImageFeatures()
|
|
|
|
def add_images(
|
|
self,
|
|
p: StableDiffusionProcessing,
|
|
builder: ProcessedBuilder,
|
|
extracted_features: MultiImageFeatures[TInfo],
|
|
add_average: bool,
|
|
color: Colorizer
|
|
):
|
|
|
|
if not self.enabled:
|
|
return
|
|
|
|
if shared.state.interrupted:
|
|
return
|
|
|
|
sorted_step_features = list(sorted_values(extracted_features))
|
|
assert len(builder.items) == len(sorted_step_features), E(f"#images={len(builder.items)}, #features={len(sorted_step_features)}")
|
|
|
|
t0 = int(time.time()) # for binary files' name
|
|
shared.total_tqdm.clear()
|
|
shared.total_tqdm.updateTotal(len(sorted_step_features) * len(self.steps) * len(self.layers))
|
|
|
|
for idx, step_features in enumerate(sorted_step_features):
|
|
for step, features in sorted_items(step_features):
|
|
for layer, feature in features:
|
|
|
|
if shared.state.interrupted:
|
|
break
|
|
|
|
canvases = self.feature_to_grid_images(feature, layer, idx, step, p.width, p.height, add_average, color)
|
|
for canvas in canvases:
|
|
builder.add_ref(idx, canvas, None, {"Layer Name": layer, "Feature Steps": step})
|
|
|
|
if self.path is not None:
|
|
basename = f"{idx:03}-{layer}-{step:03}-{{ch:04}}-{t0}"
|
|
self.save_features(feature, layer, idx, step, p.width, p.height, self.path, basename)
|
|
|
|
if hasattr(shared.total_tqdm, "_tqdm"):
|
|
shared.total_tqdm._tqdm.set_postfix_str(layer.ljust(5)) # type: ignore
|
|
shared.total_tqdm.update()
|
|
|
|
def feature_to_grid_images(self, feature: TInfo, layer: str, img_idx: int, step: int, width: int, height: int, add_average: bool, color: Colorizer):
|
|
raise NotImplementedError(f"{self.__class__.__name__}.feature_to_grid_images")
|
|
|
|
def save_features(self, feature: TInfo, layer: str, img_idx: int, step: int, width: int, height: int, path: str, basename: str):
|
|
raise NotImplementedError(f"{self.__class__.__name__}.save_features")
|
|
|
|
def _fixup(self, proc: Processed):
|
|
# For Dynamic Prompt Extension
|
|
# which is not append subseeds...
|
|
while len(proc.all_subseeds) < len(proc.all_seeds):
|
|
proc.all_subseeds.append(proc.all_subseeds[0] if 0 < len(proc.all_subseeds) else 0)
|
|
return proc
|
|
|