From 13eab2805adb7c4e15d83dc27d021b35a9e5694e Mon Sep 17 00:00:00 2001
From: liji <896126514@qq.com>
Date: Fri, 23 Aug 2024 10:54:16 +0800
Subject: [PATCH] feat: Fooocus Expansion Nodes

---
 py/fooocusNodes.py | 35 +++++++++++++++++++++++++++++++++++
 pyproject.toml     |  2 +-
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/py/fooocusNodes.py b/py/fooocusNodes.py
index 73c47f9..c34171a 100644
--- a/py/fooocusNodes.py
+++ b/py/fooocusNodes.py
@@ -16,6 +16,7 @@
 from extras.expansion import safe_str
 import extras.face_crop as face_crop
 import modules.advanced_parameters as advanced_parameters
+from extras.expansion import FooocusExpansion as Expansion
 import extras.preprocessors as preprocessors
 import extras.ip_adapter as ip_adapter
 from nodes import SaveImage, PreviewImage, MAX_RESOLUTION, NODE_CLASS_MAPPINGS as ALL_NODE_CLASS_MAPPINGS
@@ -1458,6 +1459,39 @@ def describe(self, image, image_type):
         print(tags)
         return {"ui": {"tags": tags}, "result": (tags,)}
 
+
+class FooocusExpansion:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt": ("STRING", {"multiline": True}),
+            },
+            "optional": {
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
+                "log_prompt": ("BOOLEAN", {"default": False},),
+            },
+        }
+    RETURN_TYPES = ("STRING", "INT",)
+    RETURN_NAMES = ("final_prompt", "seed",)
+    FUNCTION = "expansion"
+
+    CATEGORY = "Fooocus"
+
+    def expansion(self, prompt, seed, log_prompt):
+        expansion = Expansion()
+        prompt = remove_empty_str([safe_str(prompt)], default='')[0]
+        max_seed = int(1024 * 1024 * 1024)
+        if not isinstance(seed, int):
+            seed = random.randint(1, max_seed)
+        if seed < 0:
+            seed = - seed
+        seed = seed % max_seed
+        expansion = expansion(prompt, seed)
+        if log_prompt:
+            print(f'[Prompt Expansion] {expansion}')
+        return {"ui": {"expansion": expansion}, "result": (expansion,)}
+
 NODE_CLASS_MAPPINGS = {
 
     "Fooocus Loader": FooocusLoader,
@@ -1471,6 +1505,7 @@ def describe(self, image, image_type):
     "Fooocus Inpaint": FooocusInpaint,
     "Fooocus PipeOut": FooocusPipeOut,
     "Fooocus Describe": FooocusDescribe,
+    "Fooocus Expansion": FooocusExpansion,
     # fix
     "Fooocus preDetailerFix": preDetailerFix,
     "Fooocus ultralyticsDetectorPipe": ultralyticsDetectorForDetailerFix,
diff --git a/pyproject.toml b/pyproject.toml
index 191b146..38c7c1c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ name = "fooocus_nodes"
 description = "This extension provides image generation features based on Fooocus."
 version = "1.0.1"
 license = "LICENSE"
-dependencies = ["accelerate==0.32.1", "pytorch_lightning==2.3.3", "pygit2==1.15.1", "opencv-contrib-python-headless==4.10.0.84", "httpx==0.27.0", "onnxruntime==1.16.3", "timm==1.0.7"]
+dependencies = ["accelerate==0.32.1", "pytorch_lightning==2.3.3", "pygit2==1.15.1", "opencv-contrib-python-headless==4.10.0.84", "httpx==0.27.0", "onnxruntime", "timm==1.0.7"]
 
 [project.urls]
 Repository = "https://github.com/Seedsa/Fooocus_Nodes"