bugfix

README.md

@@ -1,5 +1,5 @@
 ---
-title: Flux Inpaint
+title: Flux-dev controlnet inpainting with lora
 emoji: 🖼
 colorFrom: purple
 colorTo: red

app.py

@@ -22,42 +22,36 @@ from diffusers.utils import load_image, make_image_grid
 
 import json
 from preprocessor import Preprocessor
-from diffusers.pipelines.flux.pipeline_flux_controlnet_inpaint import FluxControlNetInpaintPipeline
-from diffusers.models.controlnet_flux import FluxControlNetModel
-from diffusers.models import FluxMultiControlNetModel
+from diffusers import FluxControlNetInpaintPipeline
+from diffusers.models import FluxControlNetModel
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 login(token=HF_TOKEN)
 
 MAX_SEED = np.iinfo(np.int32).max
-IMAGE_SIZE = 768
+IMAGE_SIZE = 512
 
 # init
 device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = "black-forest-labs/FLUX.1-dev"
 
-controlnet_model = 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro'
+controlnet_model = 'InstantX/FLUX.1-dev-Controlnet-Union'
 controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
-controlnet = FluxMultiControlNetModel([controlnet])
 
 
 pipe = FluxControlNetInpaintPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16).to(device)
-torch.backends.cuda.matmul.allow_tf32 = True
-pipe.vae.enable_tiling()
-pipe.vae.enable_slicing()
+
 # pipe.enable_model_cpu_offload() # for saving memory
 
 control_mode_ids = {
-    "scribble_hed": 0,
     "canny": 0,        # supported
-    "mlsd": 0,         # supported
     "tile": 1,         # supported
-    "depth_midas": 2,  # supported
+    "depth": 2,        # supported
     "blur": 3,         # supported
-    "openpose": 4,     # supported
+    "pose": 4,         # supported
     "gray": 5,         # supported
-    "low_quality": 6,  # supported
+    "lq": 6,           # supported
 }
 
 def clear_cuda_cache():
