initial commit.

app.py

@@ -0,0 +1,80 @@
+import gradio as gr
+import torch
+from torch import nn
+import torchvision
+from torchvision.transforms import ToTensor
+from types import SimpleNamespace
+import matplotlib.pyplot as plt
+
+
+unet = torch.load("unet_checkpoint13.pt", map_location=torch.device('cpu')).to("cpu")
+unet.eval()
+
+@torch.no_grad
+def linear_sched(betamin=0.0001,betamax=0.02,n_steps=1000):
+    beta = torch.linspace(betamin, betamax, n_steps)
+    return SimpleNamespace(a=1.-beta, abar=(1.-beta).cumprod(dim=0), sig=beta.sqrt())
+n_steps = 1000
+lin_abar = linear_sched(betamax=0.01)
+alphabar = lin_abar.abar
+alpha = lin_abar.a
+sigma = lin_abar.sig
+
+@torch.no_grad()  
+def generate():
+    model = unet
+    sz = (1, 1, 32, 32)
+
+    ps = next(model.parameters())
+    x_t = torch.randn(sz).to(ps)
+    sample_at = {t for t in range(n_steps) if (t+101)%((t+101)//100)==0}
+    preds = []
+    img_final = img_799 = img_399 = (x_t[0].float().cpu()+0.5).squeeze().clamp(-1,1).detach().numpy()
+    for t in reversed(range(n_steps)):
+        t_batch = torch.full((x_t.shape[0],), t, device=ps.device, dtype=torch.long)
+        z = (torch.randn(x_t.shape) if t > 0 else torch.zeros(x_t.shape)).to(ps)
+        ᾱ_t1 = alphabar[t-1]  if t > 0 else torch.tensor(1)
+        b̄_t = 1-alphabar[t]
+        b̄_t1 = 1-ᾱ_t1
+        if t in sample_at: noise = model(x_t, t_batch).sample
+        x_0_hat = ((x_t - b̄_t.sqrt() * noise)/alphabar[t].sqrt())
+        x_t = x_0_hat * ᾱ_t1.sqrt()*(1-alpha[t])/b̄_t + x_t * alpha[t].sqrt()*b̄_t1/b̄_t + sigma[t]*z
+        if t in sample_at: 
+            preds.append(x_t.float().cpu())
+            img = (x_t[0].float().cpu()+0.5).squeeze().clamp(-1,1).detach().numpy()
+
+            if t >= 799: 
+                img_final = img_799 = img_399 = img
+            elif t >= 50: 
+                img_final = img_399 = img
+            else:
+                img_final = img
+
+            yield(img_799,img_399,img_final)
+    
+with gr.Blocks() as demo:
+    gr.HTML("""<h1 align="center">UNet with DPPM</h1>""")
+    gr.HTML("""<h1 align="center">trained with FashionMNIST</h1>""")
+    session_data = gr.State([])
+
+    sampling_button = gr.Button("Unconditional image generation")
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            gr.HTML("""<h3 align="left">image at step 800</h1>""")
+            step_800_image = gr.Image(height=250,width=200)             
+        with gr.Column(scale=2):
+            gr.HTML("""<h3 align="left">image at step 50</h1>""")
+            step_50_image = gr.Image(height=250,width=200) 
+        with gr.Column(scale=2):
+            gr.HTML("""<h3 align="left">final image</h1>""")
+            step_final_image = gr.Image(height=250,width=200) 
+     
+
+    sampling_button.click(
+        generate,
+        [],
+        [step_800_image, step_50_image, step_final_image],
+    )
+
+demo.queue().launch(share=False, inbrowser=True)

requirements.txt

@@ -0,0 +1,2 @@
+gradio
+torchvision

unet_checkpoint13.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22929146d48ed7a302cb8ac6166dd7c4ff3dd36a686441a8be01caaa59c786c8
+size 27239114