Update app.py

app.py

@@ -3,23 +3,15 @@ import random
 
 import gradio as gr
 import numpy as np
-import PIL.Image
+from PIL import Image, ImageOps
 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 = """
@@ -27,6 +19,7 @@ function refresh() {
     const url = new URL(window.location);
 }
 """
+
 def nms(x, t, s):
     x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
 
@@ -65,7 +58,7 @@ def HWC3(x):
 DESCRIPTION = ''''''
 
 if not torch.cuda.is_available():
-    DESCRIPTION += ""
+    DESCRIPTION += "GPU not available. Using CPU."
 
 style_list = [
     {
@@ -124,26 +117,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]:
+def apply_style(style_name, positive, negative=""):
     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(
@@ -154,7 +138,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,
@@ -163,60 +147,51 @@ 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)
 
-    z = np.zeros_like(y, dtype=np.uint8)
-    z[y > t] = 255
-    return z
-
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+def randomize_seed_fn(seed, randomize_seed):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-@spaces.GPU
+@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,
+)
 def run(
-    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,
+    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,
     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)
@@ -234,11 +209,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,
@@ -270,58 +245,50 @@ 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", 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)
-
+        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)
 
     inputs = [
         image,
@@ -336,6 +303,7 @@ 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],
@@ -345,7 +313,5 @@ 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(show_error=True)
+demo.queue().launch(share=True, show_error=True)