Update app.py

app.py

@@ -20,13 +20,6 @@ canny_checkpoint = "models/control_sd15_canny.pth"
 scribble_checkpoint = "models/control_sd15_scribble.pth"
 pose_checkpoint = "models/control_sd15_openpose.pth"
 
-canny_model = create_model('./models/cldm_v15.yaml').cpu()
-canny_model.load_state_dict(load_state_dict(cached_download(
-    hf_hub_url(REPO_ID, canny_checkpoint)
-), location='cuda'))
-canny_model = canny_model.cuda()
-ddim_sampler = DDIMSampler(canny_model)
-
 pose_model = create_model('./models/cldm_v15.yaml').cpu()
 pose_model.load_state_dict(load_state_dict(cached_download(
     hf_hub_url(REPO_ID, pose_checkpoint)
@@ -34,10 +27,10 @@ pose_model.load_state_dict(load_state_dict(cached_download(
 pose_model = pose_model.cuda()
 ddim_sampler_pose = DDIMSampler(pose_model)
 
-scribble_model = create_model('./models/cldm_v15.yaml')
+scribble_model = create_model('./models/cldm_v15.yaml').cpu()
 scribble_model.load_state_dict(load_state_dict(cached_download(
     hf_hub_url(REPO_ID, scribble_checkpoint)
-), location='cpu'))
+), location='cuda'))
 scribble_model = canny_model.cuda()
 ddim_sampler_scribble = DDIMSampler(scribble_model)
 
@@ -45,38 +38,9 @@ def process(input_image, prompt, input_control, a_prompt, n_prompt, num_samples,
     # TODO: Add other control tasks
     if input_control == "Scribble":
         return process_scribble(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta)
-    elif input_control == "Pose":
+    else:
         return process_pose(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, image_resolution, ddim_steps, scale, seed, eta)
-    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
-
-        detected_map = apply_canny(img, low_threshold, high_threshold)
-        detected_map = HWC3(detected_map)
-
-        control = torch.from_numpy(detected_map.copy()).float().cuda() / 255.0
-        control = torch.stack([control for _ in range(num_samples)], dim=0)
-        control = einops.rearrange(control, 'b h w c -> b c h w').clone()
-
-        seed_everything(seed)
-
-        cond = {"c_concat": [control], "c_crossattn": [canny_model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
-        un_cond = {"c_concat": [control], "c_crossattn": [canny_model.get_learned_conditioning([n_prompt] * num_samples)]}
-        shape = (4, H // 8, W // 8)
-
-        samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
-                                                     shape, cond, verbose=False, eta=eta,
-                                                     unconditional_guidance_scale=scale,
-                                                     unconditional_conditioning=un_cond)
-        x_samples = canny_model.decode_first_stage(samples)
-        x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
-
-        results = [x_samples[i] for i in range(num_samples)]
-    return [255 - detected_map] + results
-
 def process_scribble(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta):
     with torch.no_grad():
         img = resize_image(HWC3(input_image), image_resolution)
@@ -150,7 +114,6 @@ def create_canvas(w, h):
     
 block = gr.Blocks().queue()
 control_task_list = [
-    "Canny Edge Map",
     "Scribble", 
     "Pose"
 ]
@@ -172,8 +135,6 @@ with block:
             with gr.Accordion("Advanced options", open=False):
                 num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
                 image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=256)
-                low_threshold = gr.Slider(label="Canny low threshold", minimum=1, maximum=255, value=100, step=1)
-                high_threshold = gr.Slider(label="Canny high threshold", minimum=1, maximum=255, value=200, step=1)
                 ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
                 scale = gr.Slider(label="Guidance Scale", minimum=0.1, maximum=30.0, value=9.0, step=0.1)
                 seed = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True)
@@ -183,25 +144,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, input_control, 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]
     run_button.click(fn=process, inputs=ips, outputs=[result_gallery])
     examples_list = [
-        [
-            "bird.png", 
-            "bird",
-            "Canny Edge Map",
-            "best quality, extremely detailed",
-            'longbody, lowres, bad anatomy, bad hands, missing fingers, pubic hair,extra digit, fewer digits, cropped, worst quality, low quality',
-             1, 
-            512,
-            20, 
-            9.0, 
-            123490213,
-            0.0, 
-            100, 
-            200
-            
-        ],
+        
                 [
             "turtle.png", 
             "turtle",
@@ -213,9 +159,7 @@ with block:
             20, 
             9.0, 
             123490213,
-            0.0, 
-            100, 
-            200
+            0.0
             
         ]
     ]