sample_ddpm_context

app.py

@@ -198,7 +198,7 @@ def denoise_add_noise(x, t, pred_noise, z=None):
 
 # sample using standard algorithm
 @torch.no_grad()
-def sample_ddpm(n_sample, save_rate=20):
+def sample_ddpm(n_sample, context, save_rate=20):
     # x_T ~ N(0, 1), sample initial noise
     samples = torch.randn(n_sample, 3, height, height).to(device)
 
@@ -221,6 +221,30 @@ def sample_ddpm(n_sample, save_rate=20):
     intermediate = np.stack(intermediate)
     return samples, intermediate
 
+@torch.no_grad()
+def sample_ddpm_context(n_sample, save_rate=20):
+    # x_T ~ N(0, 1), sample initial noise
+    samples = torch.randn(n_sample, 3, height, height).to(device)
+
+    # array to keep track of generated steps for plotting
+    intermediate = []
+    for i in range(timesteps, 0, -1):
+        print(f'sampling timestep {i:3d}', end='\r')
+
+        # reshape time tensor
+        t = torch.tensor([i / timesteps])[:, None, None, None].to(device)
+
+        # sample some random noise to inject back in. For i = 1, don't add back in noise
+        z = torch.randn_like(samples) if i > 1 else 0
+
+        eps = nn_model(samples, t, c=context)    # predict noise e_(x_t,t)
+        samples = denoise_add_noise(samples, i, eps, z)
+        if i % save_rate ==0 or i==timesteps or i<8:
+            intermediate.append(samples.detach().cpu().numpy())
+
+    intermediate = np.stack(intermediate)
+    return samples, intermediate
+
 def greet(input):
 	steps = int(input)
 
@@ -233,7 +257,7 @@ def greet(input):
 	ctx = torch.from_numpy(mtx_2d).to(device=device).float()
 
 	#samples, intermediate = sample_ddim_context(32, ctx, n=steps)
-	samples, intermediate = sample_ddpm(32, steps)
+	samples, intermediate = sample_ddpm_context(32, ctx, steps)
 
 	#samples, intermediate = sample_ddim(32, n=steps)
 	#ctx = F.one_hot(torch.randint(0, 5, (32,)), 5).to(device=device).float()