From 6d8f4b46e57025ad0e6db112c22fa7f3049add03 Mon Sep 17 00:00:00 2001 From: nkchocoai Date: Sat, 17 Feb 2024 17:42:29 +0900 Subject: [PATCH] Support LatentUpscale --- py/capture.py | 8 ++++---- py/defs/captures.py | 13 ++++++++++++- py/defs/ext/size_from_presets.py | 4 ++-- py/defs/formatters.py | 12 +++++++++++- py/nodes/node.py | 5 +++-- 5 files changed, 32 insertions(+), 10 deletions(-) diff --git a/py/capture.py b/py/capture.py index 0616862..59e2091 100644 --- a/py/capture.py +++ b/py/capture.py @@ -18,7 +18,7 @@ class Capture: class_type = obj["class_type"] obj_class = NODE_CLASS_MAPPINGS[class_type] node_inputs = prompt[node_id]["inputs"] - input_data_all = get_input_data( + input_data = get_input_data( node_inputs, obj_class, node_id, outputs, prompt, extra_data ) for node_class, metas in CAPTURE_FIELD_LIST.items(): @@ -26,7 +26,7 @@ class Capture: for meta, field_data in metas.items(): validate = field_data.get("validate") if validate is not None and not validate( - node_id, obj, prompt, extra_data, outputs, input_data_all + node_id, obj, prompt, extra_data, outputs, input_data ): print("validate failed", meta, node_id, obj) continue @@ -34,14 +34,14 @@ class Capture: if meta not in inputs: inputs[meta] = [] field_name = field_data["field_name"] - value = input_data_all.get(field_name) + value = input_data.get(field_name) if value is not None: format = field_data.get("format") v = value if isinstance(value, list) and len(value) > 0: v = value[0] if format is not None: - v = format(v) + v = format(v, input_data) inputs[meta].append((node_id, v)) return inputs diff --git a/py/defs/captures.py b/py/defs/captures.py index bfbdc87..41cf3c3 100644 --- a/py/defs/captures.py +++ b/py/defs/captures.py @@ -1,6 +1,6 @@ from .meta import MetaField from .validators import is_positive_prompt, is_negative_prompt -from .formatters import calc_model_hash +from .formatters import calc_model_hash, get_scaled_width, get_scaled_height CAPTURE_FIELD_LIST = { @@ -36,4 +36,15 @@ CAPTURE_FIELD_LIST = { MetaField.SAMPLER_NAME: {"field_name": "sampler_name"}, MetaField.SCHEDULER: {"field_name": "scheduler"}, }, + "LatentUpscale": { + MetaField.IMAGE_WIDTH: {"field_name": "width"}, + MetaField.IMAGE_HEIGHT: {"field_name": "height"}, + }, + "LatentUpscaleBy": { + MetaField.IMAGE_WIDTH: {"field_name": "scale_by", "format": get_scaled_width}, + MetaField.IMAGE_HEIGHT: { + "field_name": "scale_by", + "format": get_scaled_height, + }, + }, } diff --git a/py/defs/ext/size_from_presets.py b/py/defs/ext/size_from_presets.py index f667a5b..af5b68a 100644 --- a/py/defs/ext/size_from_presets.py +++ b/py/defs/ext/size_from_presets.py @@ -1,11 +1,11 @@ from ..meta import MetaField -def get_width(preset): +def get_width(preset, input_data): return preset.split("x")[0].strip() -def get_height(preset): +def get_height(preset, input_data): return preset.split("x")[1].strip() diff --git a/py/defs/formatters.py b/py/defs/formatters.py index 1d80d34..d85cc57 100644 --- a/py/defs/formatters.py +++ b/py/defs/formatters.py @@ -5,7 +5,7 @@ import folder_paths cache_model_hash = {} -def calc_model_hash(model_name): +def calc_model_hash(model_name, input_data): if model_name in cache_model_hash: return cache_model_hash[model_name] @@ -19,3 +19,13 @@ def calc_model_hash(model_name): cache_model_hash[model_name] = model_hash return model_hash + + +def get_scaled_width(scaled_by, input_data): + samples = input_data["samples"][0]["samples"] + return round(samples.shape[3] * scaled_by * 8) + + +def get_scaled_height(scaled_by, input_data): + samples = input_data["samples"][0]["samples"] + return round(samples.shape[2] * scaled_by * 8) diff --git a/py/nodes/node.py b/py/nodes/node.py index 1f9a832..7b58dab 100644 --- a/py/nodes/node.py +++ b/py/nodes/node.py @@ -69,7 +69,7 @@ class SaveImageWithMetaData(BaseNode): i = 255.0 * image.cpu().numpy() img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - pnginfo_dict = self.gen_pnginfo() + pnginfo_dict = self.gen_pnginfo(sampler_find_method, sampler_find_node_id) if len(images) >= 2: pnginfo_dict["Batch index"] = index pnginfo_dict["Batch size"] = len(images) @@ -101,7 +101,7 @@ class SaveImageWithMetaData(BaseNode): return {"ui": {"images": results}} @classmethod - def gen_pnginfo(cls): + def gen_pnginfo(cls, sampler_find_method, sampler_find_node_id): # get all node inputs inputs = Capture.get_inputs() @@ -118,6 +118,7 @@ class SaveImageWithMetaData(BaseNode): filtered_inputs = Trace.filter_inputs_by_trace_tree( inputs, trace_tree_from_sampler_node ) + print(filtered_inputs) # generate PNGInfo from inputs pnginfo_dict = Capture.gen_pnginfo_dict(filtered_inputs)