Update app.py

app.py

@@ -3,15 +3,23 @@ import random
 
 import gradio as gr
 import numpy as np
-from PIL import Image, ImageOps
+import PIL.Image
 import torch
 import torchvision.transforms.functional as TF
+
 from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
 from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
 from controlnet_aux import PidiNetDetector, HEDdetector
 from diffusers.utils import load_image
 from huggingface_hub import HfApi
 from pathlib import Path
+from PIL import Image, ImageOps
+import torch
+import numpy as np
+import cv2
+import os
+import random
+import spaces
 from gradio_imageslider import ImageSlider
 
 js_func = """
@@ -19,7 +27,6 @@ function refresh() {
     const url = new URL(window.location);
 }
 """
-
 def nms(x, t, s):
     x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
 
@@ -58,7 +65,7 @@ def HWC3(x):
 DESCRIPTION = ''''''
 
 if not torch.cuda.is_available():
-    DESCRIPTION += "GPU not available. Using CPU."
+    DESCRIPTION += ""
 
 style_list = [
     {
@@ -117,17 +124,26 @@ 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, positive, negative=""):
+
+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)
-controlnet_canny = ControlNetModel.from_pretrained("xinsir/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16)
 
+controlnet = ControlNetModel.from_pretrained(
+    "xinsir/controlnet-scribble-sdxl-1.0",
+    torch_dtype=torch.float16
+)
+controlnet_canny = ControlNetModel.from_pretrained(
+    "xinsir/controlnet-canny-sdxl-1.0",
+    torch_dtype=torch.float16
+)
+# when test with other base model, you need to change the vae also.
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 
 pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
@@ -138,7 +154,7 @@ pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
     scheduler=eulera_scheduler,
 )
 pipe.to(device)
-
+# Load model.
 pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
     controlnet=controlnet_canny,
@@ -147,51 +163,60 @@ pipe_canny = StableDiffusionXLControlNetPipeline.from_pretrained(
     torch_dtype=torch.float16,
     scheduler=eulera_scheduler,
 )
+
 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)
 
-def randomize_seed_fn(seed, randomize_seed):
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@gr.annotations(
-    image=dict,
-    prompt=str,
-    negative_prompt=str,
-    style_name=str,
-    num_steps=int,
-    guidance_scale=float,
-    controlnet_conditioning_scale=float,
-    seed=int,
-    use_hed=bool,
-    use_canny=bool,
-    controlnet_img=Image,
-    out=Image,
-)
+@spaces.GPU
 def run(
-    image,
-    prompt,
-    negative_prompt,
-    style_name=DEFAULT_STYLE_NAME,
-    num_steps=25,
-    guidance_scale=5,
-    controlnet_conditioning_scale=1.0,
-    seed=0,
-    use_hed=False,
-    use_canny=False,
+    image: dict,
+    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:
+    # Get the composite image from the EditorValue dict
     composite_image = image['composite']
     width, height = composite_image.size
+    
+    # Calculate new dimensions to fit within 1024x1024 while maintaining aspect ratio
     max_size = 1024
     ratio = min(max_size / width, max_size / height)
     new_width = int(width * ratio)
     new_height = int(height * ratio)
+    
+    # Resize the image
     resized_image = composite_image.resize((new_width, new_height), Image.LANCZOS)
-
+    
     if use_canny:
         controlnet_img = np.array(resized_image)
         controlnet_img = cv2.Canny(controlnet_img, 100, 200)
@@ -209,11 +234,11 @@ def run(
         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,
@@ -245,50 +270,58 @@ 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_Multiplier: gr.ImageEditor(type="pil", label="Sketch your image or upload one", width=512, height=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")
-    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)
+        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", label="Sketch your image or upload one", width=512, height=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,
@@ -303,7 +336,6 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
         use_canny
     ]
     outputs = [image_slider]
-
     run_button.click(
         fn=randomize_seed_fn,
         inputs=[seed, randomize_seed],
@@ -313,5 +345,7 @@ with gr.Blocks(css="style.css", js=js_func) as demo:
     ).then(lambda x: None, inputs=None, outputs=image_slider).then(
         fn=run, inputs=inputs, outputs=outputs
     )
+    
+    
 
-demo.queue().launch(share=True, show_error=True)
+demo.queue().launch(show_error=True)