Update pipeline.py

pipeline.py

@@ -72,13 +72,13 @@ class OkkhorDiffusionPipeline(DiffusionPipeline):
 
         # set step values
         self.scheduler.set_timesteps(num_inference_steps)
+        with torch.no_grad():    
+            for t in self.progress_bar(self.scheduler.timesteps):
+                # 1. predict noise model_output
+                model_output = self.unet(image, t,class_labels=self.embedding).sample
 
-        for t in self.progress_bar(self.scheduler.timesteps):
-            # 1. predict noise model_output
-            model_output = self.unet(image, t,class_labels=self.embedding).sample
-
-            # 2. compute previous image: x_t -> x_t-1
-            image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
+                # 2. compute previous image: x_t -> x_t-1
+                image = self.scheduler.step(model_output, t, image, generator=generator).prev_sample
 
         image = (image / 2 + 0.5).clamp(0, 1)
         image = image.cpu().permute(0, 2, 3, 1).numpy()