Support LatentUpscale

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
 

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,
+        },
+    },
 }

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()
 
 

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)

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)