Make it run on long videos

utils.py

@@ -154,20 +154,28 @@ def prepare_video(video_path:str, resolution:int, device, dtype, normalize=True,
     else:
         video = video.asnumpy()
     _, h, w, _ = video.shape
-    video = rearrange(video, "f h w c -> f c h w")
-    video = torch.Tensor(video).to(device).to(dtype)
-
-    # Use max if you want the larger side to be equal to resolution (e.g. 512)
-    # k = float(resolution) / min(h, w)
-    k = float(resolution) / max(h, w)
-    h *= k
-    w *= k
-    h = int(np.round(h / 64.0)) * 64
-    w = int(np.round(w / 64.0)) * 64
-
-    video = Resize((h, w), interpolation=InterpolationMode.BILINEAR, antialias=True)(video)
-    if normalize:
-        video = video / 127.5 - 1.0
+
+    video_resized = []
+    for f in range(video.shape[0]):
+        frame = video[f:f+1, ...]
+
+        frame = rearrange(frame, "f h w c -> f c h w")
+        frame = torch.Tensor(frame).to(device).to(dtype)
+
+        # Use max if you want the larger side to be equal to resolution (e.g. 512)
+        # k = float(resolution) / min(h, w)
+        k = float(resolution) / max(h, w)
+        h *= k
+        w *= k
+        h = int(np.round(h / 64.0)) * 64
+        w = int(np.round(w / 64.0)) * 64
+
+        frame = Resize((h, w), interpolation=InterpolationMode.BILINEAR, antialias=True)(frame)
+        if normalize:
+            frame = frame / 127.5 - 1.0
+        video_resized.append(frame)
+    video = torch.cat(video_resized)
+
     return video, output_fps