Update app.py

app.py

@@ -1,40 +1,37 @@
-#!/usr/bin/env pythona
+#!/usr/bin/env python
 
 from __future__ import annotations
 
 import requests
 import os
 import random
-import random
-import string
 
 import gradio as gr
 import numpy as np
 import spaces
 import torch
-import gc
 import cv2
+import xformers
+import triton
 from PIL import Image
-from accelerate import init_empty_weights
 from io import BytesIO
 from diffusers.utils import load_image
-from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting, UNet2DConditionModel
-from controlnet_aux import HEDdetector
-from compel import Compel, ReturnedEmbeddingsType
-import threading
+from diffusers import StableDiffusionXLControlNetPipeline, StableDiffusionXLControlNetInpaintPipeline, ControlNetModel, AutoencoderKL, DiffusionPipeline, AutoPipelineForImage2Image, AutoPipelineForInpainting, EulerDiscreteScheduler, DPMSolverMultistepScheduler
 
 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
-CUDA_LAUNCH_BLOCKING=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_IMG2IMG", "1") == "1"
+ENABLE_USE_CONTROLNET = os.getenv("ENABLE_USE_CONTROLNET", "1") == "1"
+ENABLE_USE_CONTROLNETIMG2IMG = os.getenv("ENABLE_USE_CONTROLNET", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -43,112 +40,181 @@ def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
         seed = random.randint(0, MAX_SEED)
     return seed
 
-cached_pipelines = {}  # Dicionário para armazenar os pipelines
-cached_loras = {}
-# Crie um objeto Lock
-pipeline_lock = threading.Lock()
 
 @spaces.GPU
 def generate(
-    prompt: str = "",
+    prompt: str,
     negative_prompt: str = "",
+    prompt_2: str = "",
+    negative_prompt_2: str = "",
     use_negative_prompt: bool = False,
+    use_prompt_2: bool = False,
+    use_negative_prompt_2: bool = False,
     seed: int = 0,
     width: int = 1024,
     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_controlnetimg2img: bool = False,
     url = '',
-):    
-    global cached_pipelines, cached_loras
-
+    controlnet_img = '',
+    controlnet_img2img = '',
+):
     if torch.cuda.is_available():
-        # Construa a chave do dicionário baseada no modelo e no tipo de pipeline
-        pipeline_key = (model, use_img2img)
 
-        if pipeline_key not in cached_pipelines:
-            if not use_img2img:
-                cached_pipelines[pipeline_key] = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-            elif use_img2img:
-                cached_pipelines[pipeline_key] = AutoPipelineForImage2Image.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-    
-        pipe = cached_pipelines[pipeline_key]  # Usa o pipeline carregado da memória
-        
+        if not use_img2img:
+            scheduler = DPMSolverMultistepScheduler.from_pretrained(model, subfolder="scheduler")
+            pipe = DiffusionPipeline.from_pretrained(model, scheduler=scheduler, torch_dtype=torch.float16)
+
+            pipe.to(device)
+            
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+
+                pipe.to(device)
+                
         if use_img2img:
+            pipe = AutoPipelineForImage2Image.from_pretrained(model, torch_dtype=torch.float16)
+
             init_image = load_image(url)
-                
-        if use_lora:
-            lora_key = (lora, lora_scale)
-            if lora_key not in cached_loras:
-                adapter_name = ''.join(random.choice(string.ascii_letters) for _ in range(5))
-                pipe.load_lora_weights(lora, adapter_name=adapter_name)
-                cached_loras[lora_key] = adapter_name
-            else:
-                adapter_name = cached_loras[lora_key]
-            pipe.set_adapters(adapter_name, adapter_weights=[lora_scale])
+
+            if use_vae:  
+                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
+                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, torch_dtype=torch.float16)
+
+        if use_controlnet:
+            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
+            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, controlnet=controlnet, torch_dtype=torch.float16)
+
+            image = load_image(controlnet_img)
             
-        if use_lora2:
-            lora_key1 = (lora, lora_scale)
-            lora_key2 = (lora2, lora_scale2)
-            if lora_key1 not in cached_loras:
-                adapter_name1 = ''.join(random.choice(string.ascii_letters) for _ in range(5))
-                pipe.load_lora_weights(lora, adapter_name=adapter_name1)
-                cached_loras[lora_key1] = adapter_name1
-            else:
-                adapter_name1 = cached_loras[lora_key1]
-            if lora_key2 not in cached_loras:
-                adapter_name2 = ''.join(random.choice(string.ascii_letters) for _ in range(5))
-                pipe.load_lora_weights(lora2, adapter_name=adapter_name2)
-                cached_loras[lora_key2] = adapter_name2
-            else:
-                adapter_name2 = cached_loras[lora_key2]
-            pipe.set_adapters([adapter_name1, adapter_name2], adapter_weights=[lora_scale, lora_scale2])
+            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 = StableDiffusionXLControlNetPipeline.from_pretrained(model, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
 
-        pipe.enable_model_cpu_offload()
-        generator = torch.Generator().manual_seed(seed)
+        if use_controlnetimg2img:
+            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16)
+            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, 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, controlnet=controlnet, vae=vae, torch_dtype=torch.float16)
+                
+        if use_lora:
+            pipe.load_lora_weights(lora)
+            pipe.fuse_lora(lora_scale)
         
