Update app.py

app.py

@@ -34,6 +34,7 @@ function refresh() {
     }
 }
 """
+
 def nms(x, t, s):
     x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
 
@@ -69,7 +70,8 @@ def HWC3(x):
         y = y.clip(0, 255).astype(np.uint8)
         return y
 
-DESCRIPTION = '''
+DESCRIPTION = '''# Scribble SDXL 🖋️🌄
+sketch to image with SDXL, using [@xinsir](https://huggingface.co/xinsir) [scribble sdxl controlnet](https://huggingface.co/xinsir/controlnet-scribble-sdxl-1.0), [sdxl controlnet canny](https://huggingface.co/xinsir/controlnet-canny-sdxl-1.0)
 '''
 
 if not torch.cuda.is_available():
@@ -132,17 +134,14 @@ 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
@@ -174,181 +173,98 @@ pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
 
 pipe_canny.to(device)
 
-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)
-
-    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
+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;
+}
+"""
+)
 
-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
-    )
-    
+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
+
+hed = HEDdetector.from_pretrained('lllyasviel/ControlNet')
 
-demo.queue().launch()
+if __name__ == "__main__":
+    demo = create_demo()
+    demo.launch(debug=True)