support fluxcontrolnet (#196)

* support fluxcontrolnet * refine * refine
siliconflow · Nov 7, 2024 · ab5bd4e · ab5bd4e
1 parent 7a784a1
commit ab5bd4e
Show file tree

Hide file tree

Showing 8 changed files with 118 additions and 14 deletions.
diff --git a/bizyair_extras/__init__.py b/bizyair_extras/__init__.py
@@ -1,6 +1,7 @@
 from .nodes_comfyui_instantid import *
 from .nodes_comfyui_pulid_flux import *
 from .nodes_controlnet import *
+from .nodes_custom_sampler import *
 from .nodes_flux import *
 from .nodes_image_utils import *
 from .nodes_ipadapter_plus.nodes_ipadapter_plus import *

diff --git a/bizyair_extras/nodes_controlnet.py b/bizyair_extras/nodes_controlnet.py
@@ -17,3 +17,35 @@ def INPUT_TYPES(s):
     RETURN_TYPES = (data_types.CONTROL_NET,)
 
     # FUNCTION = "set_controlnet_type"
+
+
+class ControlNetInpaintingAliMamaApply(BizyAirBaseNode):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "positive": (data_types.CONDITIONING,),
+                "negative": (data_types.CONDITIONING,),
+                "control_net": (data_types.CONTROL_NET,),
+                "vae": (data_types.VAE,),
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+                "strength": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01},
+                ),
+                "start_percent": (
+                    "FLOAT",
+                    {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+                "end_percent": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+            }
+        }
+
+    # FUNCTION = "apply_inpaint_controlnet"
+    CATEGORY = "conditioning/controlnet"
+    RETURN_TYPES = (data_types.CONDITIONING, data_types.CONDITIONING)
+    RETURN_NAMES = ("positive", "negative")
diff --git a/bizyair_extras/nodes_custom_sampler.py b/bizyair_extras/nodes_custom_sampler.py
@@ -0,0 +1,28 @@
+from bizyair import BizyAirBaseNode, data_types
+
+
+class CFGGuider(BizyAirBaseNode):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "model": (data_types.MODEL,),
+                "positive": (data_types.CONDITIONING,),
+                "negative": (data_types.CONDITIONING,),
+                "cfg": (
+                    "FLOAT",
+                    {
+                        "default": 8.0,
+                        "min": 0.0,
+                        "max": 100.0,
+                        "step": 0.1,
+                        "round": 0.01,
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("GUIDER",)
+
+    # FUNCTION = "get_guider"
+    CATEGORY = "sampling/custom_sampling/guiders"
diff --git a/bizyair_extras/nodes_sd3.py b/bizyair_extras/nodes_sd3.py
@@ -18,3 +18,35 @@ def INPUT_TYPES(s):
     # FUNCTION = "load_clip"
 
     CATEGORY = "advanced/loaders"
+
+
+class ControlNetApplySD3(BizyAirBaseNode):
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "positive": (data_types.CONDITIONING,),
+                "negative": (data_types.CONDITIONING,),
+                "control_net": (data_types.CONTROL_NET,),
+                "vae": (data_types.VAE,),
+                "image": ("IMAGE",),
+                "strength": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01},
+                ),
+                "start_percent": (
+                    "FLOAT",
+                    {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+                "end_percent": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+            }
+        }
+
+    CATEGORY = "conditioning/controlnet"
+    # DEPRECATED = True
+    NODE_DISPLAY_NAME = "Apply Controlnet with VAE"
+    RETURN_TYPES = (data_types.CONDITIONING, data_types.CONDITIONING)
+    RETURN_NAMES = ("positive", "negative")
diff --git a/src/bizy_server/utils.py b/src/bizy_server/utils.py
@@ -6,6 +6,7 @@
 
 TYPE_OPTIONS = {
     "lora": "bizyair/lora",
+    "controlnet": "bizyair/controlnet",
     # "other": "other",
 }
 ALLOW_TYPES = list(TYPE_OPTIONS.values())

diff --git a/src/bizyair/common/client.py b/src/bizyair/common/client.py
@@ -92,7 +92,7 @@ def process_response_data(response_data: dict) -> dict:
         try:
             msg = json.loads(response_data["result"])
         except json.JSONDecodeError:
-            raise ValueError("Failed to decode JSON from response.")
+            raise ValueError(f"Failed to decode JSON from response. {response_data=}")
     else:
         # Handle local response directly
         msg = response_data

diff --git a/src/bizyair/configs/models.yaml b/src/bizyair/configs/models.yaml
@@ -5,6 +5,7 @@ model_hub:
 
 model_types:
   loras: bizyair/lora
+  controlnet: bizyair/controlnet
   # folder_name, server_folder_name
   # checkpoints: bizyair/checkpoint
   # vae: bizyair/vae
@@ -18,6 +19,8 @@ routing_configs:
     route: /supernode/flux-dev-bizyair-comfy-ksampler-speedup
   flux-schnell:
     route: /supernode/flux-bizyair-sdxl-comfy-ksampler
+  flux-vae:
+    route: /supernode/flux-vae-bizyair-comfy-ksampler
   ultimatesdupscale:
     route: /supernode/bizyair-ultimate-sd-upscale-ksampler
   sd3_5_large:
@@ -45,7 +48,7 @@ routing_rules:
     config: flux-schnell
   - name: flux-vae
     pattern: ^flux/ae.sft$
-    config: flux-schnell
+    config: flux-vae
   - name: sd3_5_large
     pattern: ^sd3.5_large.safetensors$
     config: sd3_5_large

diff --git a/src/bizyair/path_utils/path_manager.py b/src/bizyair/path_utils/path_manager.py
@@ -90,7 +90,9 @@ def guess_url_from_node(
                                     + "/supernode/bizyair-flux-dev-comfy-pulid"
                                 )
 
-                            if class_type_table.get("LoraLoader", False):
+                            if class_type_table.get(
+                                "LoraLoader", False
+                            ) or class_type_table.get("ControlNetLoader", False):
                                 node["inputs"][
                                     "weight_dtype"
                                 ] = "fp8_e4m3fn"  # set to fp8_e4m3fn for lora
@@ -174,17 +176,22 @@ def convert_prompt_label_path_to_real_path(prompt: dict[str, dict[str, any]]) ->
     for unique_id in prompt:
         new_prompt[unique_id] = copy.copy(prompt[unique_id])
         inputs = copy.copy(prompt[unique_id]["inputs"])
-        if "lora_name" in inputs:
-            lora_name = inputs["lora_name"]
-            new_lora_name = filename_path_mapping.get("loras", {}).get(lora_name, None)
-            if new_lora_name:
-                inputs["lora_name"] = new_lora_name
-            else:
-                file_list = get_filename_list("loras")
-                if lora_name not in file_list:
-                    raise ValueError(
-                        f"Lora name '{lora_name}' not found in file list. Available lora names: {', '.join(file_list)}"
-                    )
+
+        for key, folder_name in [
+            ("lora_name", "loras"),
+            ("control_net_name", "controlnet"),
+        ]:
+            if key in inputs:
+                value = inputs[key]
+                new_value = filename_path_mapping.get(folder_name, {}).get(value, None)
+                if new_value:
+                    inputs[key] = new_value
+                else:
+                    file_list = get_filename_list(folder_name)
+                    if value not in file_list:
+                        raise ValueError(
+                            f"{key} '{value}' not found in file list. Available {key} names: {', '.join(file_list)}"
+                        )
 
         new_prompt[unique_id]["inputs"] = inputs
     return new_prompt