Update skyreelsinfer/pipelines/pipeline_skyreels_video.py

skyreelsinfer/pipelines/pipeline_skyreels_video.py

@@ -14,7 +14,7 @@ from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import MultiPipeli
 from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import PipelineCallback
 from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import retrieve_timesteps
 from PIL import Image
-#import gc
+import gc
 
 def resizecrop(image, th, tw):
     w, h = image.size
@@ -241,8 +241,11 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
         else:
             batch_size = prompt_embeds.shape[0]
         if self.text_encoder.device.type == 'cpu':
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats() 
+            self.vae.to("cuda")
             self.text_encoder.to("cuda")
-
+            gc.collect()
         # 3. Encode input prompt
         (
             prompt_embeds,
@@ -313,6 +316,7 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
         )
         
         self.text_encoder.to("cpu")
+        gc.collect()
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()  
         
@@ -345,6 +349,7 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
         if hasattr(self, "text_encoder_to_cpu"):
             self.text_encoder_to_cpu()
         self.vae.to("cpu")
+        gc.collect()
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()   
         with self.progress_bar(total=num_inference_steps) as progress_bar:
@@ -420,7 +425,10 @@ class SkyreelsVideoPipeline(HunyuanVideoPipeline):
 
         if not output_type == "latent":
             if self.vae.device.type == 'cpu':
+                torch.cuda.empty_cache()
+                torch.cuda.reset_peak_memory_stats() 
                 self.vae.to("cuda")
+                gc.collect()
             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)