Update app.py

app.py

@@ -10,22 +10,27 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
+import cv2
 from PIL import Image
 from io import BytesIO
-from diffusers import AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image
+from diffusers.utils import load_image
+from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting, UNet2DConditionModel
+from controlnet_aux import HEDdetector
 
 DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>⚠️ This space is running on the CPU. This demo doesn't work on CPU 😞! Run on a GPU by duplicating this space or test our website for free and unlimited by <a href='https://squaadai.com'>clicking here</a>, which provides these and more options.</p>"
 
 MAX_SEED = np.iinfo(np.int32).max
-CACHE_EXAMPLES = torch.cuda.is_available() and os.getenv("CACHE_EXAMPLES") == "1"
 MAX_IMAGE_SIZE = int(os.getenv("MAX_IMAGE_SIZE", "1824"))
 USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD") == "1"
 ENABLE_USE_LORA = os.getenv("ENABLE_USE_LORA", "1") == "1"
+ENABLE_USE_LORA2 = os.getenv("ENABLE_USE_LORA2", "1") == "1"
 ENABLE_USE_VAE = os.getenv("ENABLE_USE_VAE", "1") == "1"
-ENABLE_USE_IMG2IMG = os.getenv("ENABLE_USE_VAE", "1") == "1"
+ENABLE_USE_IMG2IMG = os.getenv("ENABLE_USE_IMG2IMG", "1") == "1"
+ENABLE_USE_CONTROLNET = os.getenv("ENABLE_USE_CONTROLNET", "1") == "1"
+ENABLE_USE_CONTROLNETINPAINT = os.getenv("ENABLE_USE_CONTROLNETINPAINT", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -34,10 +39,9 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-
 @spaces.GPU
 def generate(
-    prompt: str,
+    prompt: str = "",
     negative_prompt: str = "",
     prompt_2: str = "",
     negative_prompt_2: str = "",
@@ -49,40 +53,88 @@ def generate(
     height: int = 1024,
     guidance_scale_base: float = 5.0,
     num_inference_steps_base: int = 25,
+    controlnet_conditioning_scale: float = 1,
+    control_guidance_start: float = 0,
+    control_guidance_end: float = 1,
     strength_img2img: float = 0.7,
     use_vae: bool = False,
     use_lora: bool = False,
+    use_lora2: bool = False,
     model = 'stabilityai/stable-diffusion-xl-base-1.0',
     vaecall = 'madebyollin/sdxl-vae-fp16-fix',
     lora = '',
+    lora2 = '',
+    controlnet_model = 'diffusers/controlnet-canny-sdxl-1.0',
     lora_scale: float = 0.7,
+    lora_scale2: float = 0.7,
     use_img2img: bool = False,
+    use_controlnet: bool = False,
+    use_controlnetinpaint: bool = False,
     url = '',
+    controlnet_img = '',
+    controlnet_inpaint = '',
 ):
     if torch.cuda.is_available():
-            
-        if not use_img2img:
-            pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch.float16)
 
+        if not use_img2img:
+            pipe = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+                        
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
-            
+                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+                
         if use_img2img:
-            pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
+            pipe = AutoPipelineForImage2Image.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+
+            init_image = load_image(url)
 
             if use_vae:  
                 vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
-                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
-                    
-            response = requests.get(url)
-            init_image = Image.open(BytesIO(response.content)).convert("RGB")
-            init_image = init_image.resize((width, height))
+                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16)
+
+        if use_controlnet:
+            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
+            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16)
+
+            image = load_image(controlnet_img)
+            
+            image = np.array(image)
+            image = cv2.Canny(image, 250, 255)
+            image = image[:, :, None]
+            image = np.concatenate([image, image, image], axis=2)
+            image = Image.fromarray(image)
+
+            if use_vae:
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
+                
+        if use_controlnetinpaint:
+            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
+            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16)
 
+            image_start = load_image(controlnet_img)
+            image = load_image(controlnet_img)
+            image_mask = load_image(controlnet_img2img)
+            
+            image = np.array(image)
+            image = cv2.Canny(image, 100, 200)
+            image = image[:, :, None]
+            image = np.concatenate([image, image, image], axis=2)
+            image = Image.fromarray(image)
+            
+            if use_vae:
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
+                
         if use_lora:
