diff --git a/generate_image3.py b/generate_image3.py index 5bacbd1..bfd9c5f 100644 --- a/generate_image3.py +++ b/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) \ No newline at end of file