Update app.py

app.py

@@ -20,6 +20,7 @@ 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
 
 DESCRIPTION = "# Run any LoRA or SD Model"
@@ -27,6 +28,7 @@ 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"
@@ -50,17 +52,14 @@ pipeline_lock = threading.Lock()
 def generate(
     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,
     strength_img2img: float = 0.7,
+    is_sdxl: bool = False,
     use_lora: bool = False,
     use_lora2: bool = False,
     model = 'stabilityai/stable-diffusion-xl-base-1.0',
@@ -84,7 +83,7 @@ def generate(
                 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)
                 
@@ -115,49 +114,95 @@ def generate(
                 adapter_name2 = cached_loras[lora_key2]
             pipe.set_adapters([adapter_name1, adapter_name2], adapter_weights=[lora_scale, lora_scale2])
 
-    pipe.enable_model_cpu_offload()
-    generator = torch.Generator().manual_seed(seed)
+    # É SDXL 1.0 (NÃO SD 1.5)
+    if is_sdxl:
+        pipe.enable_model_cpu_offload()
+        compel = Compel(
+        tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+        text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+        returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+        requires_pooled=[False, True],
+        truncate_long_prompts=False
+        )
+        conditioning, pooled = compel(prompt)
+        generator = torch.Generator().manual_seed(seed)
+        
+        if not use_negative_prompt:
+            negative_prompt = None  # type: ignore
+            
+        with pipeline_lock:
+            if use_img2img:
+                result = pipe(
+                    prompt_embeds=conditioning,
+                    pooled_prompt_embeds=pooled,
+                    image=init_image,
+                    strength=strength_img2img,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale_base,
+                    num_inference_steps=num_inference_steps_base,
+                    generator=generator,
+                ).images[0]
+            else:
+                result = pipe(
+                    prompt_embeds=conditioning,
+                    pooled_prompt_embeds=pooled,
+                    negative_prompt=negative_prompt,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale_base,
+                    num_inference_steps=num_inference_steps_base,
+                    generator=generator,
+                ).images[0]
     
-    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
-
-    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]
+        # Limpeza de memória
+        del pipe, conditioning, pooled
+        torch.cuda.empty_cache()
+        gc.collect()
         return result
 
-    # Limpeza de memória
-    del pipe
-    torch.cuda.empty_cache()
-    gc.collect()
+    # NÃO É SDXL (E SIM SD 1.5)
+    if not is_sdxl:
+        pipe.enable_model_cpu_offload()
+        compel = Compel(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder, truncate_long_prompts=False)
+        conditioning = compel.build_conditioning_tensor(prompt)
+        negative_conditioning = compel.build_conditioning_tensor(negative_prompt)
+        [conditioning, negative_conditioning] = compel.pad_conditioning_tensors_to_same_length([conditioning, negative_conditioning])
+        generator = torch.Generator().manual_seed(seed)
+        
+        if not use_negative_prompt:
+            negative_prompt = None  # type: ignore
+            
+        with pipeline_lock:
+            if use_img2img:
+                result = pipe(
+                    prompt_embeds=conditioning,
+                    image=init_image,
+                    strength=strength_img2img,
+                    negative_prompt_embeds=negative_conditioning,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale_base,
+                    num_inference_steps=num_inference_steps_base,
+                    generator=generator,
+                ).images[0]
+            else:
+                result = pipe(
+                    prompt_embeds=conditioning,
+                    negative_prompt_embeds=negative_conditioning,
+                    width=width,
+                    height=height,
+                    guidance_scale=guidance_scale_base,
+                    num_inference_steps=num_inference_steps_base,
+                    generator=generator,
+                ).images[0]
+    
+        # Limpeza de memória
+        del pipe, conditioning, negative_conditioning
+        torch.cuda.empty_cache()
+        gc.collect()
+        return result
 
 with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     gr.HTML(
@@ -198,30 +243,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
     with gr.Accordion("Advanced options", open=False):
         with gr.Row():
             use_img2img = gr.Checkbox(label='Use Img2Img', value=False, visible=ENABLE_USE_IMG2IMG)
+            is_sdxl = gr.Checkbox(label='Is SDXL?', value=False)
             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,
@@ -281,20 +312,6 @@ 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_lora.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_lora,
@@ -309,6 +326,13 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    is_sdxl.change(
+        fn=lambda x: gr.update(visible=x),
+        inputs=is_sdxl,
+        outputs=is_sdxl,
+        queue=False,
+        api_name=False,
+    )
     use_img2img.change(
         fn=lambda x: gr.update(visible=x),
         inputs=use_img2img,
@@ -321,8 +345,6 @@ with gr.Blocks(theme=gr.themes.Soft(), css="style.css") as demo:
         triggers=[
             prompt.submit,
             negative_prompt.submit,
-            prompt_2.submit,
-            negative_prompt_2.submit,
             run_button.click,
         ],
         fn=randomize_seed_fn,
@@ -335,17 +357,14 @@ 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,
             strength_img2img,
+            is_sdxl,
             use_lora,
             use_lora2,
             model,