app.py

import torch
import gradio as gr
from diffusers import StableDiffusionXLPipeline

# =========================
# CARGA DEL MODELO (UNA VEZ)
# =========================

pipe = StableDiffusionXLPipeline.from_pretrained(
    "inclusionAI/TwinFlow-Z-Image-Turbo",
    torch_dtype=torch.float16,
    variant="fp16",
    use_safetensors=True
)

pipe = pipe.to("cuda")
pipe.enable_xformers_memory_efficient_attention()
pipe.enable_model_cpu_offload()

# =========================
# FUNCIÓN DE GENERACIÓN
# =========================

def generate_image(prompt, steps, seed):
    generator = None
    if seed != -1:
        generator = torch.Generator("cuda").manual_seed(seed)

    image = pipe(
        prompt=prompt,
        num_inference_steps=steps,
        guidance_scale=0.0,
        generator=generator
    ).images[0]

    return image

# =========================
# INTERFAZ GRADIO
# =========================

with gr.Blocks(title="TwinFlow-Z Image Turbo") as demo:
    gr.Markdown("## ⚡ TwinFlow-Z Image Turbo (SDXL)")

    with gr.Row():
        prompt = gr.Textbox(
            label="Prompt",
            placeholder="Describe la imagen…",
            lines=4
        )

    with gr.Row():
        steps = gr.Slider(
            minimum=2,
            maximum=10,
            value=6,
            step=1,
            label="Inference Steps"
        )
        seed = gr.Number(
            value=-1,
            label="Seed (-1 = random)",
            precision=0
        )

    generate_btn = gr.Button("Generar")
    output = gr.Image(label="Resultado", type="pil")

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, steps, seed],
        outputs=output
    )

demo.launch()