update

projects/video_diffusion_sr/infer.py

@@ -123,12 +123,12 @@ class VideoDiffusionInfer():
     def configure_diffusion(self):
         self.schedule = create_schedule_from_config(
             config=self.config.diffusion.schedule,
-            device=get_device(),
+            device="cuda",
         )
         self.sampling_timesteps = create_sampling_timesteps_from_config(
             config=self.config.diffusion.timesteps.sampling,
             schedule=self.schedule,
-            device=get_device(),
+            device="cuda",
         )
         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 = get_device()
+            device = "cuda"
             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 = get_device()
+            device = "cuda"
             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(get_device())
+        self.vae.to("cuda")
         samples = self.vae_decode(latents)
 
         if dit_offload:
-            self.dit.to(get_device())
+            self.dit.to("cuda")
         return samples