optimized

generate_image3.py

@@ -1,21 +1,33 @@
 import hashlib
 from pathlib import Path
 import torch
-from diffusers import DiffusionPipeline
+
 from diffusers import StableDiffusionXLPipeline
+from diffusers import DPMSolverMultistepScheduler
 from peft import PeftModel
 
-
-from project.generate_image import negative_prompt
-
 CACHE_DIR = Path("images/cache")
 CACHE_DIR.mkdir(parents=True, exist_ok=True)
 
+DEVICE = "cuda"
+
 pipe = StableDiffusionXLPipeline.from_pretrained(
     "stabilityai/stable-diffusion-xl-base-1.0",
-    dtype=torch.float32
-).to("cuda")
+    torch_dtype=torch.float32,
+    use_safetensors=True
+).to(DEVICE)
+
+# scheduler mais rápido
+pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+
+# otimizações de memória
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+
+# VAE precisa FP32
 pipe.vae.to(torch.float32)
+
+# carregar LoRA
 pipe.unet = PeftModel.from_pretrained(pipe.unet, "sdxl_lora")
 pipe.unet = pipe.unet.merge_and_unload()
 pipe.unet.eval()
@@ -44,26 +56,31 @@ def generate_slide_image(prompt: str) -> str:
 
     print("Generating image...")
 
-    image = pipe(
-        prompt,
-        negative_prompt=NEGATIVE,
-        num_inference_steps=30,
-        guidance_scale=7,
-        height=1024,
-        width=1024
-    ).images[0]
+    with torch.no_grad():
+
+        image = pipe(
+            prompt,
+            negative_prompt=NEGATIVE,
+            num_inference_steps=25,
+            guidance_scale=7,
+            height=1024,
+            width=1024
+        ).images[0]
 
     image.save(image_path)
 
     return str(image_path)
 
+
 if __name__ == "__main__":
+
     prompt = """
     clean enterprise illustration,
     corporate presentation slide,
     minimal design,
     white background,
     blue gradient accents,
-    expose oracle integration connected to autonomous database and exposing through api gateway
+    oracle integration connected to autonomous database through api gateway
     """
+
     generate_slide_image(prompt)