update

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
@@ -69,7 +69,7 @@ class VideoDiffusionInfer():
         raise NotImplementedError
 
     # @log_on_entry
-    @log_runtime
+    # @log_runtime
     def configure_dit_model(self, device="cpu", checkpoint=None):
         # Load dit checkpoint.
         # For fast init & resume,
@@ -90,7 +90,7 @@ 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")
@@ -100,7 +100,7 @@ class VideoDiffusionInfer():
         print(f"DiT trainable parameters: {num_params:,}")
 
     # @log_on_entry
-    @log_runtime
+    # @log_runtime
     def configure_vae_model(self):
         # Create vae model.
         dtype = getattr(torch, self.config.vae.dtype)