Update app.py

app.py

@@ -15,17 +15,17 @@ high_threshold = 200
 
 
 # Models
-controlnet_canny = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
-pipe_canny = StableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", controlnet=controlnet_canny, safety_checker=None, torch_dtype=torch.float16
-)
-pipe_canny.scheduler = UniPCMultistepScheduler.from_config(pipe_canny.scheduler.config)
+# controlnet_canny = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
+# pipe_canny = StableDiffusionControlNetPipeline.from_pretrained(
+#     "runwayml/stable-diffusion-v1-5", controlnet=controlnet_canny, safety_checker=None, torch_dtype=torch.float16
+# )
+# pipe_canny.scheduler = UniPCMultistepScheduler.from_config(pipe_canny.scheduler.config)
 
-# This command loads the individual model components on GPU on-demand. So, we don't
-# need to explicitly call pipe.to("cuda").
-pipe_canny.enable_model_cpu_offload()
+# # This command loads the individual model components on GPU on-demand. So, we don't
+# # need to explicitly call pipe.to("cuda").
+# pipe_canny.enable_model_cpu_offload()
 
-pipe_canny.enable_xformers_memory_efficient_attention()
+# pipe_canny.enable_xformers_memory_efficient_attention()
 
 # Generator seed,
 generator = torch.manual_seed(0)
@@ -47,38 +47,38 @@ pipe_pose.enable_model_cpu_offload()
 pipe_pose.enable_xformers_memory_efficient_attention()
 
 
-def get_canny_filter(image):
+# def get_canny_filter(image):
     
-    if not isinstance(image, np.ndarray):
-        image = np.array(image) 
+#     if not isinstance(image, np.ndarray):
+#         image = np.array(image) 
         
-    image = cv2.Canny(image, low_threshold, high_threshold)
-    image = image[:
-, :, None]
-    image = np.concatenate([image, image, image], axis=2)
-    canny_image = Image.fromarray(image)
-    return canny_image
+#     image = cv2.Canny(image, low_threshold, high_threshold)
+#     image = image[:
+# , :, None]
+#     image = np.concatenate([image, image, image], axis=2)
+#     canny_image = Image.fromarray(image)
+#     return canny_image
 
 def get_pose(image):
     return pose_model(image) 
     
 def process(input_image, prompt, input_control):
     # TODO: Add other control tasks
-    if input_control == "Pose":
-        return process_pose(input_image, prompt)
-    else:    
-        return process_canny(input_image, prompt)
-
-def process_canny(input_image, prompt):
-    canny_image = get_canny_filter(input_image)
-    output = pipe_canny(
-        prompt,
-        canny_image,
-        generator=generator,
-        num_images_per_prompt=1,
-        num_inference_steps=20,
-    )
-    return [canny_image,output.images[0]]
+    #if input_control == "Pose":
+    return process_pose(input_image, prompt)
+    # else:    
+    #     return process_canny(input_image, prompt)
+
+# def process_canny(input_image, prompt):
+#     canny_image = get_canny_filter(input_image)
+#     output = pipe_canny(
+#         prompt,
+#         canny_image,
+#         generator=generator,
+#         num_images_per_prompt=1,
+#         num_inference_steps=20,
+#     )
+#     return [canny_image,output.images[0]]
 
 
 def process_pose(input_image, prompt):
@@ -109,22 +109,22 @@ with block:
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(source='upload', type="numpy")
-            input_control = gr.Dropdown(control_task_list, value="Scribble", label="Control Task")
+            # input_control = gr.Dropdown(control_task_list, value="Scribble", label="Control Task")
             prompt = gr.Textbox(label="Prompt")
             run_button = gr.Button(label="Run")
             
             
         with gr.Column():
             result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
-    ips = [input_image, prompt, input_control]
+    ips = [input_image, prompt]
     run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
     examples_list = [
-                [
-            "bird.png", 
-            "bird",
-            "Canny Edge Map"
+        #         [
+        #     "bird.png", 
+        #     "bird",
+        #     "Canny Edge Map"
             
-        ],
+        # ],
         
         #         [
         #     "turtle.png", 
@@ -159,7 +159,7 @@ with block:
             
          ]
     ]
-    examples = gr.Examples(examples=examples_list,inputs = [input_image, prompt, input_control], outputs = [result_gallery], cache_examples = True, fn = process)
+    examples = gr.Examples(examples=examples_list,inputs = [input_image, prompt], outputs = [result_gallery], cache_examples = True, fn = process)
     gr.Markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=RamAnanth1.ControlNet)")  
 
 block.launch(debug = True)