add auraSR node (#63)

* add auraSR node * refine * add exception catch * add docs * refine * refine docs * add news
siliconflow · Jul 31, 2024 · 6e63ed3 · 6e63ed3
1 parent 6e741c8
commit 6e63ed3
Show file tree

Hide file tree

Showing 10 changed files with 283 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -1,5 +1,6 @@
 # BizyAir
 
+- [2024/07/31] 🌩️ The super-resolution node has been released, capable of enlarging images four times their original size. [BizyAir Photorealistic Image Super Resolution](https://siliconflow.github.io/BizyAir/others/index.html#bizyair-photorealistic-image-super-resolution)
 - [2024/07/25] 🌩️ Users can load BizyAir workflow examples directly by clicking the "☁️BizyAir Workflow Examples" button. [Example GIF](./docs/docs/getting-started/imgs/run-bizyair-examples.gif)
 - [2024/07/23] 🌩️ [BizyAir ChatGLM3 Text Encode](./examples/bizyair_showcase_run_with_local_nodes.json) node is released.
 - [2024/07/16] 🌩️ [BizyAir Controlnet Union SDXL 1.0](https://siliconflow.github.io/BizyAir/controlnet-union/introduce.html) node is released.

diff --git a/docs/docs/ksampler/introduce.md b/docs/docs/ksampler/introduce.md
diff --git a/docs/docs/others/imgs/super-resolution.png b/docs/docs/others/imgs/super-resolution.png
diff --git a/docs/docs/others/index.md b/docs/docs/others/index.md
@@ -1,3 +1,9 @@
+## ☁️BizyAir Photorealistic Image Super Resolution
+
+This node helps you enlarge the image 4 times, suitable for non-cartoon images.
+
+![](./imgs/super-resolution.png)
+
 ## ☁️BizyAir Anime Image Super Resolution
 
 This node helps you enhance the resolution of your input cartoon images. You can choose to upscale them by 2x or 4x.

diff --git a/examples/bizyair_showcase_realistic_superresolution.json b/examples/bizyair_showcase_realistic_superresolution.json
@@ -0,0 +1,169 @@
+{
+  "last_node_id": 4,
+  "last_link_id": 3,
+  "nodes": [
+    {
+      "id": 4,
+      "type": "PreviewImage",
+      "pos": [
+        895,
+        225
+      ],
+      "size": {
+        "0": 794.525390625,
+        "1": 737.365966796875
+      },
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 3
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
+    },
+    {
+      "id": 1,
+      "type": "BizyAirGenerateLightningImage",
+      "pos": [
+        547,
+        226
+      ],
+      "size": {
+        "0": 336,
+        "1": 220
+      },
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            1,
+            3
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "BizyAirGenerateLightningImage"
+      },
+      "widgets_values": [
+        "a dog",
+        5,
+        "fixed",
+        512,
+        512,
+        1.5,
+        1
+      ]
+    },
+    {
+      "id": 2,
+      "type": "BizyAirAuraSR",
+      "pos": [
+        489,
+        507
+      ],
+      "size": [
+        394.79998779296875,
+        26
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            2
+          ],
+          "shape": 3,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "BizyAirAuraSR"
+      }
+    },
+    {
+      "id": 3,
+      "type": "PreviewImage",
+      "pos": [
+        1725,
+        233
+      ],
+      "size": {
+        "0": 1438.3111572265625,
+        "1": 1453.90234375
+      },
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 2
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      2,
+      2,
+      0,
+      3,
+      0,
+      "IMAGE"
+    ],
+    [
+      3,
+      1,
+      0,
+      4,
+      0,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.42409761837248505,
+      "offset": [
+        1360.4125112955187,
+        349.01045420727985
+      ]
+    }
+  },
+  "version": 0.4
+}
diff --git a/image_utils.py b/image_utils.py
@@ -307,21 +307,22 @@ def decode_comfy_image(img_data: Union[List, str], image_format="png") -> torch.
     return output
 
 
-def tensor_to_base64(tensor: torch.Tensor) -> str:
+def tensor_to_base64(tensor: torch.Tensor, compress=True) -> str:
     tensor_np = tensor.cpu().detach().numpy()
 
     tensor_bytes = pickle.dumps(tensor_np)
-
-    tensor_bytes = zlib.compress(tensor_bytes)
+    if compress:
+        tensor_bytes = zlib.compress(tensor_bytes)
 
     tensor_b64 = base64.b64encode(tensor_bytes).decode("utf-8")
     return tensor_b64
 
 
-def base64_to_tensor(tensor_b64: str) -> torch.Tensor:
+def base64_to_tensor(tensor_b64: str, compress=True) -> torch.Tensor:
     tensor_bytes = base64.b64decode(tensor_b64)
 
-    tensor_bytes = zlib.decompress(tensor_bytes)
+    if compress:
+        tensor_bytes = zlib.decompress(tensor_bytes)
 
     tensor_np = pickle.loads(tensor_bytes)
 

diff --git a/js/deprecated.js b/js/deprecated.js
@@ -0,0 +1,23 @@
+import { app } from "../../scripts/app.js";
+
+
+app.registerExtension({
+    name: "bizyair.deprecated.nodes",
+
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        const warning_msg = {
+            "BizyAirSetAPIKey": "It will be available until 2024/08/15. Please click \"BizyAir Key\" button to set API key instead.",
+            //"BizyAirSuperResolution": "It will be available until 2024/08/31. Please use \"☁️BizyAir Image Super Resolution\" node instead.",
+        }
+        if (Object.keys(warning_msg).includes(nodeData.name)) {
+            async function alert_deprecated(node_name, display_name) {
+                alert(`${display_name}: ${warning_msg[node_name]}`);
+            }
+            const onNodeCreated = nodeType.prototype.onNodeCreated;
+            nodeType.prototype.onNodeCreated = function (message) {
+                onNodeCreated?.apply(this, arguments);
+                alert_deprecated.call(this, nodeData.name, nodeData.display_name);
+            };
+        }
+    },
+});
diff --git a/js/set_api_key.js b/js/set_api_key.js
@@ -95,10 +95,4 @@ or you can only use nodes locally.`);
             };
         }
     },
-
-    async nodeCreated(node) {
-        if (node.widgets[0].name === 'API_KEY' && node.widgets[0].value === 'YOUR_API_KEY') {
-            alert(`The Set API Key node will be deprecated, available until 2024/08/15. Please click "BizyAir Key" button to set API key instead.`);
-        }
-    }
 });
diff --git a/showcase.py b/showcase.py
@@ -8,6 +8,10 @@
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
 
 SHOW_CASES = [
+    {
+        "title": "Super Resolution",
+        "file": "bizyair_showcase_realistic_superresolution.json",
+    },
     {
         "title": "Remove the background from the image",
         "file": "bizyair_showcase_remove_background.json",

diff --git a/supernode.py b/supernode.py
@@ -1,3 +1,4 @@
+import json
 import os
 import uuid
 
@@ -10,6 +11,8 @@
     get_api_key,
 )
 
+from .image_utils import encode_data, decode_data
+
 BIZYAIR_SERVER_ADDRESS = os.getenv(
     "BIZYAIR_SERVER_ADDRESS", "https://api.siliconflow.cn"
 )
@@ -30,7 +33,7 @@ def INPUT_TYPES(s):
     RETURN_TYPES = ("IMAGE",)
     FUNCTION = "super_resolution"
 
-    CATEGORY = "☁️BizyAir"
+    CATEGORY = "☁️BizyAir/Super Resolution"
 
     def super_resolution(self, image, scale="2x"):
         API_KEY = get_api_key()
@@ -176,12 +179,82 @@ def generate_image(self, prompt, seed, width, height, cfg, batch_size):
         return (tensors,)
 
 
+class AuraSR:
+    API_URL = f"{BIZYAIR_SERVER_ADDRESS}/supernode/aurasr"
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "super_resolution"
+
+    CATEGORY = "☁️BizyAir/Super Resolution"
+
+    def super_resolution(self, image):
+        API_KEY = get_api_key()
+        SIZE_LIMIT = 1536
+        device = image.device
+        _, w, h, c = image.shape
+        assert (
+            w <= SIZE_LIMIT and h <= SIZE_LIMIT
+        ), f"width and height must be less than {SIZE_LIMIT}x{SIZE_LIMIT}, but got {w} and {h}"
+
+        # support RGB mode only now
+        image = image[:, :, :, :3]
+
+        payload = {
+            "is_compress": True,
+            "image": None,
+        }
+        auth = f"Bearer {API_KEY}"
+        headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "authorization": auth,
+        }
+        input_image = encode_data(image)
+        payload["image"] = input_image
+        payload["is_compress"] = True
+
+        ret: str = send_post_request(self.API_URL, payload=payload, headers=headers)
+        ret = json.loads(ret)
+
+        try:
+            if "result" in ret:
+                ret = json.loads(ret["result"])
+        except Exception as e:
+            raise Exception(f"Unexpected response: {ret}")
+
+        if ret["status"] == "error":
+            raise Exception(ret["message"])
+
+        msg = ret["data"]
+        if msg["type"] not in (
+            "comfyair",
+            "bizyair",
+        ):
+            raise Exception(f"Unexpected response type: {msg}")
+
+        image_b64 = msg["data"]["payload"]
+
+        image = decode_data(image_b64)
+        image = image.to(device)
+        return (image,)
+
+
 NODE_CLASS_MAPPINGS = {
     "BizyAirSuperResolution": SuperResolution,
     "BizyAirRemoveBackground": RemoveBackground,
     "BizyAirGenerateLightningImage": GenerateLightningImage,
+    "BizyAirAuraSR": AuraSR,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
+    "BizyAirAuraSR": "☁️BizyAir Photorealistic Image Super Resolution",
     "BizyAirSuperResolution": "☁️BizyAir Anime Image Super Resolution",
     "BizyAirRemoveBackground": "☁️BizyAir Remove Image Background",
     "BizyAirGenerateLightningImage": "☁️BizyAir Generate Photorealistic Images",