split

app.py

@@ -75,9 +75,8 @@ MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
 @spaces.GPU
-def infer(prompt, 
+def infer_depth(prompt, 
           image = None, 
-          controlnet_type = "Depth", 
           negative_prompt = "", 
           seed = 0, 
           randomize_seed = False,
@@ -91,14 +90,8 @@ def infer(prompt,
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
     init_image = resize_image(image,  MAX_IMAGE_SIZE)
-    if controlnet_type == "Depth":
-        pipe = pipe_depth.to("cuda")
-        condi_img = process_depth_condition_midas( np.array(init_image), MAX_IMAGE_SIZE)
-    elif controlnet_type == "Canny":
-        pipe = pipe_canny.to("cuda")
-        condi_img = process_canny_condition(np.array(init_image))
-    else:
-        return None
+    pipe = pipe_depth.to("cuda")
+    condi_img = process_depth_condition_midas( np.array(init_image), MAX_IMAGE_SIZE)
     image = pipe(
         prompt= prompt ,
         image = init_image,
@@ -114,8 +107,38 @@ def infer(prompt,
     ).images[0]
     return [condi_img, image]
 
-def change_type(type):
-    return type
+@spaces.GPU
+def infer_canny(prompt, 
+          image = None, 
+          negative_prompt = "", 
+          seed = 0, 
+          randomize_seed = False,
+          guidance_scale = 6.0, 
+          num_inference_steps = 50,
+          controlnet_conditioning_scale = 0.7,
+          control_guidance_end = 0.9,
+          strength = 1.0
+        ):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator().manual_seed(seed)
+    init_image = resize_image(image,  MAX_IMAGE_SIZE)
+    pipe = pipe_canny.to("cuda")
+    condi_img = process_canny_condition(np.array(init_image))
+    image = pipe(
+        prompt= prompt ,
+        image = init_image,
+        controlnet_conditioning_scale = controlnet_conditioning_scale,
+        control_guidance_end = control_guidance_end, 
+        strength= strength , 
+        control_image = condi_img,
+        negative_prompt= negative_prompt , 
+        num_inference_steps= num_inference_steps, 
+        guidance_scale= guidance_scale,
+        num_images_per_prompt=1,
+        generator=generator,
+    ).images[0]
+    return [condi_img, image]
 
 canny_examples = [
     ["一个漂亮的女孩，高品质，超清晰，色彩鲜艳，超高分辨率，最佳品质，8k，高清，4K",
@@ -228,7 +251,7 @@ with gr.Blocks(css=css) as Kolors:
     
     with gr.Row():
         gr.Examples(
-                fn = infer,
+                fn = infer_canny,
                 examples = canny_examples,
                 inputs = [prompt, image],
                 outputs = [result],
@@ -236,29 +259,22 @@ with gr.Blocks(css=css) as Kolors:
             )
     with gr.Row():
         gr.Examples(
-                fn = infer,
+                fn = infer_depth,
                 examples = depth_examples,
                 inputs = [prompt, image],
                 outputs = [result],
                 label = "Depth"
             )
 
+
     canny_button.click(
-        fn = change_type,
-        inputs = "Canny",
-        outputs = controlnet_type
-    ).then(
-        fn = infer,
+        fn = infer_canny,
         inputs = [prompt, image, controlnet_type, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, controlnet_conditioning_scale, control_guidance_end, strength],
         outputs = [result]
     )
 
     depth_button.click(
-        fn = change_type,
-        inputs = "Depth",
-        outputs = controlnet_type
-    ).then(
-        fn = infer,
+        fn = infer_depth,
         inputs = [prompt, image, controlnet_type, negative_prompt, seed, randomize_seed, guidance_scale, num_inference_steps, controlnet_conditioning_scale, control_guidance_end, strength],
         outputs = [result]
     )