Update app.py

app.py

@@ -159,16 +159,21 @@ def generate_input_image(image_path,input_choice):
     print(f"input_choice: '{input_choice}'")
     if input_choice == 0 :
         conditioning_image = generate_edgeImage(image_path)
-        pipe = pipe_edge
     elif input_choice == 1 :
         conditioning_image = generate_segmt_img(image_path)
-        pipe = pipe_seg
     else:
         conditioning_image = extract_depth_map(image_path)
-        pipe = pipe_depth
-    return conditioning_image, pipe
-def generate_output_image(prompt,pipe):
+    return conditioning_image
+    
+def generate_output_image(prompt,input_choice):
     #---------------------
+    if input_choice == 0 :
+        pipe = pipe_edge
+    elif input_choice == 1 :
+        pipe = pipe_seg
+    else:
+        pipe = pipe_depth
+        
     pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
     pipe.enable_attention_slicing()
     #---------------------
@@ -255,7 +260,7 @@ with gr.Blocks() as demo:
     )
     generate_btn.click(
         fn=generate_output_image,
-        inputs=[prompt_input,pipe],
+        inputs=[prompt_input,image_dropdown],
         outputs=[output_image]
     )