import os, torch
import gradio as gr
import torchvision.utils as vutils
from t2i_text import Trainer as T2ITrainer, build_parser as t2i_parser
from sr_diffusion import SRTrainer, build_parser as sr_parser

# Load Stage-1
t2i_args = t2i_parser().parse_args([])
t2i_trainer = T2ITrainer(t2i_args)
t2i_trainer.load_ckpt("outputs/ckpt_latest.pt")

# Load Stage-2
sr_args = sr_parser().parse_args([])
sr_trainer = SRTrainer(sr_args)
sr_ckpts = sorted([f for f in os.listdir("outputs_sr") if f.endswith(".pt")])
if sr_ckpts:
    sr_trainer.load_ckpt(os.path.join("outputs_sr", sr_ckpts[-1]))

def generate(prompt, steps_stage1=50, steps_stage2=50, n_per_prompt=1):
    # Stage-1: text → 32×32
    lr_imgs = t2i_trainer.sample(prompts=[prompt], steps=steps_stage1, n_per_prompt=n_per_prompt)
    lr_imgs = lr_imgs.to(sr_trainer.device) * 2 - 1

    # Stage-2: 32×32 → 128×128
    sr_imgs = sr_trainer.super_resolve(lr_imgs, steps=steps_stage2)

    out_path = "sr_preview.png"
    vutils.save_image(sr_imgs, out_path, nrow=n_per_prompt)
    return out_path

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", value="a brown dog"),
        gr.Slider(10, 100, 50, step=10, label="Steps Stage-1"),
        gr.Slider(10, 100, 50, step=10, label="Steps Stage-2"),
        gr.Slider(1, 4, 1, step=1, label="Images per Prompt"),
    ],
    outputs=gr.Image(type="filepath"),
    title="Text→Image (32×32) + Super-Resolution (128×128)"
)

if __name__ == "__main__":
    demo.launch()