Update app.py

app.py

@@ -14,7 +14,7 @@ import random
 
 os.system("pip install -e ./controlnet_aux")
 
-from controlnet_aux import OpenposeDetector, CannyDetector
+from controlnet_aux import OpenposeDetector, CannyDetector, AnylineDetector
 from depth_anything_v2.dpt import DepthAnythingV2
 
 from huggingface_hub import hf_hub_download
@@ -51,16 +51,14 @@ from diffusers import FluxControlNetPipeline, FluxControlNetModel
 from diffusers.models import FluxMultiControlNetModel
 
 base_model = 'black-forest-labs/FLUX.1-dev'
-controlnet_model = 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro'
+controlnet_model = 'Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0'
 controlnet = FluxControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.bfloat16)
 controlnet = FluxMultiControlNetModel([controlnet])
 pipe = FluxControlNetPipeline.from_pretrained(base_model, controlnet=controlnet, torch_dtype=torch.bfloat16)
 pipe.to("cuda")
 
-mode_mapping = {"canny":0, "tile":1, "depth":2, "blur":3, "openpose":4, "gray":5, "low quality": 6}
-strength_mapping = {"canny":0.65, "tile":0.45, "depth":0.55, "blur":0.45, "openpose":0.55, "gray":0.45, "low quality": 0.4}
-
 canny = CannyDetector()
+anyline = AnylineDetector.from_pretrained("TheMistoAI/MistoLine/Anyline", filename="MTEED.pth")
 open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -90,6 +88,10 @@ def extract_canny(image):
     processed_image_canny = canny(image)
     return processed_image_canny
 
+def extract_soft_edge(image):
+    processed_image_soft_edge = anyline(image, detect_resolution=1280)
+    return processed_image_soft_edge
+
 def apply_gaussian_blur(image, kernel_size=(21, 21)):
     image = convert_from_image_to_cv2(image)
     blurred_image = convert_from_cv2_to_image(cv2.GaussianBlur(image, kernel_size, 0))
@@ -144,27 +146,21 @@ def resize_img(input_image, max_side=768, min_side=512, size=None,
     return input_image
 
 @spaces.GPU(duration=180)
-def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed, progress=gr.Progress(track_tqdm=True)):
-        
-    control_mode_num = mode_mapping[control_mode]
-    
+def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, control_guidance_end, seed, progress=gr.Progress(track_tqdm=True)):
+            
     if cond_in is None:
         if image_in is not None:
             image_in = resize_img(load_image(image_in))
             if control_mode == "canny":
                 control_image = extract_canny(image_in)
+            elif control_mode == "soft edge":
+                control_image = extract_soft_edge(image_in)
             elif control_mode == "depth":
                 control_image = extract_depth(image_in)
             elif control_mode == "openpose":
                 control_image = extract_openpose(image_in)
-            elif control_mode == "blur":
-                control_image = apply_gaussian_blur(image_in)
-            elif control_mode == "low quality":
-                control_image = add_gaussian_noise(image_in)
             elif control_mode == "gray":
                 control_image = convert_to_grayscale(image_in)
-            elif control_mode == "tile":
-                control_image = tile(image_in)
     else:
         control_image = resize_img(load_image(cond_in))
 
@@ -173,7 +169,7 @@ def infer(cond_in, image_in, prompt, inference_steps, guidance_scale, control_mo
     image = pipe(
         prompt, 
         control_image=[control_image],
-        control_mode=[control_mode_num],
+        control_guidance_end=[control_guidance_end],
         width=width,
         height=height,
         controlnet_conditioning_scale=[control_strength],
@@ -196,9 +192,9 @@ css="""
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("""
-        # FLUX.1-dev-ControlNet-Union-Pro
-        A unified ControlNet for FLUX.1-dev model from the InstantX team and Shakker Labs. Model card: [Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro](https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro). <br />
-        The recommended strength: {"canny":0.65, "tile":0.45, "depth":0.55, "blur":0.45, "openpose":0.55, "gray":0.45, "low quality": 0.4}. Long prompt is preferred by FLUX.1.
+        # FLUX.1-dev-ControlNet-Union-Pro-2.0
+        A unified ControlNet for FLUX.1-dev model from Shakker Labs. Model card: [Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0](https://huggingface.co/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro-2.0). <br />
+        The recommended strength: {"canny":0.70, "soft edge":0.70, "depth":0.80, "openpose":0.90, "gray":0.90}, control guidance end: {"canny":0.80, "soft edge":0.80, "depth":0.80, "openpose":0.65, "gray":0.80}
         """)
         
         with gr.Column():
@@ -214,7 +210,7 @@ with gr.Blocks(css=css) as demo:
                     
                     with gr.Accordion("Controlnet"):
                         control_mode = gr.Radio(
-                            ["canny", "depth", "openpose", "gray", "blur", "tile", "low quality"], label="Mode", value="gray",
+                            ["canny", "soft edge", "depth", "openpose", "gray"], label="Mode", value="canny",
                             info="select the control mode, one for all"
                         )
                         
@@ -223,7 +219,15 @@ with gr.Blocks(css=css) as demo:
                             minimum=0,
                             maximum=1.0,
                             step=0.05,
-                            value=0.50,
+                            value=0.70,
+                        )
+
+                        control_guidance_end = gr.Slider(
+                            label="control guidance end",
+                            minimum=0,
+                            maximum=1.0,
+                            step=0.05,
+                            value=0.80,
                         )
                     
                     seed = gr.Slider(
@@ -239,7 +243,7 @@ with gr.Blocks(css=css) as demo:
                         with gr.Column():
                             with gr.Row():
                                 inference_steps = gr.Slider(label="Inference steps", minimum=1, maximum=50, step=1, value=24)
-                                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
+                                guidance_scale = gr.Slider(label="Guidance scale", minimum=1.0, maximum=7.0, step=0.1, value=3.5)
                     
                     submit_btn = gr.Button("Submit")
                     
@@ -255,7 +259,7 @@ with gr.Blocks(css=css) as demo:
         api_name=False
     ).then(
         fn = infer,
-        inputs = [cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, seed],
+        inputs = [cond_in, image_in, prompt, inference_steps, guidance_scale, control_mode, control_strength, control_guidance_end, seed],
         outputs = [result, processed_cond],
         show_api=False
     )