From 229a5dcd08466abda352e08db1340687334a0710 Mon Sep 17 00:00:00 2001
From: Artiprocher <wangye87v5@hotmail.com>
Date: Mon, 12 Aug 2024 17:17:25 +0800
Subject: [PATCH] update comfyui-node

---
 __init__.py                 | 58 +-----------------------------------
 nodes/dashscope_flux_api.py | 59 +++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 57 deletions(-)
 create mode 100644 nodes/dashscope_flux_api.py

diff --git a/__init__.py b/__init__.py
index dbe3c39..77e23bf 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,60 +1,4 @@
-from http import HTTPStatus
-import requests, dashscope, torch
-import numpy as np
-from PIL import Image
-from io import BytesIO
-
-
-
-class DashscopeFLUXAPI:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": (
-                    [
-                        "flux-schnell", "flux-dev", "flux-merged",
-                    ],
-                    {
-                        "default": "flux-schnell"
-                    }
-                ),
-                "api_key": ("STRING", {"multiline": False, "default": "",}),
-                "prompt": ("STRING", {"multiline": True, "default": "",}),
-                "size": (
-                    [
-                        "512*1024", "768*512", "768*1024", "1024*576", "576*1024", "1024*1024"
-                    ],
-                    {
-                        "default": "1024*1024"
-                    }
-                ),
-                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                "steps": ("INT", {"default": 30, "min": 1, "max": 50, "step": 1}),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "call_api"
-    CATEGORY = "Dashscope"
-
-    def preprocess_image(self, image):
-        image = torch.Tensor(np.array(image, dtype=np.float32) * (1 / 255))
-        return image
-
-    def call_api(self, model, api_key, prompt, size, seed, steps):
-        dashscope.api_key=api_key
-        rsp = dashscope.ImageSynthesis.call(model=model, prompt=prompt, size=size, seed=seed, steps=steps)
-        if rsp.status_code == HTTPStatus.OK:
-            print(rsp.output)
-            print(rsp.usage)
-            response = requests.get(rsp.output.results[0].url)
-            img = Image.open(BytesIO(response.content))
-            img = self.preprocess_image(img)
-            return ([img],)
-        else:
-            print('Failed, status_code: %s, code: %s, message: %s' %
-                (rsp.status_code, rsp.code, rsp.message))
+from .nodes.dashscope_flux_api import DashscopeFLUXAPI
 
 
 NODE_CLASS_MAPPINGS = {
diff --git a/nodes/dashscope_flux_api.py b/nodes/dashscope_flux_api.py
new file mode 100644
index 0000000..99e64a6
--- /dev/null
+++ b/nodes/dashscope_flux_api.py
@@ -0,0 +1,59 @@
+from http import HTTPStatus
+import requests, dashscope, torch
+import numpy as np
+from PIL import Image
+from io import BytesIO
+
+
+
+class DashscopeFLUXAPI:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": (
+                    [
+                        "flux-schnell", "flux-dev", "flux-merged",
+                    ],
+                    {
+                        "default": "flux-schnell"
+                    }
+                ),
+                "api_key": ("STRING", {"multiline": False, "default": "",}),
+                "prompt": ("STRING", {"multiline": True, "default": "",}),
+                "size": (
+                    [
+                        "512*1024", "768*512", "768*1024", "1024*576", "576*1024", "1024*1024"
+                    ],
+                    {
+                        "default": "1024*1024"
+                    }
+                ),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
+                "steps": ("INT", {"default": 30, "min": 1, "max": 50, "step": 1}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "call_api"
+    CATEGORY = "Dashscope"
+
+    def preprocess_image(self, image):
+        image = torch.Tensor(np.array(image, dtype=np.float32) * (1 / 255))
+        return image
+
+    def call_api(self, model, api_key, prompt, size, seed, steps):
+        dashscope.api_key=api_key
+        rsp = dashscope.ImageSynthesis.call(model=model, prompt=prompt, size=size, seed=seed, steps=steps)
+        if rsp.status_code == HTTPStatus.OK:
+            print(rsp.output)
+            print(rsp.usage)
+            response = requests.get(rsp.output.results[0].url)
+            img = Image.open(BytesIO(response.content))
+            img = self.preprocess_image(img)
+            return ([img],)
+        else:
+            print('Failed, status_code: %s, code: %s, message: %s' %
+                (rsp.status_code, rsp.code, rsp.message))
+
+