From 2f4d6caa90cc50a2685780b4812e6ca845ceee8a Mon Sep 17 00:00:00 2001 From: foreverpiano Date: Wed, 30 Oct 2024 12:43:53 +0000 Subject: [PATCH] rename --- fastvideo/sample/{mochi_pipe.py => pipeline_mochi.py} | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) rename fastvideo/sample/{mochi_pipe.py => pipeline_mochi.py} (99%) diff --git a/fastvideo/sample/mochi_pipe.py b/fastvideo/sample/pipeline_mochi.py similarity index 99% rename from fastvideo/sample/mochi_pipe.py rename to fastvideo/sample/pipeline_mochi.py index f95221662..acf4c4ab3 100644 --- a/fastvideo/sample/mochi_pipe.py +++ b/fastvideo/sample/pipeline_mochi.py @@ -712,14 +712,12 @@ class MochiPipeline(DiffusionPipeline): latents = latents / self.vae.config.scaling_factor video = self.vae.decode(latents, return_dict=False)[0] - # TODO: Change pil hardcode - video = self.video_processor.postprocess_video(video, output_type='pil') + video = self.video_processor.postprocess_video(video, output_type=output_type) # Offload all models self.maybe_free_model_hooks() - if output_type == "latent_and_video": - return video, latents, prompt_embeds, prompt_attention_mask - + if not return_dict: return (video,) - return MochiPipelineOutput(frames=video) \ No newline at end of file + + return MochiPipelineOutput(frames=video)