Update infer.py

infer.py

@@ -44,40 +44,41 @@ def infer_pipe(pipe, images_batch, task_name, seed, device):
     else:
         autocast_ctx = torch.autocast(pipe.device.type)
 
-    with autocast_ctx:
-        # Convert list of images to tensor
-        images = [np.array(img.convert('RGB')).astype(np.float16) for img in images_batch]
-        test_images = torch.stack([torch.tensor(img).permute(2, 0, 1) for img in images])
-        test_images = test_images / 127.5 - 1.0
-        test_images = test_images.to(device)
+    with torch.no_grad():
+        with autocast_ctx:
+            # Convert list of images to tensor
+            images = [np.array(img.convert('RGB')).astype(np.float16) for img in images_batch]
+            test_images = torch.stack([torch.tensor(img).permute(2, 0, 1) for img in images])
+            test_images = test_images / 127.5 - 1.0
+            test_images = test_images.to(device)
 
-        task_emb = torch.tensor([1, 0]).float().unsqueeze(0).to(device)
-        task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1)
-        task_emb = task_emb.repeat(len(test_images), 1)
+            task_emb = torch.tensor([1, 0]).float().unsqueeze(0).to(device)
+            task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1)
+            task_emb = task_emb.repeat(len(test_images), 1)
 
-        # Run inference
-        preds = pipe(
-            rgb_in=test_images,
-            prompt='',
-            num_inference_steps=1,
-            generator=generator,
-            output_type='np',
-            timesteps=[999],
-            task_emb=task_emb,
-        ).images
+            # Run inference
+            preds = pipe(
+                rgb_in=test_images,
+                prompt='',
+                num_inference_steps=1,
+                generator=generator,
+                output_type='np',
+                timesteps=[999],
+                task_emb=task_emb,
+            ).images
 
-        # Post-process predictions
-        outputs = []
-        if task_name == 'depth':
-            for p in preds:
-                output_npy = p.mean(axis=-1)
-                output_color = colorize_depth_map(output_npy)
-                outputs.append(output_color)
-        else:
-            for p in preds:
-                output_npy = p
-                output_color = Image.fromarray((output_npy * 255).astype(np.uint8))
-                outputs.append(output_color)
+            # Post-process predictions
+            outputs = []
+            if task_name == 'depth':
+                for p in preds:
+                    output_npy = p.mean(axis=-1)
+                    output_color = colorize_depth_map(output_npy)
+                    outputs.append(output_color)
+            else:
+                for p in preds:
+                    output_npy = p
+                    output_color = Image.fromarray((output_npy * 255).astype(np.uint8))
+                    outputs.append(output_color)
 
     return outputs