Update app.py

app.py

@@ -21,7 +21,12 @@ model.load_state_dict(load_state_dict(cached_download(
 ), location='cpu'))
 ddim_sampler = DDIMSampler(model)
 
-def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold):
+def process(input_image, prompt, input_control, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold):
+    # TODO: Add other control tasks
+    return process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold)
+    
+
+def process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold):
     with torch.no_grad():
         img = resize_image(HWC3(input_image), image_resolution)
         H, W, C = img.shape
@@ -51,12 +56,16 @@ def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resoluti
 
 
 block = gr.Blocks().queue()
+control_task_list = [
+    "Canny Edge Map",
+    "Human Pose"
+]
 with block:
-    with gr.Row():
-        gr.Markdown("## Control Stable Diffusion with Canny Edge Maps")
+    gr.Markdown("## Adding Conditional Control to Text-to-Image Diffusion Models")
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(source='upload', type="numpy")
+            input_control = gr.Dropdown(control_list, value="Canny Edge Map", label="Control Task"),
             prompt = gr.Textbox(label="Prompt")
             run_button = gr.Button(label="Run")
             with gr.Accordion("Advanced options", open=False):
@@ -73,10 +82,10 @@ with block:
                                       value='longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair,extra digit, fewer digits, cropped, worst quality, low quality')
         with gr.Column():
             result_gallery = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
-    ips = [input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold]
+    ips = [input_image, prompt, input_control, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold]
     run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
 
-    examples = gr.Examples(examples=[["bird.png", "bird"]],inputs = [input_image, prompt], outputs = [result_gallery])
+    examples = gr.Examples(examples=[["bird.png", "bird","Canny Edge Map"]],inputs = [input_image, prompt, input_control], outputs = [result_gallery])
 
 
 block.launch(debug = True)