import torch
from diffusers import StableDiffusionPipeline

# Carregar o modelo Stable Diffusion
model_id = "stabilityai/stable-diffusion-2-1-base"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# Função para gerar uma imagem
def generate_image(prompt, output_path):
    with torch.autocast("cuda"):
        image = pipe(prompt).images[0]
    image.save(output_path)

if __name__ == "__main__":
    prompt = "A stick figure woman with a hypersexualized body, round head, and simple colors"
    output_path = "output_image.png"
    generate_image(prompt, output_path)
    print(f"Image generated and saved to {output_path}")