Update app.py

app.py

@@ -34,7 +34,6 @@ function refresh() {
     }
 }
 """
-
 def nms(x, t, s):
     x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
 
@@ -134,14 +133,17 @@ styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
 STYLE_NAMES = list(styles.keys())
 DEFAULT_STYLE_NAME = "(No style)"
 
+
 def apply_style(style_name: str, positive: str, negative: str = "") -> tuple[str, str]:
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + negative
 
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
 
+
 controlnet = ControlNetModel.from_pretrained(
     "xinsir/controlnet-scribble-sdxl-1.0",
     torch_dtype=torch.float16
@@ -173,98 +175,181 @@ pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
 
 pipe_canny.to(device)
 
-MAX_IMAGE_PIXELS = 100000000  # Adjust if needed.
-
-def resize_image(image, max_pixels=MAX_IMAGE_PIXELS):
-    """Resize an image to have at most max_pixels, maintaining aspect ratio."""
-    width, height = image.size
-    if width * height > max_pixels:
-        scale_factor = (max_pixels / (width * height)) ** 0.5
-        new_size = (int(width * scale_factor), int(height * scale_factor))
-        return image.resize(new_size, Image.ANTIALIAS)
-    return image
-
-def process(image, prompt, style, detector_name):
-    # Convert image to RGB mode if it's not already
-    if image.mode != 'RGB':
-        image = image.convert('RGB')
-    image = resize_image(image)
-
-    width, height = image.size
-
-    prompt, negative_prompt = apply_style(style, prompt)
-
-    if detector_name == "hed":
-        image = HWC3(np.array(image, dtype=np.uint8))
-        with torch.no_grad():
-            detected_map = hed(image, scribble=True)
-            detected_map = HWC3(detected_map)
-        image = Image.fromarray(detected_map)
-        images = pipe(prompt, negative_prompt=negative_prompt, image=image, height=height, width=width).images
-        return images[0]
-    elif detector_name == "scribble":
-        image = HWC3(np.array(image, dtype=np.uint8))
-        with torch.no_grad():
-            detected_map = nms(image, 127, 3.0)
-            detected_map = HWC3(detected_map)
-        image = Image.fromarray(detected_map)
-        images = pipe(prompt, negative_prompt=negative_prompt, image=image, height=height, width=width).images
-        return images[0]
-    elif detector_name == "canny":
-        image = np.array(image, dtype=np.uint8)
-        image = cv2.Canny(image, 100, 200)
-        image = image[:, :, None]
-        image = np.concatenate([image, image, image], axis=2)
-        detected_map = image
-        image = Image.fromarray(detected_map)
-        images = pipe_canny(prompt, negative_prompt=negative_prompt, image=image, height=height, width=width).images
-        return images[0]
-
-block_css = (
-    code := """
-#image_upload {
-    height: 100% !important;
-}
-#prompt_input {
-    height: 100% !important;
-}
-#select_style {
-    height: 100% !important;
-}
-#detect_method {
-    height: 100% !important;
-}
-#submit_button {
-    height: 100% !important;
-}
-"""
-)
+MAX_SEED = np.iinfo(np.int32).max
+processor = HEDdetector.from_pretrained('lllyasviel/Annotators')
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
 
-def create_demo():
-    """Create Gradio demo."""
-
-    with gr.Blocks(css=block_css) as demo:
-        gr.Markdown(DESCRIPTION)
-        with gr.Row():
-            with gr.Column():
-                input_image = gr.Image(source='upload', elem_id="image_upload", tool='editor', type="pil")
-                prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt", elem_id="prompt_input")
-                style = gr.Dropdown(STYLE_NAMES, value=DEFAULT_STYLE_NAME, label="Select style", elem_id="select_style")
-                detect_method = gr.Dropdown(choices=["scribble", "hed", "canny"], value="scribble", label="Select Detect Method", elem_id="detect_method")
-                submit_btn = gr.Button("Generate", elem_id="submit_button")
-            with gr.Column():
-                gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=2, height="auto")
-        
-        submit_btn.click(process, inputs=[input_image, prompt, style, detect_method], outputs=[gallery])
-
-        # Refresh button to apply the dark theme
-        refresh_btn = gr.Button("Refresh for Dark Theme")
-        refresh_btn.click(None, None, None, _js=js_func)
-    
-    return demo
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z
 
-hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    return seed
+
+@spaces.GPU
+def run(
+    image: PIL.Image.Image,
+    prompt: str,
+    negative_prompt: str,
+    style_name: str = DEFAULT_STYLE_NAME,
+    num_steps: int = 25,
+    guidance_scale: float = 5,
+    controlnet_conditioning_scale: float = 1.0,
+    seed: int = 0,
+    use_hed: bool = False,
+    use_canny: bool = False,
+    progress=gr.Progress(track_tqdm=True),
+) -> PIL.Image.Image:
+    width, height  = image['composite'].size
+    ratio = np.sqrt(1024. * 1024. / (width * height))
+    new_width, new_height = int(width * ratio), int(height * ratio)
+    image = image['composite'].resize((new_width, new_height))
+    
+    if use_canny:
+        controlnet_img = np.array(image)
+        controlnet_img = cv2.Canny(controlnet_img, 100, 200)
+        controlnet_img = HWC3(controlnet_img)
+        image = Image.fromarray(controlnet_img)
+    
+    elif not use_hed:
+          controlnet_img = image
+    else:
+        controlnet_img = processor(image, scribble=False)
+      # following is some processing to simulate human sketch draw, different threshold can generate different width of lines
+        controlnet_img = np.array(controlnet_img)
+        controlnet_img = nms(controlnet_img, 127, 3)
+        controlnet_img = cv2.GaussianBlur(controlnet_img, (0, 0), 3)
+
+        # higher threshold, thiner line
+        random_val = int(round(random.uniform(0.01, 0.10), 2) * 255)
+        controlnet_img[controlnet_img > random_val] = 255
+        controlnet_img[controlnet_img < 255] = 0
+        image = Image.fromarray(controlnet_img)
+    
+    
+    prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+    if use_canny:
+        out = pipe_canny(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        num_inference_steps=num_steps,
+        generator=generator,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        guidance_scale=guidance_scale,
+        width=new_width,
+        height=new_height,
+    ).images[0]
+    else:
+        out = pipe(
+        prompt=prompt,
+        negative_prompt=negative_prompt,
+        image=image,
+        num_inference_steps=num_steps,
+        generator=generator,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        guidance_scale=guidance_scale,
+        width=new_width,
+        height=new_height,).images[0]
+
+    return (controlnet_img, out)
+
+
+with gr.Blocks(css="style.css", js=js_func) as demo:
+    gr.Markdown(DESCRIPTION, elem_id="description")
+    gr.DuplicateButton(
+        value="Duplicate Space for private use",
+        elem_id="duplicate-button",
+        visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
+    )
+
+    with gr.Row():
+        with gr.Column():
+            with gr.Group():
+                image = gr.ImageEditor(type="pil", image_mode="L", crop_size=(512, 512))
+                prompt = gr.Textbox(label="Prompt")
+                style = gr.Dropdown(label="Style", choices=STYLE_NAMES, value=DEFAULT_STYLE_NAME)
+                use_hed = gr.Checkbox(label="use HED detector", value=False, info="check this box if you upload an image and want to turn it to a sketch")
+                use_canny = gr.Checkbox(label="use Canny", value=False, info="check this to use ControlNet canny instead of scribble")
+                run_button = gr.Button("Run")
+            with gr.Accordion("Advanced options", open=False):
+                negative_prompt = gr.Textbox(
+                    label="Negative prompt",
+                    value="longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
+                )
+                num_steps = gr.Slider(
+                    label="Number of steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=25,
+                )
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.1,
+                    maximum=10.0,
+                    step=0.1,
+                    value=5,
+                )
+                controlnet_conditioning_scale = gr.Slider(
+                    label="controlnet conditioning scale",
+                    minimum=0.5,
+                    maximum=5.0,
+                    step=0.1,
+                    value=0.9,
+                )
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
+                    step=1,
+                    value=0,
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+               
+        with gr.Column():
+            with gr.Group():
+                image_slider = ImageSlider(position=0.5)
+
+
+    inputs = [
+        image,
+        prompt,
+        negative_prompt,
+        style,
+        num_steps,
+        guidance_scale,
+        controlnet_conditioning_scale,
+        seed,
+        use_hed,
+        use_canny
+    ]
+    outputs = [image_slider]
+    run_button.click(
+        fn=randomize_seed_fn,
+        inputs=[seed, randomize_seed],
+        outputs=seed,
+        queue=False,
+        api_name=False,
+    ).then(lambda x: None, inputs=None, outputs=image_slider).then(
+        fn=run, inputs=inputs, outputs=outputs
+    )
+    
+    
 
-if __name__ == "__main__":
-    demo = create_demo()
-    demo.launch(debug=True)
+demo.queue().launch()