Spaces:
Runtime error
Runtime error
| 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() | |