Update app.py

Attempt at adding scribble checkpoint

app.py

@@ -16,7 +16,7 @@ from huggingface_hub import hf_hub_url, cached_download
 
 REPO_ID = "lllyasviel/ControlNet"
 canny_checkpoint = "models/control_sd15_canny.pth"
-pose_checkpoint = "models/control_sd15_openpose.pth"
+scribble_checkpoint = "models/control_sd15_scribble.pth"
 
 canny_model = create_model('./models/cldm_v15.yaml')
 canny_model.load_state_dict(load_state_dict(cached_download(
@@ -26,14 +26,17 @@ canny_model = canny_model.cuda()
 ddim_sampler = DDIMSampler(canny_model)
 
 
-# pose_model = create_model('./models/cldm_v15.yaml')
-# pose_model.load_state_dict(load_state_dict(cached_download(
-#     hf_hub_url(REPO_ID, pose_checkpoint)
-# ), location='cpu'))
-# ddim_sampler_pose = DDIMSampler(pose_model)
+scribble_model = create_model('./models/cldm_v15.yaml')
+sribbble_model.load_state_dict(load_state_dict(cached_download(
+    hf_hub_url(REPO_ID, scribble_checkpoint)
+), location='cpu'))
+scribble_model = canny_model.cuda()
+ddim_sampler_scribble = DDIMSampler(scribble_model)
 
 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
+    if input_control == "Scribble":
+        return process_scribble(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold) 
     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):
@@ -64,11 +67,40 @@ def process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_re
         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)
+        H, W, C = img.shape
+
+        detected_map = np.zeros_like(img, dtype=np.uint8)
+        detected_map[np.min(img, axis=2) < 127] = 255
+
+        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": [scribble_model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
+        un_cond = {"c_concat": [control], "c_crossattn": [scribble_model.get_learned_conditioning([n_prompt] * num_samples)]}
+        shape = (4, H // 8, W // 8)
+
 
+        samples, intermediates = ddim_sampler_scribble.sample(ddim_steps, num_samples,
+                                                     shape, cond, verbose=False, eta=eta,
+                                                     unconditional_guidance_scale=scale,
+                                                     unconditional_conditioning=un_cond)
+
+        x_samples = scribble_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
     
 block = gr.Blocks().queue()
 control_task_list = [
-    "Canny Edge Map"
+    "Canny Edge Map",
+    "Scribble"
 ]
 with block:
     gr.Markdown("## Adding Conditional Control to Text-to-Image Diffusion Models")