Update skyreelsinfer/pipelines/pipeline_skyreels_video.py

skyreelsinfer/pipelines/pipeline_skyreels_video.py

@@ -240,7 +240,8 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
             batch_size = len(prompt)
         else:
             batch_size = prompt_embeds.shape[0]
-        self.text_encoder.to("cuda")
+        if self.text_encoder.device.type == 'cpu':
+            self.text_encoder.to("cuda")
 
         # 3. Encode input prompt
         (
@@ -341,7 +342,7 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
         self.text_encoder.to("cpu")
         self.vae.to("cpu")
         torch.cuda.empty_cache()
-        
+        torch.cuda.reset_peak_memory_stats()   
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 if self.interrupt:
@@ -414,7 +415,8 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
                     progress_bar.update()
 
         if not output_type == "latent":
-            self.vae.to("cuda")
+            if self.vae.device.type == 'cpu':
+                self.vae.to("cuda")
             latents = latents.to(self.vae.dtype) / self.vae.config.scaling_factor
             video = self.vae.decode(latents, return_dict=False)[0]
             video = self.video_processor.postprocess_video(video, output_type=output_type)