diff --git a/examples/inference/basic/basic_causal.py b/examples/inference/basic/basic_self_forcing_causal.py similarity index 100% rename from examples/inference/basic/basic_causal.py rename to examples/inference/basic/basic_self_forcing_causal.py diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index cd8fb9938..93628a75d 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -9,6 +9,7 @@ from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig from fastvideo.configs.pipelines.stepvideo import StepVideoT2VConfig from fastvideo.configs.pipelines.wan import (FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, + SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig) from fastvideo.logger import init_logger @@ -34,6 +35,7 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "Wan-AI/Wan2.2-TI2V-5B-Diffusers": WanT2V720PConfig, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": WanT2V480PConfig, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": WanI2V480PConfig, + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig, # Add other specific weight variants } diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index 317127f66..579da6dfd 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -138,3 +138,11 @@ class Wan2_2_T2V_A14B_Config(WanT2V480PConfig): @dataclass class Wan2_2_I2V_A14B_Config(WanT2V480PConfig): pass + + +# ============================================= +# ============= Causal Self-Forcing ============= +# ============================================= +@dataclass +class SelfForcingWanT2V480PConfig(WanT2V480PConfig): + is_causal: bool = True diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index 740cbbde9..013c60ac8 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -17,6 +17,7 @@ from fastvideo.configs.sample.wan import ( WanI2V_14B_720P_SamplingParam, WanT2V_1_3B_SamplingParam, WanT2V_14B_SamplingParam, + SelfForcingWanT2V480PConfig, ) # isort: on from fastvideo.logger import init_logger @@ -28,16 +29,27 @@ logger = init_logger(__name__) SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, + + # Wan2.1 "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, - "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, + + # Wan2.2 "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, - "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, + + # FastWan2.1 + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480PConfig, + + # FastWan2.2 + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, + + # Causal Self-Forcing Wan2.1 + "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": SelfForcingWanT2V480PConfig, # Add other specific weight variants } diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 23cb0cdbf..4c435fdc8 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -141,3 +141,12 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParam): guidance_scale_2: float = 3.5 num_inference_steps: int = 40 fps: int = 16 + + +# ============================================= +# ============= Causal Self-Forcing ============= +# ============================================= +@dataclass +class SelfForcingWanT2V480PConfig(WanT2V_1_3B_SamplingParam): + num_frames_per_block: int = 3 + sliding_window_num_frames: int = 21 diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 117b1a728..72763c18b 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/fastvideo/models/loader/fsdp_load.py @@ -248,6 +248,7 @@ def load_model_from_full_model_state_dict( sharded_sd = {} custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict( full_sd_iterator, param_names_mapping) # type: ignore + print(custom_param_sd.keys()) for target_param_name, full_tensor in custom_param_sd.items(): meta_sharded_param = meta_sd.get(target_param_name) if meta_sharded_param is None: diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 5c585fa16..68be9bc13 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -25,12 +25,14 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoTransformer3DModel"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), + "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "StepVideoModel": ("dits", "stepvideo", "StepVideoModel") } _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), + "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), } _TEXT_ENCODER_MODELS = { diff --git a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py deleted file mode 100644 index 35983417c..000000000 --- a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py +++ /dev/null @@ -1,71 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -""" -Wan causal pipeline implementation. - -This module contains an implementation of the Wan causal pipeline -using the modular pipeline architecture. -""" - -from fastvideo.fastvideo_args import FastVideoArgs -from fastvideo.logger import init_logger -from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( - FlowMatchEulerDiscreteScheduler) -from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline - -# isort: off -from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, - CausalDenoisingStage, - InputValidationStage, - LatentPreparationStage, - TextEncodingStage, - TimestepPreparationStage) -# isort: on - -logger = init_logger(__name__) - - -class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase): - - _required_config_modules = [ - "text_encoder", "tokenizer", "vae", "transformer", "scheduler" - ] - - def initialize_pipeline(self, fastvideo_args: FastVideoArgs): - - self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( - shift=fastvideo_args.pipeline_config.flow_shift) - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - """Set up pipeline stages with proper dependency injection.""" - - self.add_stage(stage_name="input_validation_stage", - stage=InputValidationStage()) - - self.add_stage(stage_name="prompt_encoding_stage", - stage=TextEncodingStage( - text_encoders=[self.get_module("text_encoder")], - tokenizers=[self.get_module("tokenizer")], - )) - - self.add_stage(stage_name="conditioning_stage", - stage=ConditioningStage()) - - self.add_stage(stage_name="timestep_preparation_stage", - stage=TimestepPreparationStage( - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="latent_preparation_stage", - stage=LatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer", None))) - - self.add_stage(stage_name="denoising_stage", - stage=CausalDenoisingStage( - transformer=self.get_module("transformer"), - scheduler=self.get_module("scheduler"))) - - self.add_stage(stage_name="decoding_stage", - stage=DecodingStage(vae=self.get_module("vae"))) - - -EntryClass = WanCausalPipeline diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index d80fd40b5..9955c3729 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -21,6 +21,7 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "WanPipeline": "wan", "WanDMDPipeline": "wan", "WanImageToVideoPipeline": "wan", + "WanCausalDMDPipeline": "wan", "StepVideoPipeline": "stepvideo", "HunyuanVideoPipeline": "hunyuan", } diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 29571910b..9566b426f 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -7,12 +7,11 @@ complete diffusion pipelines. """ from fastvideo.pipelines.stages.base import PipelineStage +from fastvideo.pipelines.stages.causal_denoising import CausalDMDDenosingStage from fastvideo.pipelines.stages.conditioning import ConditioningStage from fastvideo.pipelines.stages.decoding import DecodingStage from fastvideo.pipelines.stages.denoising import (DenoisingStage, DmdDenoisingStage) -from fastvideo.pipelines.stages.causal_denoising import ( - CausalDMDDenosingStage, CausalDenoisingStage) from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, ImageVAEEncodingStage) @@ -32,7 +31,6 @@ __all__ = [ "ConditioningStage", "DenoisingStage", "DmdDenoisingStage", - "CausalDenoisingStage", "CausalDMDDenosingStage", "EncodingStage", "DecodingStage", diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index 11407039e..1fec923db 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -1,4 +1,5 @@ from huggingface_hub import save_torch_state_dict, load_state_dict_from_file +# from safetensors import safetensors from safetensors.torch import save_file import torch import re @@ -67,7 +68,8 @@ _param_names_mapping: dict = { r"blocks.\1.norm2.\2", } -state_dict = load_state_dict_from_file("FastVideo/diffusion_pytorch_model.safetensors") +state_dict = load_state_dict_from_file("checkpoints/self_forcing_dmd.pt") +# state_dict = state_dict["generator_ema"] new_state_dict = OrderedDict() for k, v in state_dict.items(): new_key = k @@ -84,5 +86,6 @@ del state_dict save_torch_state_dict( new_state_dict, - "Wan2.1-Fun-1.3B-InP-Diffusers/transformer" + "new_diffusers/", + max_shard_size="10GB" ) \ No newline at end of file diff --git a/scripts/huggingface/upload_hf.py b/scripts/huggingface/upload_hf.py index 1e4ce88b7..9bbad77f1 100644 --- a/scripts/huggingface/upload_hf.py +++ b/scripts/huggingface/upload_hf.py @@ -3,7 +3,7 @@ from huggingface_hub import HfApi api = HfApi() api.upload_folder( - folder_path="Wan2.2-TI2V-5B-Diffusers", - repo_id="FastVideo/FastWan2.2-TI2V-5B-Diffusers", + folder_path="wow", + repo_id="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers", repo_type="model", )