Image to video script: make determinist by random seed.

xora/examples/image_to_video.py

@@ -1,3 +1,4 @@
+import time
 import torch
 from xora.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
 from xora.models.transformers.transformer3d import Transformer3DModel
@@ -14,6 +15,8 @@ import os
 import numpy as np
 import cv2
 from PIL import Image
+from tqdm import tqdm
+import random
 
 def load_vae(vae_dir):
     vae_ckpt_path = vae_dir / "diffusion_pytorch_model.safetensors"
@@ -65,9 +68,8 @@ def load_video_to_tensor_with_resize(video_path, target_height=512, target_width
         frame_resized = center_crop_and_resize(frame_rgb, target_height, target_width)
         frames.append(frame_resized)
     cap.release()
-    video_np = np.array(frames)
+    video_np = (np.array(frames) / 127.5) - 1.0
     video_tensor = torch.tensor(video_np).permute(3, 0, 1, 2).float()
-    video_tensor = (video_tensor / 127.5) - 1.0
     return video_tensor
 
 def load_image_to_tensor_with_resize(image_path, target_height=512, target_width=768):
@@ -154,9 +156,13 @@ def main():
         'media_items': media_items,
     }
 
-    generator = torch.Generator(device="cpu").manual_seed(args.seed)
+    start_time = time.time()
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed(args.seed)
+    generator = torch.Generator(device="cuda").manual_seed(args.seed)
 
-    # Run the pipeline
     images = pipeline(
         num_inference_steps=args.num_inference_steps,
         num_images_per_prompt=args.num_images_per_prompt,
@@ -173,20 +179,27 @@ def main():
         vae_per_channel_normalize=True,
         conditioning_method=ConditioningMethod.FIRST_FRAME
     ).images
-
     # Save output video
+    def get_unique_filename(base, ext, dir='.', index_range=1000):
+        for i in range(index_range):
+            filename = os.path.join(dir, f"{base}_{i}{ext}")
+            if not os.path.exists(filename):
+                return filename
+        raise FileExistsError(f"Could not find a unique filename after {index_range} attempts.")
+
+
     for i in range(images.shape[0]):
         video_np = images.squeeze(0).permute(1, 2, 3, 0).cpu().float().numpy()
         video_np = (video_np * 255).astype(np.uint8)
         fps = args.frame_rate
         height, width = video_np.shape[1:3]
-        filename = lambda base, ext, dir='.': next(
-            os.path.join(dir, f"{base}_{i}{ext}") for i in range(1000) if
-            not os.path.exists(os.path.join(dir, f"{base}_{i}{ext}")))
-        out = cv2.VideoWriter(filename(f"video_output_{i}", ".mp4", "."), cv2.VideoWriter_fourcc(*'mp4v'), fps,
-                              (width, height))
+        output_filename = get_unique_filename(f"video_output_{i}", ".mp4", ".")
+
+        out = cv2.VideoWriter(output_filename, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+
         for frame in video_np[..., ::-1]:
             out.write(frame)
+
         out.release()