import os
import torch
import gradio as gr
from diffusers import StableDiffusion3Pipeline
from huggingface_hub import login

# Login via token secreto nos "Settings > Secrets"
login(os.environ["HF_TOKEN"])

# Carregar pipeline com bfloat16 (mais leve, compatível com CPU/GPU Hugging Face)
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.bfloat16,
    variant="bf16"
)

# Envia para GPU se disponível
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = pipe.to(device)

# Função para gerar imagem
def gerar_imagem(prompt):
    image = pipe(
        prompt,
        num_inference_steps=30,
        guidance_scale=4.5
    ).images[0]
    return image

# Interface Gradio
demo = gr.Interface(
    fn=gerar_imagem,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(label="Imagem Gerada"),
    title="Stable Diffusion 3.5 Medium",
    description="Geração de imagens com o modelo SD 3.5 da StabilityAI"
)

# Lançar com compartilhamento público
demo.launch(share=True)