-            pipe.load_lora_weights(lora)
-            pipe.fuse_lora(lora_scale)
-                                   
+            pipe.load_lora_weights(lora, adapter_name="1")
+            pipe.set_adapters("1", adapter_weights=[lora_scale])
+            
+        if use_lora2:
+            pipe.load_lora_weights(lora, adapter_name="1")
+            pipe.load_lora_weights(lora2, adapter_name="2")
+            pipe.set_adapters(["1", "2"], adapter_weights=[lora_scale, lora_scale2])
+
         if ENABLE_CPU_OFFLOAD:
             pipe.enable_model_cpu_offload()
             
@@ -91,9 +143,9 @@ def generate(
 
         if USE_TORCH_COMPILE:
             pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        
+            
     generator = torch.Generator().manual_seed(seed)
-
+    
     if not use_negative_prompt:
         negative_prompt = None  # type: ignore
     if not use_prompt_2:
@@ -101,20 +153,41 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-    if not use_img2img:
-        return pipe(
+    if use_controlnetinpaint:
+        image = pipe(
+            prompt=prompt,
+            strength=strength_img2img,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            eta=0.0,
+            mask_image=image_mask,
+            image=image_start,
+            control_image=image,
+            negative_prompt=negative_prompt,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale_base,
+            num_inference_steps=num_inference_steps_base,
+            generator=generator,
+        ).images[0]
+        return image
+    if use_controlnet:
+        image = pipe(
             prompt=prompt,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            control_guidance_start=control_guidance_start,
+            control_guidance_end=control_guidance_end,
+            image=image,
             negative_prompt=negative_prompt,
             prompt_2=prompt_2,
-            negative_prompt_2=negative_prompt_2,
             width=width,
             height=height,
+            negative_prompt_2=negative_prompt_2,
             guidance_scale=guidance_scale_base,
             num_inference_steps=num_inference_steps_base,
             generator=generator,
-            output_type="pil",
         ).images[0]
-    else:
+        return image
+    elif use_img2img:
         images = pipe(
             prompt=prompt,
             image=init_image,
@@ -127,14 +200,20 @@ def generate(
             guidance_scale=guidance_scale_base,
             num_inference_steps=num_inference_steps_base,
             generator=generator,
-            output_type="pil",
         ).images[0]
         return images
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-]
+    else:
+        return pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            prompt_2=prompt_2,
+            negative_prompt_2=negative_prompt_2,
+            width=width,
+            height=height,
+            guidance_scale=guidance_scale_base,
+            num_inference_steps=num_inference_steps_base,
+            generator=generator,
+        ).images[0]
 
 with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.HTML(
@@ -144,16 +223,28 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     with gr.Group():
         model = gr.Text(label='Model', placeholder='e.g. stabilityai/stable-diffusion-xl-base-1.0')
         vaecall = gr.Text(label='VAE', placeholder='e.g. madebyollin/sdxl-vae-fp16-fix')
-        lora = gr.Text(label='LoRA', placeholder='e.g. nerijs/pixel-art-xl')
+        lora = gr.Text(label='LoRA 1', placeholder='e.g. nerijs/pixel-art-xl')
+        lora2 = gr.Text(label='LoRA 2', placeholder='e.g. nerijs/pixel-art-xl')
+        controlnet_model = gr.Text(label='Controlnet', placeholder='e.g diffusers/controlnet-canny-sdxl-1.0')
         lora_scale = gr.Slider(
                 info="The closer to 1, the more it will resemble LoRA, but errors may be visible.",
-                label="Lora Scale",
+                label="Lora Scale 1",
+                minimum=0.01,
+                maximum=1,
+                step=0.01,
+                value=0.7,
+            )
+        lora_scale2 = gr.Slider(
+                info="The closer to 1, the more it will resemble LoRA, but errors may be visible.",
+                label="Lora Scale 2",
                 minimum=0.01,
                 maximum=1,
                 step=0.01,
                 value=0.7,
             )
-        url = gr.Text(label='URL (Img2Img)', placeholder='e.g https://example.com/image.png')
+        url = gr.Text(label='URL (Img2Img)')
+        controlnet_img = gr.Text(label='URL (Controlnet)', placeholder='e.g https://example.com/image.png')
+        controlnet_inpaint = gr.Text(label='URL (Controlnet - IMG2IMG)', placeholder='e.g https://example.com/image.png')
         with gr.Row():
             prompt = gr.Text(
                 placeholder="Input prompt",
@@ -166,9 +257,12 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         result = gr.Image(label="Result", show_label=False)
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
+            use_controlnet = gr.Checkbox(label='Use Controlnet', value=False, visible=ENABLE_USE_CONTROLNET)
+            use_controlnetinpaint = gr.Checkbox(label='Use Controlnet Img2Img', value=False, visible=ENABLE_USE_CONTROLNETINPAINT)
             use_img2img = gr.Checkbox(label='Use Img2Img', value=False, visible=ENABLE_USE_IMG2IMG)
             use_vae = gr.Checkbox(label='Use VAE', value=False, visible=ENABLE_USE_VAE)
-            use_lora = gr.Checkbox(label='Use Lora', value=False, visible=ENABLE_USE_LORA)
+            use_lora = gr.Checkbox(label='Use Lora 1', value=False, visible=ENABLE_USE_LORA)
+            use_lora2 = gr.Checkbox(label='Use Lora 2', value=False, visible=ENABLE_USE_LORA2)
             use_negative_prompt = gr.Checkbox(label="Use negative prompt", value=False)
             use_prompt_2 = gr.Checkbox(label="Use prompt 2", value=False)
             use_negative_prompt_2 = gr.Checkbox(label="Use negative prompt 2", value=False)
@@ -233,6 +327,33 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 step=1,
                 value=25,
             )
+        with gr.Row():
+            controlnet_conditioning_scale = gr.Slider(
+                info="controlnet_conditioning_scale",
+                label="controlnet_conditioning_scale",
+                minimum=0.01,
+                maximum=2,
+                step=0.01,
+                value=1,
+            )
+        with gr.Row():
+            control_guidance_start = gr.Slider(
+                info="control_guidance_start",
+                label="control_guidance_start",
+                minimum=0.01,
+                maximum=1,
+                step=0.01,
+                value=0,
+            )
+        with gr.Row():
+            control_guidance_end = gr.Slider(
+                info="control_guidance_end",
+                label="control_guidance_end",
+                minimum=0.01,
+                maximum=1,
+                step=0.01,
+                value=1,
+            )
         with gr.Row():
             strength_img2img = gr.Slider(
                 info="Strength for Img2Img",
@@ -243,14 +364,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 value=0.7,
             )
 
-    gr.Examples(
-        examples=examples,
-        inputs=prompt,
-        outputs=result,
-        fn=generate,
-        cache_examples=CACHE_EXAMPLES,
-    )
-
     use_negative_prompt.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_negative_prompt,
@@ -286,6 +399,13 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    use_lora2.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_lora2,
+        outputs=lora2,
+        queue=False,
+        api_name=False,
+    )
     use_img2img.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_img2img,
@@ -293,6 +413,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    use_controlnet.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_controlnet,
+        outputs=controlnet_img,
+        queue=False,
+        api_name=False,
+    )
+    use_controlnetinpaint.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_controlnetinpaint,
+        outputs=controlnet_inpaint,
+        queue=False,
+        api_name=False,
+    )
 
     gr.on(
         triggers=[
@@ -322,19 +456,30 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
             height,
             guidance_scale_base,
             num_inference_steps_base,
+            controlnet_conditioning_scale,
+            control_guidance_start,
+            control_guidance_end,
             strength_img2img,
             use_vae,
             use_lora,
+            use_lora2,
             model,
             vaecall,
             lora,
+            lora2,
+            controlnet_model,
             lora_scale,
+            lora_scale2,
             use_img2img,
+            use_controlnet,
+            use_controlnetinpaint,
             url,
+            controlnet_img,
+            controlnet_inpaint,
         ],
         outputs=result,
         api_name="run",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=20).launch()
+    demo.queue(max_size=20, default_concurrency_limit=5).launch()