update

app.py

@@ -124,9 +124,9 @@ torch.hub.download_url_to_file(
 	'https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/2_1_lq.mp4',
 	'03.mp4')
 
-def configure_sequence_parallel(sp_size):
-    if sp_size > 1:
-        init_sequence_parallel(sp_size)
+# def configure_sequence_parallel(sp_size):
+#     if sp_size > 1:
+#         init_sequence_parallel(sp_size)
 
 @spaces.GPU(duration=120)
 def configure_runner(sp_size):
@@ -135,8 +135,8 @@ def configure_runner(sp_size):
     runner = VideoDiffusionInfer(config)
     OmegaConf.set_readonly(runner.config, False)
     
-    init_torch(cudnn_benchmark=False, timeout=datetime.timedelta(seconds=3600))
-    configure_sequence_parallel(sp_size)
+    # init_torch(cudnn_benchmark=False, timeout=datetime.timedelta(seconds=3600))
+    # configure_sequence_parallel(sp_size)
     runner.configure_dit_model(device="cuda", checkpoint='./ckpts/seedvr_ema_3b.pth')
     runner.configure_vae_model()
     # Set memory limit.
@@ -279,14 +279,9 @@ def generation_loop(video_path='./test_videos', output_dir='./results', seed=666
     original_videos, _, _ = _build_test_prompts(video_path)
 
     # divide the prompts into different groups
-    original_videos_group = partition_by_groups(
-        original_videos,
-        get_data_parallel_world_size() // get_sequence_parallel_world_size(),
-    )
+    original_videos_group = original_videos
     # store prompt mapping
-    original_videos_local = original_videos_group[
-        get_data_parallel_rank() // get_sequence_parallel_world_size()
-    ]
+    original_videos_local = original_videos_group
     original_videos_local = partition_by_size(original_videos_local, batch_size)
 
     # pre-extract the text embeddings
@@ -345,32 +340,31 @@ def generation_loop(video_path='./test_videos', output_dir='./results', seed=666
         del cond_latents
 
         # dump samples to the output directory
-        if get_sequence_parallel_rank() == 0:
-            for path, input, sample, ori_length in zip(
-                videos, input_videos, samples, ori_lengths
-            ):
-                if ori_length < sample.shape[0]:
-                    sample = sample[:ori_length]
-                filename = os.path.join(tgt_path, os.path.basename(path))
-                # color fix
-                input = (
-                    rearrange(input[:, None], "c t h w -> t c h w")
-                    if input.ndim == 3
-                    else rearrange(input, "c t h w -> t c h w")
-                )
-                if use_colorfix:
-                    sample = wavelet_reconstruction(
-                        sample.to("cpu"), input[: sample.size(0)].to("cpu")
-                    )
-                else:
-                    sample = sample.to("cpu")
-                sample = (
-                    rearrange(sample[:, None], "t c h w -> t h w c")
-                    if sample.ndim == 3
-                    else rearrange(sample, "t c h w -> t h w c")
+        for path, input, sample, ori_length in zip(
+            videos, input_videos, samples, ori_lengths
+        ):
+            if ori_length < sample.shape[0]:
+                sample = sample[:ori_length]
+            filename = os.path.join(tgt_path, os.path.basename(path))
+            # color fix
+            input = (
+                rearrange(input[:, None], "c t h w -> t c h w")
+                if input.ndim == 3
+                else rearrange(input, "c t h w -> t c h w")
+            )
+            if use_colorfix:
+                sample = wavelet_reconstruction(
+                    sample.to("cpu"), input[: sample.size(0)].to("cpu")
                 )
-                sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
-                sample = sample.to(torch.uint8).numpy()
+            else:
+                sample = sample.to("cpu")
+            sample = (
+                rearrange(sample[:, None], "t c h w -> t h w c")
+                if sample.ndim == 3
+                else rearrange(sample, "t c h w -> t h w c")
+            )
+            sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
+            sample = sample.to(torch.uint8).numpy()
 
         gc.collect()
         torch.cuda.empty_cache()