import torch
from diffusers import StableDiffusionPipeline
import matplotlib.pyplot as plt
1. Configuração do Modelo e Device
Usamos float16 para economizar VRAM e ganhar velocidade
model_id = "runwayml/stable-diffusion-v1-5"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Carregando o modelo no device: {device}...")
2. Inicialização do Pipeline
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16 if device == "cuda" else torch.float32
)
pipe = pipe.to(device)
3. Definição dos Parâmetros de Geração
prompt = "A high-tech laboratory in the middle of a tropical forest, cinematic lighting, 8k, detailed"
negative_prompt = "low quality, blurry, distorted, deformed, text"
4. Execução da Inferência
num_inference_steps: define a qualidade (20-50 é o ideal)
guidance_scale: quanto o modelo deve obedecer ao prompt (7.5 é o padrão)
print("Gerando imagem... Isso pode levar alguns segundos.")
image = pipe(
prompt,
negative_prompt=negative_prompt,
num_inference_steps=30,
guidance_scale=7.5
).images[0]
5. Exibição e Salvamento
plt.figure(figsize=(10, 10))
plt.imshow(image)
plt.axis('off')
plt.title(f"Prompt: {prompt[:50]}...")
plt.show()
image.save("resultado_geracao.png")
print("Imagem salva como 'resultado_geracao.png'.")
Conte com o apoio da comunidade Alura na sua jornada. Abraços e bons estudos!