demo

visualcloze.py

@@ -407,20 +407,21 @@ class VisualClozeModel:
                 )[-1]
 
             # Get query row
-            samples = samples[:1]
-            row_samples = []
-            start = 0
-            for size in sliced_subimage:
-                end = start + (size[0] * size[1] // 256)
-                latent_h = size[0] // 8
-                latent_w = size[1] // 8
-                row_sample = samples[:, start:end, :]
-                row_sample = rearrange(row_sample, "b (h w) (c ph pw) -> b c (h ph) (w pw)", ph=2, pw=2, h=latent_h//2, w=latent_w//2)
-                row_sample = self.ae.decode(row_sample / self.ae.config.scaling_factor + self.ae.config.shift_factor)[0]
-                row_sample = (row_sample + 1.0) / 2.0
-                row_sample.clamp_(0.0, 1.0)
-                row_samples.append(row_sample[0])
-                start = end
+            with torch.no_grad():
+                samples = samples[:1]
+                row_samples = []
+                start = 0
+                for size in sliced_subimage:
+                    end = start + (size[0] * size[1] // 256)
+                    latent_h = size[0] // 8
+                    latent_w = size[1] // 8
+                    row_sample = samples[:, start:end, :]
+                    row_sample = rearrange(row_sample, "b (h w) (c ph pw) -> b c (h ph) (w pw)", ph=2, pw=2, h=latent_h//2, w=latent_w//2)
+                    row_sample = self.ae.decode(row_sample / self.ae.config.scaling_factor + self.ae.config.shift_factor)[0]
+                    row_sample = (row_sample + 1.0) / 2.0
+                    row_sample.clamp_(0.0, 1.0)
+                    row_samples.append(row_sample[0])
+                    start = end
             
             # Convert all samples to PIL images
             output_images = []