-        if not use_negative_prompt:
-            negative_prompt = None  # type: ignore
+        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])
             
-        with pipeline_lock:
-            if use_img2img:
-                result = pipe(
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    image=init_image,
-                    strength=strength_img2img,
-                    width=width,
-                    height=height,
-                    guidance_scale=guidance_scale_base,
-                    num_inference_steps=num_inference_steps_base,
-                    generator=generator,
-                ).images[0]
-            else:
-                result = pipe(
-                    prompt=prompt,
-                    negative_prompt=negative_prompt,
-                    width=width,
-                    height=height,
-                    guidance_scale=guidance_scale_base,
-                    num_inference_steps=num_inference_steps_base,
-                    generator=generator,
-                ).images[0]
+    generator = torch.Generator().manual_seed(seed)
     
-        # Limpeza de memória
-        del pipe
-        torch.cuda.empty_cache()
-        gc.collect()
-        return result
+    if not use_negative_prompt:
+        negative_prompt = None  # type: ignore
+    if not use_prompt_2:
+        prompt_2 = None  # type: ignore
+    if not use_negative_prompt_2:
+        negative_prompt_2 = None  # type: ignore
 
+    if use_controlnetimg2img:
+        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,
+            width=width,
+            height=height,
+            negative_prompt_2=negative_prompt_2,
+            guidance_scale=guidance_scale_base,
+            num_inference_steps=num_inference_steps_base,
+            generator=generator,
+        ).images[0]
+        return image
+    elif use_img2img:
+        images = pipe(
+            prompt=prompt,
+            image=init_image,
+            strength=strength_img2img,
+            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,
+            output_type="pil",
+        ).images[0]
+        return images
+    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,
+            output_type="pil",
+        ).images[0]
+        
 with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.HTML(
         "<p><center>📙 For any additional support, join our <a href='https://discord.gg/JprjXpjt9K'>Discord</a></center></p>"
@@ -156,8 +222,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.Markdown(DESCRIPTION, elem_id="description")
     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 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 1",
@@ -174,7 +242,9 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 step=0.01,
                 value=0.7,
             )
-        url = gr.Text(label='URL (Img2Img)')
+        url = gr.Text(label='URL (Img2Img)', placeholder='e.g https://example.com/image.png')
+        controlnet_img = gr.Text(label='URL (Controlnet)', placeholder='e.g https://example.com/image.png')
+        controlnet_img2img = gr.Text(label='URL (Controlnet - IMG2IMG)', placeholder='e.g https://example.com/image.png')
         with gr.Row():
             prompt = gr.Text(
                 placeholder="Input prompt",
@@ -187,16 +257,34 @@ 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_controlnetimg2img = gr.Checkbox(label='Use Controlnet Img2Img', value=False, visible=ENABLE_USE_CONTROLNETIMG2IMG)
             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 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)
         negative_prompt = gr.Text(
             placeholder="Input Negative Prompt",
             label="Negative prompt",
             max_lines=1,
             visible=False,
         )
+        prompt_2 = gr.Text(
+            placeholder="Input Prompt 2",
+            label="Prompt 2",
+            max_lines=1,
+            visible=False,
+        )
+        negative_prompt_2 = gr.Text(
+            placeholder="Input Negative Prompt 2",
+            label="Negative prompt 2",
+            max_lines=1,
+            visible=False,
+        )
+
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -239,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",
@@ -256,6 +371,27 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    use_prompt_2.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_prompt_2,
+        outputs=prompt_2,
+        queue=False,
+        api_name=False,
+    )
+    use_negative_prompt_2.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_negative_prompt_2,
+        outputs=negative_prompt_2,
+        queue=False,
+        api_name=False,
+    )
+    use_vae.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_vae,
+        outputs=vaecall,
+        queue=False,
+        api_name=False,
+    )
     use_lora.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_lora,
@@ -277,11 +413,27 @@ 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_controlnetimg2img.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=use_controlnetimg2img,
+        outputs=controlnet_img2img,
+        queue=False,
+        api_name=False,
+    )
 
     gr.on(
         triggers=[
             prompt.submit,
             negative_prompt.submit,
+            prompt_2.submit,
+            negative_prompt_2.submit,
             run_button.click,
         ],
         fn=randomize_seed_fn,
@@ -294,26 +446,40 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         inputs=[
             prompt,
             negative_prompt,
+            prompt_2,
+            negative_prompt_2,
             use_negative_prompt,
+            use_prompt_2,
+            use_negative_prompt_2,
             seed,
             width,
             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_controlnetimg2img,
             url,
+            controlnet_img,
+            controlnet_img2img,
         ],
         outputs=result,
         api_name="run",
     )
 
 if __name__ == "__main__":
-    demo.queue(max_size=4, default_concurrency_limit=4).launch()
+    demo.queue(max_size=20).launch()