Update app.py

app.py

@@ -172,7 +172,7 @@ def p_sample_loop(model, shape):
     img = torch.randn(shape, device=device)
     imgs = []
 
-    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=10):
+    for i in tqdm(reversed(range(0, timesteps)), desc='sampling loop time step', total=1):
         img = p_sample(model, img, torch.full((b,), i, device=device, dtype=torch.long), 3)
         imgs.append(img.cpu().numpy())
     return imgs
@@ -191,7 +191,7 @@ model.load_state_dict(torch.load("new_linear_model_1090.pt", map_location=torch.
 
 if(st.button("Click to generate image")):
     samples = sample(model, image_size=img_size, batch_size=64, channels=3)
-    for i in range(10):
+    for i in range(1):
         reverse_transforms = transforms.Compose([
             transforms.Lambda(lambda t: (t + 1) / 2),
             transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC