import os
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image
import gradio as gr

# Definir dispositivo de processamento
device = "cuda" if torch.cuda.is_available() else "cpu"

# Carregar o modelo Stable Diffusion Img2Img
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    torch_dtype=torch.float16
)
pipe = pipe.to(device)

# Função para gerar imagem
def generate_image(prompt, image):
    # Processar a imagem com o modelo
    result_image = pipe(
        prompt=prompt,
        image=image,
        num_inference_steps=50,
        strength=0.75,
        guidance_scale=7.5
    ).images[0]
    return result_image

# Configurar a interface Gradio
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Digite o seu prompt aqui..."),
        gr.Image(label="Imagem Inicial", type="pil")  # Espera uma imagem no formato PIL
    ],
    outputs=gr.Image(label="Imagem Gerada"),
    title="Gerador de Imagens com Stable Diffusion Img2Img",
    description="Insira um prompt e uma imagem inicial para gerar uma nova imagem."
)

# Executar a interface
if __name__ == "__main__":
    iface.launch()