context

app.py

@@ -161,9 +161,38 @@ def sample_ddim(n_sample, n=20):
     intermediate = np.stack(intermediate)
     return samples, intermediate
 
+# load in model weights and set to eval mode
+nn_model.load_state_dict(torch.load(f"{save_dir}/context_model_31.pth", map_location=device))
+nn_model.eval()
+print("Loaded in Context Model")
+
+# fast sampling algorithm with context
+@torch.no_grad()
+def sample_ddim_context(n_sample, context, n=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 = []
+    step_size = timesteps // n
+    for i in range(timesteps, 0, -step_size):
+        print(f'sampling timestep {i:3d}', end='\r')
+
+        # reshape time tensor
+        t = torch.tensor([i / timesteps])[:, None, None, None].to(device)
+
+        eps = nn_model(samples, t, c=context)    # predict noise e_(x_t,t)
+        samples = denoise_ddim(samples, i, i - step_size, eps)
+        intermediate.append(samples.detach().cpu().numpy())
+
+    intermediate = np.stack(intermediate)
+    return samples, intermediate
+
 def greet(input):
 	steps = int(input)
-	samples, intermediate = sample_ddim(32, n=steps)
+	#samples, intermediate = sample_ddim(32, n=steps)
+	ctx = F.one_hot(torch.randint(0, 5, (32,)), 5).to(device=device).float()
+	samples, intermediate = sample_ddim_context(32, ctx)
 	#response = transform2(transform(np.moveaxis(samples.detach().cpu().numpy(),1,3)[-1]))
 	#response2 = transform2(transform(np.moveaxis(samples.detach().cpu().numpy(),1,3)[1]))
 	#response = im.fromarray(intermediate[24][0][1]).convert("RGB")