import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import os

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

# Função para carregar imagens de referência
def load_reference_images(reference_folder):
    reference_images = []
    for filename in os.listdir(reference_folder):
        if filename.endswith(".png") or filename.endswith(".jpg"):
            image_path = os.path.join(reference_folder, filename)
            reference_images.append(Image.open(image_path).convert("RGB"))
    return reference_images

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

if __name__ == "__main__":
    reference_folder = ".../bimbo maker/artes-ref-galeria/pessoas/0-stick-melons"
    reference_images = load_reference_images(reference_folder)
    
    if not reference_images:
        print("No reference images found in the specified folder.")
    else:
        prompt = "A stick figure woman with a hypersexualized body, round head, and simple colors"
        output_path = "output_image.png"
        generate_image_with_reference(prompt, reference_images, output_path)
        print(f"Image generated and saved to {output_path}")