From 0d8acdad867a0a7f32355a5518acc387604a804e Mon Sep 17 00:00:00 2001
From: Tropfchen <6117944+Tropfchen@users.noreply.github.com>
Date: Sat, 20 Jan 2024 12:46:04 +0100
Subject: [PATCH] add resolution printout

---
 .gitignore       |  5 +++++
 js/quickNodes.js | 38 ++++++++++++++++++++++++++++++++++++--
 nodes.py         | 24 ++++++++++++++++++++----
 3 files changed, 61 insertions(+), 6 deletions(-)

diff --git a/.gitignore b/.gitignore
index 68bc17f..a7d2512 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,8 @@
+# pre-commit
+.flake8
+.pre-commit-config.yaml
+pyproject.toml
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]
diff --git a/js/quickNodes.js b/js/quickNodes.js
index d1ee0ab..013faba 100644
--- a/js/quickNodes.js
+++ b/js/quickNodes.js
@@ -1,4 +1,5 @@
 import { app } from "../../scripts/app.js";
+import { ComfyWidgets } from "/scripts/widgets.js";
 
 // ----------- ComfyUI\web\extensions\core\widgetInputs.js copypaste -----------
 const CONVERTED_TYPE = "converted-widget";
@@ -139,9 +140,42 @@ function prependNewNode(node, nodeData, newNodeName, newNodewidgetNames) {
 app.registerExtension({
   name: "trop.YARS",
   async beforeRegisterNodeDef(nodeType, nodeData, app) {
-    if (
+    // add resolution printout
+    if (nodeData.name === "YARSAdv" || nodeData.name === "YARS") {
+      const onExecuted = nodeType.prototype.onExecuted;
+      nodeType.prototype.onExecuted = function (message) {
+        const r = onExecuted?.apply?.(this, arguments);
+
+        let w = this.widgets.find((w) => w.name === "resolution_printout");
+        if (!w) {
+          w = ComfyWidgets["STRING"](
+            this,
+            "resolution_printout",
+            ["STRING", { multiline: true }],
+            app,
+          ).widget;
+          w.inputEl.readOnly = true;
+          w.inputEl.style.opacity = 0.8;
+        }
+
+        const width = message.width[0];
+        const height = message.height[0];
+        const mpx = ((width * height) / 1048576).toFixed(2);
+        const res_msg = `resolution: ${width}x${height} (~${mpx} Mpx)`;
+        const ratio_msg = `ratio: ~${message.ratio[0].toFixed(2)}`;
+        w.value = `${res_msg}\n${ratio_msg}`;
+
+        this.onResize?.(this.size);
+
+        return r;
+      };
+    }
+
+    // add quick nodes
+    else if (
       nodeData.name === "EmptyLatentImage" ||
-      nodeData.name === "ImageScale"
+      nodeData.name === "ImageScale" ||
+      nodeData.name === "LatentUpscale"
     ) {
       addMenuHandler(nodeType, function (_, options) {
         const entries = [
diff --git a/nodes.py b/nodes.py
index 70efa2c..416ac4e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -106,12 +106,20 @@ class YARS:
         base_resolution: int,
         aspect_ratio: str,
         overextend: bool,
-    ) -> tuple[int, int]:
+    ):
         if m := re.search(r"(\d+):(\d+)", aspect_ratio):
             ratio: float = int(m.group(2)) / int(m.group(1))
             d = calculate_aspect_ratio(base_resolution, ratio, overextend)
 
-            return (d.width, d.height)
+            # return as dict with `ui` key to trigger onExecuted
+            return {
+                "ui": {
+                    "width": [d.width],
+                    "height": [d.height],
+                    "ratio": [ratio],
+                },
+                "result": (d.width, d.height),
+            }
 
         raise ValueError(f"Could't find aspect ratio in string `{aspect_ratio}`")
 
@@ -175,7 +183,7 @@ class YARSAdv:
         height_ratio: int,
         overextend: bool,
         constant_resolution: bool,
-    ) -> tuple[int, int]:
+    ):
         ratio: float = height_ratio / width_ratio
 
         if constant_resolution:
@@ -183,7 +191,15 @@ class YARSAdv:
         else:
             d = calculate_aspect_ratio(base_resolution, ratio, overextend)
 
-        return (d.width, d.height)
+        # return as dict with `ui` key to trigger onExecuted
+        return {
+            "ui": {
+                "width": [d.width],
+                "height": [d.height],
+                "ratio": [width_ratio / height_ratio],
+            },
+            "result": (d.width, d.height),
+        }
 
 
 if __name__ == "__main__":