update

app.py

@@ -139,9 +139,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):
@@ -150,8 +150,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/seedvr2_ema_3b.pth')
     runner.configure_vae_model()
     # Set memory limit.

common/distributed/basic.py

@@ -66,11 +66,11 @@ def init_torch(cudnn_benchmark=True, timeout=timedelta(seconds=600)):
     torch.backends.cuda.matmul.allow_tf32 = True
     torch.backends.cudnn.allow_tf32 = True
     torch.backends.cudnn.benchmark = cudnn_benchmark
-    torch.cuda.set_device(get_local_rank())
+    torch.cuda.set_device(0)
     dist.init_process_group(
         backend="nccl",
-        rank=get_global_rank(),
-        world_size=get_world_size(),
+        rank=0,
+        world_size=1,
         timeout=timeout,
     )
 

projects/video_diffusion_sr/infer.py

@@ -26,14 +26,14 @@ from common.diffusion import (
     create_sampling_timesteps_from_config,
     create_schedule_from_config,
 )
-# from common.distributed import (
-#     get_device,
-#     get_global_rank,
-# )
-
-# from common.distributed.meta_init_utils import (
-#     meta_non_persistent_buffer_init_fn,
-# )
+from common.distributed import (
+    get_device,
+    get_global_rank,
+)
+
+from common.distributed.meta_init_utils import (
+    meta_non_persistent_buffer_init_fn,
+)
 # from common.fs import download
 
 from models.dit_v2 import na
@@ -68,20 +68,20 @@ class VideoDiffusionInfer():
             return cond
         raise NotImplementedError
 
-    # @log_on_entry
-    # @log_runtime
+    @log_on_entry
+    @log_runtime
     def configure_dit_model(self, device="cpu", checkpoint=None):
         # Load dit checkpoint.
         # For fast init & resume,
         #   when training from scratch, rank0 init DiT on cpu, then sync to other ranks with FSDP.
         #   otherwise, all ranks init DiT on meta device, then load_state_dict with assign=True.
-        # if self.config.dit.get("init_with_meta_device", False):
-        #     init_device = "cpu" if get_global_rank() == 0 and checkpoint is None else "meta"
-        # else:
-        #     init_device = "cpu"
+        if self.config.dit.get("init_with_meta_device", False):
+            init_device = "cpu" if get_global_rank() == 0 and checkpoint is None else "meta"
+        else:
+            init_device = "cpu"
 
         # Create dit model.
-        with torch.device("cpu"):
+        with torch.device(init_device):
             self.dit = create_object(self.config.dit.model)
         self.dit.set_gradient_checkpointing(self.config.dit.gradient_checkpoint)
 
@@ -90,27 +90,27 @@ class VideoDiffusionInfer():
             loading_info = self.dit.load_state_dict(state, strict=True, assign=True)
             print(f"Loading pretrained ckpt from {checkpoint}")
             print(f"Loading info: {loading_info}")
-            # self.dit = meta_non_persistent_buffer_init_fn(self.dit)
+            self.dit = meta_non_persistent_buffer_init_fn(self.dit)
 
-        # if device in [get_device(), "cuda"]:
-        self.dit.to("cuda")
+        if device in [get_device(), "cuda"]:
+            self.dit.to(get_device())
 
         # Print model size.
         num_params = sum(p.numel() for p in self.dit.parameters() if p.requires_grad)
         print(f"DiT trainable parameters: {num_params:,}")
 
-    # @log_on_entry
-    # @log_runtime
+    @log_on_entry
+    @log_runtime
     def configure_vae_model(self):
         # Create vae model.
         dtype = getattr(torch, self.config.vae.dtype)
         self.vae = create_object(self.config.vae.model)
         self.vae.requires_grad_(False).eval()
-        self.vae.to(device="cuda", dtype=dtype)
+        self.vae.to(device=get_device(), dtype=dtype)
 
         # Load vae checkpoint.
         state = torch.load(
-            self.config.vae.checkpoint, map_location="cuda", mmap=True
+            self.config.vae.checkpoint, map_location=get_device(), mmap=True
         )
         self.vae.load_state_dict(state)
 
@@ -123,12 +123,12 @@ class VideoDiffusionInfer():
     def configure_diffusion(self):
         self.schedule = create_schedule_from_config(
             config=self.config.diffusion.schedule,
-            device="cuda",
+            device=get_device(),
         )
         self.sampling_timesteps = create_sampling_timesteps_from_config(
             config=self.config.diffusion.timesteps.sampling,
             schedule=self.schedule,
-            device="cuda",
+            device=get_device(),
         )
         self.sampler = create_sampler_from_config(
             config=self.config.diffusion.sampler,
@@ -143,7 +143,7 @@ class VideoDiffusionInfer():
         use_sample = self.config.vae.get("use_sample", True)
         latents = []
         if len(samples) > 0:
-            device = "cuda"
+            device = get_device()
             dtype = getattr(torch, self.config.vae.dtype)
             scale = self.config.vae.scaling_factor
             shift = self.config.vae.get("shifting_factor", 0.0)
@@ -186,7 +186,7 @@ class VideoDiffusionInfer():
     def vae_decode(self, latents: List[Tensor]) -> List[Tensor]:
         samples = []
         if len(latents) > 0:
-            device = "cuda"
+            device = get_device()
             dtype = getattr(torch, self.config.vae.dtype)
             scale = self.config.vae.scaling_factor
             shift = self.config.vae.get("shifting_factor", 0.0)
@@ -340,9 +340,9 @@ class VideoDiffusionInfer():
             self.dit.to("cpu")
 
         # Vae decode.
-        self.vae.to("cuda")
+        self.vae.to(get_device())
         samples = self.vae_decode(latents)
 
         if dit_offload:
-            self.dit.to("cuda")
+            self.dit.to(get_device())
         return samples