Update app.py

app.py

@@ -1,21 +1,26 @@
-#!/usr/bin/env python
+#!/usr/bin/env pythona
 
 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
 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
+import threading
 
 DESCRIPTION = "# Run any LoRA or SD Model"
 if not torch.cuda.is_available():
@@ -27,10 +32,7 @@ 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_CONTROLNETINPAINT = os.getenv("ENABLE_USE_CONTROLNETINPAINT", "1") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
@@ -39,6 +41,11 @@ 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 = "",
@@ -53,88 +60,62 @@ 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 = '',
-):
+):    
+    global cached_pipelines, cached_loras
+
     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 not use_img2img:
-            pipe = DiffusionPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-                        
-            if use_vae:  
-                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-                pipe = DiffusionPipeline.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-                
-        if use_img2img:
-            pipe = AutoPipelineForImage2Image.from_pretrained(model, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
+        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 use_img2img:
             init_image = load_image(url)
-
-            if use_vae:  
-                vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-                pipe = AutoPipelineForImage2Image.from_pretrained(model, vae=vae, safety_checker=None, requires_safety_checker=False, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-
-        if use_controlnet:
-            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-            pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-
-            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, device_map="balanced", low_cpu_mem_usage=True)
-                pipe = StableDiffusionXLControlNetPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-                
-        if use_controlnetinpaint:
-            controlnet = ControlNetModel.from_pretrained(controlnet_model, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-            pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
-
-            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, device_map="balanced", low_cpu_mem_usage=True)
-                pipe = StableDiffusionXLControlNetInpaintPipeline.from_pretrained(model, safety_checker=None, requires_safety_checker=False, controlnet=controlnet, vae=vae, torch_dtype=torch.float16, device_map="balanced", low_cpu_mem_usage=True)
                 
         if use_lora:
-            pipe.load_lora_weights(lora, adapter_name="1")
-            pipe.set_adapters("1", adapter_weights=[lora_scale])
+            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_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])
-
+            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])
+    
+    pipe.to("cuda")
     generator = torch.Generator().manual_seed(seed)
     
     if not use_negative_prompt:
@@ -144,67 +125,39 @@ def generate(
     if not use_negative_prompt_2:
         negative_prompt_2 = None  # type: ignore
 
-    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,
-            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,
-        ).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,
-        ).images[0]
+    with pipeline_lock:
+        if use_img2img:
+            result = 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,
+            ).images[0]
+        else:
+            result = 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]
+        return result
+
+    # Limpeza de memória
+    del pipe
+    torch.cuda.empty_cache()
+    gc.collect()
 
 with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.HTML(
@@ -213,10 +166,8 @@ 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",
@@ -234,8 +185,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
                 value=0.7,
             )
         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",
@@ -248,10 +197,7 @@ 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 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)
@@ -318,33 +264,6 @@ 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",
@@ -376,13 +295,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         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,
@@ -404,20 +316,6 @@ 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=[
@@ -447,30 +345,20 @@ 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, default_concurrency_limit=2).launch()
+    demo.queue(max_size=4, default_concurrency_limit=4).launch()