Update kandinsky/t2v_pipeline.py

kandinsky/t2v_pipeline.py

@@ -111,9 +111,9 @@ class Kandinsky4T2VPipeline:
             width: int = 512,
             height: int = 512,
             seed: int = None,
-            return_frames: bool = True
+            return_frames: bool = False
     ):
-        num_steps = 8
+        num_steps = 4
 
         # SEED 
         if seed is None:
@@ -135,7 +135,7 @@ class Kandinsky4T2VPipeline:
         if (height, width) not in self.RESOLUTIONS[self.resolution]:
             raise ValueError(f"Wrong height, width pair. Available (height, width) are: {self.RESOLUTIONS[self.resolution]}")
 
-        if num_steps != 8:
+        if num_steps != 4:
             raise NotImplementedError(f"In the distilled version number of steps have to be strictly equal to 4")
 
         # PREPARATION
@@ -198,4 +198,4 @@ class Kandinsky4T2VPipeline:
                         return_images.append(ToPILImage()(image))
                     return return_images
                 else:  
-                    torchvision.io.write_video(save_path, 255. * images.squeeze(0).float().permute(1, 2, 3, 0).cpu().numpy(), fps=15, options = {"crf": "5"})
+                    torchvision.io.write_video(save_path, 255. * images.squeeze(0).float().permute(1, 2, 3, 0).cpu().numpy(), fps=23, options = {"crf": "5"})