Refactor ControlNetReq class to remove unused import and add controlnets, control_images, and controlnet_conditioning_scale attributes

modules/events/flux_events.py

@@ -5,14 +5,15 @@ import spaces
 import gradio as gr
 from huggingface_hub import ModelCard
 
-from modules.helpers.flux_helpers import FluxReq, FluxImg2ImgReq, FluxInpaintReq, ControlNetReq, gen_img
-from config import flux_models, flux_loras
+from modules.helpers.common_helpers import ControlNetReq, BaseReq, BaseImg2ImgReq, BaseInpaintReq
+from modules.helpers.flux_helpers import gen_img
+from config import flux_loras
 
 loras = flux_loras
 
 
 # Event functions
-def update_fast_generation(model, fast_generation):
+def update_fast_generation(fast_generation):
     if fast_generation:
         return (
             gr.update(
@@ -125,7 +126,7 @@ def update_selected_lora(custom_lora):
         )
 
 
-def add_to_enabled_loras(model, selected_lora, enabled_loras):
+def add_to_enabled_loras(selected_lora, enabled_loras):
     lora_data = loras
     try:
         selected_lora = int(selected_lora)
@@ -233,7 +234,7 @@ def generate_image(
             "vae": vae,
             "controlnet_config": None,
         }
-        base_args = FluxReq(**base_args)
+        base_args = BaseReq(**base_args)
         
         if len(enabled_loras) > 0:
             base_args.loras = []
@@ -252,7 +253,7 @@ def generate_image(
             image = img2img_image
             strength = float(img2img_strength)
             
-            base_args = FluxImg2ImgReq(
+            base_args = BaseImg2ImgReq(
                 **base_args.__dict__,
                 image=image,
                 strength=strength
@@ -263,7 +264,7 @@ def generate_image(
             strength = float(inpaint_strength)
             
             if image and mask_image:
-                base_args = FluxInpaintReq(
+                base_args = BaseInpaintReq(
                     **base_args.__dict__,
                     image=image,
                     mask_image=mask_image,
@@ -289,7 +290,7 @@ def generate_image(
                 base_args.controlnet_config.control_images.append(depth_image)
                 base_args.controlnet_config.controlnet_conditioning_scale.append(float(depth_strength))
         else:
-            base_args = FluxReq(**base_args.__dict__)
+            base_args = BaseReq(**base_args.__dict__)
         
         return gr.update(
             value=gen_img(base_args),

modules/helpers/flux_helpers.py

@@ -6,10 +6,6 @@ from diffusers import (
     AutoPipelineForText2Image,
     AutoPipelineForImage2Image,
     AutoPipelineForInpainting,
-    DiffusionPipeline,
-    AutoencoderKL,
-    FluxControlNetModel,
-    FluxMultiControlNetModel,
 )
 from huggingface_hub import hf_hub_download
 from diffusers.schedulers import *
@@ -17,56 +13,8 @@ from huggingface_hub import hf_hub_download
 from sd_embed.embedding_funcs import get_weighted_text_embeddings_flux1
 
 from .common_helpers import ControlNetReq, BaseReq, BaseImg2ImgReq, BaseInpaintReq, cleanup, get_controlnet_images, resize_images
-
-
-def load_sd():
-    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-
-    # Models
-    models = [
-        {
-            "repo_id": "black-forest-labs/FLUX.1-dev",
-            "loader": "flux",
-            "compute_type": torch.bfloat16,
-        }
-    ]
-
-    for model in models:
-        try:
-            model["pipeline"] = AutoPipelineForText2Image.from_pretrained(
-                model['repo_id'],
-                vae=AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device),
-                torch_dtype=model['compute_type'],
-                safety_checker=None,
-                variant="fp16"
-            ).to(device)
-        except:
-            model["pipeline"] = AutoPipelineForText2Image.from_pretrained(
-                model['repo_id'],
-                vae=AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device),
-                torch_dtype=model['compute_type'],
-                safety_checker=None
-            ).to(device)
-
-        model["pipeline"].enable_model_cpu_offload()
-
-    # VAE n Refiner
-    flux_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
-    sdxl_vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device)
-    refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=sdxl_vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to(device)
-    refiner.enable_model_cpu_offload()
-
-    # ControlNet
-    controlnet = FluxMultiControlNetModel([FluxControlNetModel.from_pretrained(
-        "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro",
-        torch_dtype=torch.bfloat16
-    ).to(device)])
-
-    return device, models, flux_vae, sdxl_vae, refiner, controlnet
-
-
-device, models, flux_vae, sdxl_vae, refiner, controlnet = load_sd()
+from modules.pipelines.flux_pipelines import device, models, flux_vae, controlnet
+from modules.pipelines.common_pipelines import refiner
 
 
 def get_control_mode(controlnet_config: ControlNetReq):

modules/pipelines/common_pipelines.py

@@ -0,0 +1,19 @@
+import torch
+from diffusers import (
+    DiffusionPipeline,
+    AutoencoderKL,
+)
+from diffusers.schedulers import *
+
+
+def load_common():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # VAE n Refiner
+    sdxl_vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16).to(device)
+    refiner = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0", vae=sdxl_vae, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to(device)
+    refiner.enable_model_cpu_offload()
+    
+    return refiner, sdxl_vae
+
+refiner, sdxl_vae = load_common()

modules/pipelines/flux_pipelines.py

@@ -1,19 +1,58 @@
-# modules/pipelines/flux_pipelines.py
 
 import torch
-from diffusers import AutoPipelineForText2Image, AutoencoderKL
+from diffusers import (
+    AutoPipelineForText2Image,
+    DiffusionPipeline,
+    AutoencoderKL,
+    FluxControlNetModel,
+    FluxMultiControlNetModel,
+)
+from diffusers.schedulers import *
+
+
 
 def load_flux():
-    # Load FLUX models and pipelines
-    # ...
-    return device, models, flux_vae, controlnet
+    # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# modules/pipelines/sdxl_pipelines.py
+    # Models
+    models = [
+        {
+            "repo_id": "black-forest-labs/FLUX.1-dev",
+            "loader": "flux",
+            "compute_type": torch.bfloat16,
+        }
+    ]
+
+    for model in models:
+        try:
+            model["pipeline"] = AutoPipelineForText2Image.from_pretrained(
+                model['repo_id'],
+                vae=AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device),
+                torch_dtype=model['compute_type'],
+                safety_checker=None,
+                variant="fp16"
+            ).to(device)
+        except:
+            model["pipeline"] = AutoPipelineForText2Image.from_pretrained(
+                model['repo_id'],
+                vae=AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device),
+                torch_dtype=model['compute_type'],
+                safety_checker=None
+            ).to(device)
+
+        model["pipeline"].enable_model_cpu_offload()
+
+    # VAE n Refiner
+    flux_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+
+    # ControlNet
+    controlnet = FluxMultiControlNetModel([FluxControlNetModel.from_pretrained(
+        "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro",
+        torch_dtype=torch.bfloat16
+    ).to(device)])
+
+    return device, models, flux_vae, controlnet
 
-import torch
-from diffusers import AutoPipelineForText2Image, AutoencoderKL
 
-def load_sdxl():
-    # Load SDXL models and pipelines
-    # ...
-    return device, models, sdxl_vae, controlnet
+device, models, flux_vae, controlnet = load_flux()

tabs/image_tab.py

@@ -144,13 +144,13 @@ def flux_tab():
     
     # Events
     # Base Options
-    fast_generation.change(update_fast_generation, [model, fast_generation], [image_guidance_scale, image_num_inference_steps]) # Fast Generation # type: ignore
+    fast_generation.change(update_fast_generation, [fast_generation], [image_guidance_scale, image_num_inference_steps]) # Fast Generation # type: ignore
     
 
     # Lora Gallery
     lora_gallery.select(selected_lora_from_gallery, None, selected_lora)
     custom_lora.change(update_selected_lora, custom_lora, [custom_lora, selected_lora])
-    add_lora.click(add_to_enabled_loras, [model, selected_lora, enabled_loras], [selected_lora, custom_lora_info, enabled_loras])
+    add_lora.click(add_to_enabled_loras, [selected_lora, enabled_loras], [selected_lora, custom_lora_info, enabled_loras])
     enabled_loras.change(update_lora_sliders, enabled_loras, [lora_slider_0, lora_slider_1, lora_slider_2, lora_slider_3, lora_slider_4, lora_slider_5, lora_remove_0, lora_remove_1, lora_remove_2, lora_remove_3, lora_remove_4, lora_remove_5]) # type: ignore
 
     for i in range(6):