Update skyreelsinfer/pipelines/pipeline_skyreels_video.py

skyreelsinfer/pipelines/pipeline_skyreels_video.py

@@ -7,12 +7,12 @@ from typing import Union
 
 import numpy as np
 import torch
-from diffusers import HunyuanSkyreelsImageToVideoPipeline
-from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_skyreels_image2video import DEFAULT_PROMPT_TEMPLATE
-from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_skyreels_image2video import HunyuanVideoPipelineOutput
-from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_skyreels_image2video import MultiPipelineCallbacks
-from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_skyreels_image2video import PipelineCallback
-from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_skyreels_image2video import retrieve_timesteps
+from diffusers import HunyuanVideoPipeline
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import DEFAULT_PROMPT_TEMPLATE
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import HunyuanVideoPipelineOutput
+from diffusers.pipelines.hunyuan_video.pipeline_hunyuan_video import MultiPipelineCallbacks
+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
 
@@ -46,7 +46,7 @@ def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0):
     return noise_cfg
 
 
-class SkyreelsVideoPipeline(HunyuanSkyreelsImageToVideoPipeline):
+class SkyreelsVideoPipeline(HunyuanVideoPipeline):
     """
     support i2v and t2v
     support true_cfg
@@ -300,8 +300,7 @@ class SkyreelsVideoPipeline(HunyuanSkyreelsImageToVideoPipeline):
                 device, dtype=prompt_embeds.dtype
             )
         num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
-        latents, image_latents = self.prepare_latents(
-            image,
+        latents = self.prepare_latents(
             batch_size * num_videos_per_prompt,
             num_channels_latents,
             height,
@@ -316,7 +315,7 @@ class SkyreelsVideoPipeline(HunyuanSkyreelsImageToVideoPipeline):
         self.text_encoder.to("cpu")
         torch.cuda.empty_cache()
         torch.cuda.reset_peak_memory_stats()  
-        '''
+        
         # add image latents
         if image is not None:
             image_latents = self.image_latents(
@@ -326,7 +325,7 @@ class SkyreelsVideoPipeline(HunyuanSkyreelsImageToVideoPipeline):
             image_latents = image_latents.to(transformer_dtype)
         else:
             image_latents = None
-        '''
+
         # 6. Prepare guidance condition
         if self.do_classifier_free_guidance:
             guidance = (
@@ -361,7 +360,7 @@ class SkyreelsVideoPipeline(HunyuanSkyreelsImageToVideoPipeline):
                     latent_image_input = (
                         torch.cat([image_latents] * 2) if self.do_classifier_free_guidance else image_latents
                     )
-                    latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=1).to('cuda',torch.bfloat16)
+                    latent_model_input = torch.cat([latent_model_input, latent_image_input], dim=1)
                 timestep = t.repeat(latent_model_input.shape[0]).to(torch.float32)
                 if cfg_for and self.do_classifier_free_guidance:
                     noise_pred_list = []
@@ -434,4 +433,4 @@ class SkyreelsVideoPipeline(HunyuanSkyreelsImageToVideoPipeline):
         if not return_dict:
             return (video,)
 
-        return HunyuanVideoPipelineOutput(frames=video)
+        return HunyuanVideoPipelineOutput(frames=video)