feat: Fooocus Expansion Nodes

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,

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"