wb-droid/Conditional_Diffusion

1import gradio as gr2import torch3from torch import nn4import torchvision5from diffusers import UNet2DModel, UNet2DConditionModel, DDPMScheduler, DDPMPipeline, DDIMScheduler6from fastprogress.fastprogress import progress_bar7 8labels_map = {9    0: "T-Shirt",10    1: "Trouser",11    2: "Pullover",12    3: "Dress",13    4: "Coat",14    5: "Sandal",15    6: "Shirt",16    7: "Sneaker",17    8: "Bag",18    9: "Ankle Boot",19}20 21l2i = {l:i for i,l in labels_map.items()}22 23def label2idx(l):24    return l2i[l]25    26 27unet = torch.load("unconditional01.pt", map_location=torch.device('cpu')).to("cpu")28Emb = torch.load("unconditional_emb_01.pt", map_location=torch.device('cpu')).to("cpu")29unet.eval()30 31sched = DDIMScheduler(beta_end=0.01)32sched.set_timesteps(20)33 34@torch.no_grad35def diff_sample(model, sz, sched, hidden, **kwargs):36    x_t = torch.randn(sz)37    preds = []38    for t in progress_bar(sched.timesteps):39        with torch.no_grad(): noise = model(x_t, t, hidden).sample40        x_t = sched.step(noise, t, x_t, **kwargs).prev_sample41        preds.append(x_t.float().cpu())42    return preds43 44 45@torch.no_grad()  46def generate(classChoice):47    sz = (1,1,32,32)48    print(classChoice)49    hidden = Emb(torch.tensor([label2idx(classChoice)]*1)[:,None]).detach().to("cpu")50    preds = diff_sample(unet, sz, sched, hidden, eta=1.)51 52    return((preds[-1][0] + 0.5).squeeze().clamp(-1,1).detach().numpy())53    54with gr.Blocks() as demo:55    gr.HTML("""<h1 align="center">Conditional Diffusion with DDIM</h1>""")56    gr.HTML("""<h1 align="center">trained with FashionMNIST</h1>""")57    session_data = gr.State([])58 59    classChoice = gr.Radio(list(labels_map.values()), value="T-Shirt", label="Select the type of image to generate", info="")60    sampling_button = gr.Button("Conditional image generation")61    final_image = gr.Image(height=250,width=200) 62 63  64 65    sampling_button.click(66        generate,67        [classChoice],68        [final_image],69    )70 71demo.queue().launch(share=False, inbrowser=True)72