update control

app.py

@@ -30,7 +30,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 login(token=HF_TOKEN)
 
 MAX_SEED = np.iinfo(np.int32).max
-IMAGE_SIZE = 512
+IMAGE_SIZE = 1024
 
 # init
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -44,16 +44,6 @@ pipe = FluxControlNetInpaintPipeline.from_pretrained(base_model, controlnet=cont
 
 # pipe.enable_model_cpu_offload() # for saving memory
 
-control_mode_ids = {
-    "canny": 0,        # supported
-    "tile": 1,         # supported
-    "depth": 2,        # supported
-    "blur": 3,         # supported
-    "pose": 4,         # supported
-    "gray": 5,         # supported
-    "lq": 6,           # supported
-}
-
 def clear_cuda_cache():
     torch.cuda.empty_cache()
 
@@ -154,6 +144,8 @@ def run_flux(
     randomize_seed_checkbox: bool,
     strength_slider: float,
     num_inference_steps_slider: int,
+    controlnet_conditioning_scale: float,
+    guidance_scale: float,
     resolution_wh: Tuple[int, int],
     progress
 ) -> Image.Image:
@@ -173,11 +165,11 @@ def run_flux(
                 image=image,
                 mask_image=mask,
                 control_image=control_image,
-                control_mode=control_mode,
-                controlnet_conditioning_scale=0.55,
+                controlnet_conditioning_scale=controlnet_conditioning_scale,
+                strength=strength_slider,
+                guidance_scale=guidance_scale,
                 width=width,
                 height=height,
-                strength=strength_slider,
                 generator=generator,
                 num_inference_steps=num_inference_steps_slider,
             ).images[0]
@@ -212,33 +204,16 @@ def load_loras(lora_strings_json:str):
             pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
 
 
-def generate_control_image(image, mask, control_mode, width, height):
+def generate_control_image(image, mask, width, height):
         # 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,
-            )
-    
+        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
 
@@ -248,10 +223,11 @@ def process(
     inpainting_prompt_text: str,
     mask_inflation_slider: int,
     mask_blur_slider: int,
-    control_mode: str,
     seed_slicer: int,
     randomize_seed_checkbox: bool,
     strength_slider: float,
+    guidance_scale: float,
+    controlnet_conditioning_scale: float,
     num_inference_steps_slider: int,
     lora_strings_json: str,
     upload_to_r2: bool,
@@ -290,8 +266,7 @@ def process(
         mask = mask.resize((width, height), Image.LANCZOS)
         mask = process_mask(mask, mask_inflation=mask_inflation_slider, mask_blur=mask_blur_slider)
 
-    control_image = generate_control_image(image, mask, control_mode, width, height)
-    control_mode_id = control_mode_ids[control_mode]
+    control_image = generate_control_image(image, mask, width, height)
     clear_cuda_cache()
 
     load_loras(lora_strings_json=lora_strings_json)
@@ -301,12 +276,13 @@ def process(
             image=image,
             mask=mask,
             control_image=control_image, 
-            control_mode=control_mode_id,
             prompt=inpainting_prompt_text,
             seed_slicer=seed_slicer,
             randomize_seed_checkbox=randomize_seed_checkbox,
             strength_slider=strength_slider,
             num_inference_steps_slider=num_inference_steps_slider,
+            guidance_scale=guidance_scale,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
             resolution_wh=(width, height),
             progress=progress
         )
@@ -366,9 +342,7 @@ with gr.Blocks() as demo:
                     container=False,
                 )
             
-            control_mode = gr.Dropdown(
-                [ "canny"], 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)
@@ -408,6 +382,26 @@ with gr.Blocks() as demo:
                 randomize_seed_checkbox_component = gr.Checkbox(
                     label="Randomize seed", value=True)
 
+                with gr.Row():
+
+                    guidance_scale = gr.Slider(
+                        label="guidance_scale",
+                        info="guidance_scale`.",
+                        minimum=0,
+                        maximum=1,
+                        step=0.1,
+                        value=3.5,
+                    )
+
+                    controlnet_conditioning_scale = gr.Slider(
+                        label="controlnet_conditioning_scale",
+                        info="controlnet_conditioning_scale",
+                        minimum=1,
+                        maximum=10,
+                        step=0.1,
+                        value=4,
+                    )
+
                 with gr.Row():
 
                     strength_slider_component = gr.Slider(
@@ -453,10 +447,11 @@ with gr.Blocks() as demo:
             inpainting_prompt_text_component,
             mask_inflation_slider_component,
             mask_blur_slider_component,
-            control_mode,
             seed_slicer_component,
             randomize_seed_checkbox_component,
             strength_slider_component,
+            guidance_scale,
+            controlnet_conditioning_scale,
             num_inference_steps_slider_component,
             lora_strings_json,
             upload_to_r2,