Update app.py (#6)

- Update app.py (3373014c5d6a6a49ca973739e58012fa3d5d8cde)


Co-authored-by: Saptarshi Mukherjee <zombie-596@users.noreply.huggingface.co>

app.py

@@ -6,6 +6,21 @@ import numpy as np
 import tensorflow as tf
 import math
 
+
+
+def linear_beta_schedule(timesteps):
+    beta_start = 0.0001
+    beta_end = 0.02
+    return torch.linspace(beta_start, beta_end, timesteps)
+    
+alphas = 1. - betas
+alphas_cumprod = torch.cumprod(alphas, axis=0)
+alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
+sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
+sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
+posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)
+
 class Block(nn.Module):
     def __init__(self, in_ch, out_ch, time_emb_dim, up=False):
         super().__init__()
@@ -153,26 +168,24 @@ def p_sample_loop(model, shape):
 def sample(model, image_size, batch_size=16, channels=3):
     return p_sample_loop(model, shape=(batch_size, channels, image_size, image_size))
 
-samples = sample(model, image_size=img_size, batch_size=64, channels=3)
-
 
-reverse_transforms = transforms.Compose([
-        transforms.Lambda(lambda t: (t + 1) / 2),
-        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC
-        transforms.Lambda(lambda t: t * 255.),
-        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),
-        transforms.ToPILImage(),
-    ])
 
-for i in range(10):
-    img  = reverse_transforms(torch.Tensor((samples[-1][i].reshape(3, img_size, img_size))))
-    plt.imshow(img)
 model = SimpleUnet()
 
 st.title("Generatig images using a diffusion model")
 model.load_state_dict(torch.load("new_linear_model_1090.pt"))
 
-result = st.button("Click to generate image")
 
-if(result):
-    model()
+if(st.button("Click to generate image")):
+    samples = sample(model, image_size=img_size, batch_size=64, channels=3)
+    for i in range(10):
+        reverse_transforms = transforms.Compose([
+            transforms.Lambda(lambda t: (t + 1) / 2),
+            transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC
+            transforms.Lambda(lambda t: t * 255.),
+            transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),
+            transforms.ToPILImage(),
+        ])
+        img  = reverse_transforms(torch.Tensor((samples[-1][i].reshape(3, img_size, img_size))))
+        
+        st.image(plt.imshow(img))