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

# Login com token via Secrets
login(os.environ["HF_TOKEN"])

# Carregar pipeline (sem variant)
pipe = StableDiffusion3Pipeline.from_pretrained(
    "stabilityai/stable-diffusion-3.5-medium",
    torch_dtype=torch.float16
).to("cuda")

def gerar_imagem(prompt):
    image = pipe(
        prompt,
        num_inference_steps=30,
        guidance_scale=4.5
    ).images[0]
    return image

demo = gr.Interface(
    fn=gerar_imagem,
    inputs=gr.Textbox(label="Prompt"),
    outputs=gr.Image(label="Imagem"),
    title="Stable Diffusion 3.5 Medium",
    description="Gere imagens com SD3.5 diretamente no Hugging Face"
)

demo.launch(share=True)