@@ -126,37 +120,36 @@ def process_mask(
     return mask
 
 def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
-    print("upload_image_to_r2", account_id, access_key, secret_key, bucket_name)
-    connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
-
-    s3 = boto3.client(
-        's3',
-        endpoint_url=connectionUrl,
-        region_name='auto',
-        aws_access_key_id=access_key,
-        aws_secret_access_key=secret_key
-    )
+    with calculateDuration("Upload image"):
+        print("upload_image_to_r2", account_id, access_key, secret_key, bucket_name)
+        connectionUrl = f"https://{account_id}.r2.cloudflarestorage.com"
+
+        s3 = boto3.client(
+            's3',
+            endpoint_url=connectionUrl,
+            region_name='auto',
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key
+        )
 
-    current_time = datetime.now().strftime("%Y/%m/%d/%H%M%S")
-    image_file = f"generated_images/{current_time}_{random.randint(0, MAX_SEED)}.png"
-    buffer = BytesIO()
-    image.save(buffer, "PNG")
-    buffer.seek(0)
-    s3.upload_fileobj(buffer, bucket_name, image_file)
-    print("upload finish", image_file)
-    
+        current_time = datetime.now().strftime("%Y/%m/%d/%H%M%S")
+        image_file = f"generated_images/{current_time}_{random.randint(0, MAX_SEED)}.png"
+        buffer = BytesIO()
+        image.save(buffer, "PNG")
+        buffer.seek(0)
+        s3.upload_fileobj(buffer, bucket_name, image_file)
+        print("upload finish", image_file)
+        
     return image_file
 
-
+@spaces.GPU(duration=120)
+@torch.inference_mode()
 def run_flux(
     image: Image.Image,
     mask: Image.Image,
     control_image: Image.Image,
     control_mode: int,
     prompt: str,
-    lora_path: str,
-    lora_weights: str,
-    lora_scale: float,
     seed_slicer: int,
     randomize_seed_checkbox: bool,
     strength_slider: float,
@@ -165,12 +158,6 @@ def run_flux(
     progress
 ) -> Image.Image:
     print("Running FLUX...")
-    clear_cuda_cache()
-    if lora_path and lora_weights:
-        with calculateDuration("load lora"):
-            print("start to load lora", lora_path, lora_weights)
-            pipe.load_lora_weights(lora_path, weight_name=lora_weights)
-    
     width, height = resolution_wh
     if randomize_seed_checkbox:
         seed_slicer = random.randint(0, MAX_SEED)
@@ -184,22 +171,73 @@ def run_flux(
                 prompt=prompt,
                 image=image,
                 mask_image=mask,
-                control_image=[control_image],
-                control_mode=[control_mode],
+                control_image=control_image,
+                control_mode=control_mode,
                 controlnet_conditioning_scale=[0.55],
                 width=width,
                 height=height,
                 strength=strength_slider,
                 generator=generator,
                 num_inference_steps=num_inference_steps_slider,
-                # max_sequence_length=256,
-                joint_attention_kwargs={"scale": lora_scale}
             ).images[0]
     progress(99, "Generate image success!")
     return generated_image
 
-@spaces.GPU(duration=120)
-@torch.inference_mode()
+
+def load_loras(lora_strings_json:str):
+    if lora_strings_json:
+        try:
+            lora_configs = json.loads(lora_strings_json)
+        except:
+            lora_configs = None
+    if lora_configs:
+        with calculateDuration("Loading LoRA weights"):
+            pipe.unload_lora_weights()
+            adapter_names = []
+            adapter_weights = []
+            for lora_info in lora_configs:
+                    lora_repo = lora_info.get("repo")
+                    weights = lora_info.get("weights")
+                    adapter_name = lora_info.get("adapter_name")
+                    adapter_weight = lora_info.get("adapter_weight")
+                    if lora_repo and weights and adapter_name:
+                        # load lora
+                        pipe.load_lora_weights(lora_repo, weight_name=weights, adapter_name=adapter_name)
+                        adapter_names.append(adapter_name)
+                        adapter_weights.append(adapter_weight)
+            # set lora weights
+            pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
+
+
+def generate_control_image(orginal_image, mask, control_mode):
+        # generated control_
+    with calculateDuration("Generate control image"):
+        preprocessor = Preprocessor()
+        if control_mode == "depth":
+            preprocessor.load("Midas")
+            control_image = preprocessor(
+                image=image,
+                image_resolution=width,
+                detect_resolution=512,
+            )
+        if control_mode == "pose":
+            preprocessor.load("Openpose")
+            control_image = preprocessor(
+                image=image,
+                hand_and_face=False,
+                image_resolution=width,
+                detect_resolution=512,
+            )
+        if control_mode == "canny":
+            preprocessor.load("Canny")
+            control_image = preprocessor(
+                image=image,
+                image_resolution=width,
+                detect_resolution=512,
+            )
+    
+    control_image = control_image.resize((width, height), Image.LANCZOS)    
+    return control_image
 
 def process(
     image_url: str,
@@ -212,9 +250,7 @@ def process(
     randomize_seed_checkbox: bool,
     strength_slider: float,
     num_inference_steps_slider: int,
-    lora_path: str,
-    lora_weights: str,
-    lora_scale: str,
+    lora_strings_json: str,
     upload_to_r2: bool,
     account_id: str,
     access_key: str,
@@ -251,54 +287,12 @@ def process(
         mask = mask.resize((width, height), Image.LANCZOS)
         mask = process_mask(mask, mask_inflation=mask_inflation_slider, mask_blur=mask_blur_slider)
 
-
-    # generated control_
-    with calculateDuration("Preprocessor Image"):
-        print("start to generate control image")
-        preprocessor = Preprocessor()
-        if control_mode == "depth_midas":
-            preprocessor.load("Midas")
-            control_image = preprocessor(
-                image=image,
-                image_resolution=width,
-                detect_resolution=512,
-            )
-        if control_mode == "openpose":
-            preprocessor.load("Openpose")
-            control_image = preprocessor(
-                image=image,
-                hand_and_face=False,
-                image_resolution=width,
-                detect_resolution=512,
-            )
-        if control_mode == "canny":
-            preprocessor.load("Canny")
-            control_image = preprocessor(
-                image=image,
-                image_resolution=width,
-                detect_resolution=512,
-            )
-
-        if control_mode == "mlsd":
-            preprocessor.load("MLSD")
-            control_image = preprocessor(
-                image=image_before,
-                image_resolution=width,
-                detect_resolution=512,
-            )
-            
-
-        if control_mode == "scribble_hed":
-            preprocessor.load("HED")
-            control_image = preprocessor(
-                image=image_before,
-                image_resolution=image_resolution,
-                detect_resolution=preprocess_resolution,
-            )
-    
-    control_image = control_image.resize((width, height), Image.LANCZOS)    
+    control_image = generate_control_image(image, mask, control_mode)
     control_mode_id = control_mode_ids[control_mode]
     clear_cuda_cache()
+
+    load_loras(lora_strings_json=lora_strings_json)
+    
     try:
         generated_image = run_flux(
             image=image,
@@ -306,9 +300,6 @@ def process(
             control_image=control_image, 
             control_mode=control_mode_id,
             prompt=inpainting_prompt_text,
-            lora_path=lora_path,
-            lora_scale=lora_scale,
-            lora_weights=lora_weights,
             seed_slicer=seed_slicer,
             randomize_seed_checkbox=randomize_seed_checkbox,
             strength_slider=strength_slider,
@@ -321,16 +312,16 @@ def process(
         result["message"] = "generate image failed"
         print(e)
         generated_image = None
+
     clear_cuda_cache()
     print("run flux finish")
     if generated_image:
         if upload_to_r2:
-            with calculateDuration("upload image"):
-                url = upload_image_to_r2(generated_image, account_id, access_key, secret_key, bucket)
-                result = {"status": "success", "message": "upload image success", "url": url}
+            url = upload_image_to_r2(generated_image, account_id, access_key, secret_key, bucket)
+            result = {"status": "success", "message": "upload image success", "url": url}    
         else:
             result = {"status": "success", "message": "Image generated but not uploaded"}
-
+    
     clear_cuda_cache()
     final_images = []
     final_images.append(image)
@@ -344,7 +335,7 @@ def process(
 
 
 with gr.Blocks() as demo:
-    gr.Markdown("Flux inpaint with lora")
+    gr.Markdown("Flux controlnet inpaint with loras")
     with gr.Row():
         with gr.Column():
             
@@ -367,41 +358,18 @@ with gr.Blocks() as demo:
             inpainting_prompt_text_component = gr.Text(
                     label="Inpainting prompt",
                     show_label=True,
-                    max_lines=1,
+                    max_lines=5,
                     placeholder="Enter text to generate inpainting",
                     container=False,
                 )
             
             control_mode = gr.Dropdown(
-                [ "canny", "depth_midas",  "openpose", "mlsd"], label="Controlnet Model", info="choose controlnet model!", value="openpose"
+                [ "canny", "depth",  "pose"], label="Controlnet Model", info="choose controlnet model!", value="canny"
             )
+            lora_strings_json = gr.Text(label="LoRA Configs (JSON List String)", placeholder='[{"repo": "lora_repo1", "weights": "weights1", "adapter_name": "adapter_name1", "adapter_weight": 1}, {"repo": "lora_repo2", "weights": "weights2", "adapter_name": "adapter_name2", "adapter_weight": 1}]', lines=5)
 
             submit_button_component = gr.Button(value='Submit', variant='primary', scale=0)
 
-            with gr.Accordion("Lora Settings", open=True):
-                lora_path = gr.Textbox(
-                    label="Lora model path",
-                    show_label=True,
-                    max_lines=1,
-                    placeholder="Enter your model path",
-                    info="Currently, only LoRA hosted on Hugging Face'model can be loaded properly.",
-                    value=""
-                )
-                lora_weights = gr.Textbox(
-                    label="Lora weights",
-                    show_label=True,
-                    max_lines=1,
-                    placeholder="Enter your lora weights name",
-                    value=""
-                )
-                lora_scale = gr.Slider(
-                    label="Lora scale",
-                    show_label=True,
-                    minimum=0,
-                    maximum=1,
-                    step=0.1,
-                    value=0.9,
-                )
                 
             with gr.Accordion("Advanced Settings", open=False):
                
@@ -487,9 +455,7 @@ with gr.Blocks() as demo:
             randomize_seed_checkbox_component,
             strength_slider_component,
             num_inference_steps_slider_component,
-            lora_path,
-            lora_weights,
-            lora_scale,
+            lora_strings_json,
             upload_to_r2,
             account_id,
             access_key,

requirements.txt

@@ -7,7 +7,7 @@ einops
 spaces
 gradio
 opencv-python
-git+https://github.com/mylovelycodes/diffusers.git
+git+https://github.com/diffusers/diffusers.git
 boto3
 sentencepiece
 peft