import spaces
import torch
import gradio as gr
from diffusers import StableDiffusionXLPipeline

MODEL_REPO = "BinaryLight1011/Imageflow"

pipe = StableDiffusionXLPipeline.from_pretrained(
    MODEL_REPO,
    torch_dtype=torch.float16,
    use_safetensors=True,
)
pipe.to("cuda")


@spaces.GPU(duration=60)
def generate(prompt, negative_prompt, width, height, guidance_scale, steps, seed):
    generator = torch.Generator(device="cuda").manual_seed(int(seed))
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        width=int(width),
        height=int(height),
        guidance_scale=float(guidance_scale),
        num_inference_steps=int(steps),
        generator=generator,
    ).images[0]
    return image


with gr.Blocks(title="Imageflow") as demo:
    gr.Markdown("# 🖼️ Imageflow — Text-to-Image (SDXL)")

    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", lines=4, placeholder="Descreva a imagem...")
            negative_prompt = gr.Textbox(label="Negative prompt (opcional)", lines=2)
            with gr.Row():
                width = gr.Slider(512, 1536, value=1024, step=64, label="Largura")
                height = gr.Slider(512, 1536, value=1024, step=64, label="Altura")
            guidance_scale = gr.Slider(1.0, 15.0, value=7.0, step=0.5, label="Guidance scale")
            steps = gr.Slider(10, 100, value=30, step=5, label="Inference steps")
            seed = gr.Number(value=42, label="Seed")
            btn = gr.Button("Gerar imagem", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Resultado")

    btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, width, height, guidance_scale, steps, seed],
        outputs=output_image,
    )

demo.launch()