Initial Lip Forcing 14B streaming demo

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/example1_audio.wav filter=lfs diff=lfs merge=lfs -text
+examples/example1_video.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/example2_audio.wav filter=lfs diff=lfs merge=lfs -text
+examples/example2_video.mp4 filter=lfs diff=lfs merge=lfs -text

OmniAvatar/configs/model_config.py

@@ -0,0 +1,664 @@
+from typing_extensions import Literal, TypeAlias
+from ..models.wan_video_dit import WanModel
+from ..models.wan_video_text_encoder import WanTextEncoder
+from ..models.wan_video_vae import WanVideoVAE
+
+
+model_loader_configs = [
+    # These configs are provided for detecting model type automatically.
+    # The format is (state_dict_keys_hash, state_dict_keys_hash_with_shape, model_names, model_classes, model_resource)
+    (None, "9269f8db9040a9d860eaca435be61814", ["wan_video_dit"], [WanModel], "civitai"),
+    (None, "aafcfd9672c3a2456dc46e1cb6e52c70", ["wan_video_dit"], [WanModel], "civitai"),
+    (None, "6bfcfb3b342cb286ce886889d519a77e", ["wan_video_dit"], [WanModel], "civitai"),
+    (None, "cb104773c6c2cb6df4f9529ad5c60d0b", ["wan_video_dit"], [WanModel], "diffusers"),
+    (None, "9c8818c2cbea55eca56c7b447df170da", ["wan_video_text_encoder"], [WanTextEncoder], "civitai"),
+    (None, "1378ea763357eea97acdef78e65d6d96", ["wan_video_vae"], [WanVideoVAE], "civitai"),
+    (None, "ccc42284ea13e1ad04693284c7a09be6", ["wan_video_vae"], [WanVideoVAE], "civitai"),
+]
+huggingface_model_loader_configs = [
+    # These configs are provided for detecting model type automatically.
+    # The format is (architecture_in_huggingface_config, huggingface_lib, model_name, redirected_architecture)
+    ("ChatGLMModel", "diffsynth.models.kolors_text_encoder", "kolors_text_encoder", None),
+    ("MarianMTModel", "transformers.models.marian.modeling_marian", "translator", None),
+    ("BloomForCausalLM", "transformers.models.bloom.modeling_bloom", "beautiful_prompt", None),
+    ("Qwen2ForCausalLM", "transformers.models.qwen2.modeling_qwen2", "qwen_prompt", None),
+    # ("LlamaForCausalLM", "transformers.models.llama.modeling_llama", "omost_prompt", None),
+    ("T5EncoderModel", "diffsynth.models.flux_text_encoder", "flux_text_encoder_2", "FluxTextEncoder2"),
+    ("CogVideoXTransformer3DModel", "diffsynth.models.cog_dit", "cog_dit", "CogDiT"),
+    ("SiglipModel", "transformers.models.siglip.modeling_siglip", "siglip_vision_model", "SiglipVisionModel"),
+    ("LlamaForCausalLM", "diffsynth.models.hunyuan_video_text_encoder", "hunyuan_video_text_encoder_2", "HunyuanVideoLLMEncoder"),
+    ("LlavaForConditionalGeneration", "diffsynth.models.hunyuan_video_text_encoder", "hunyuan_video_text_encoder_2", "HunyuanVideoMLLMEncoder"),
+    ("Step1Model", "diffsynth.models.stepvideo_text_encoder", "stepvideo_text_encoder_2", "STEP1TextEncoder"),
+]
+
+preset_models_on_huggingface = {
+    "HunyuanDiT": [
+        ("Tencent-Hunyuan/HunyuanDiT", "t2i/clip_text_encoder/pytorch_model.bin", "models/HunyuanDiT/t2i/clip_text_encoder"),
+        ("Tencent-Hunyuan/HunyuanDiT", "t2i/mt5/pytorch_model.bin", "models/HunyuanDiT/t2i/mt5"),
+        ("Tencent-Hunyuan/HunyuanDiT", "t2i/model/pytorch_model_ema.pt", "models/HunyuanDiT/t2i/model"),
+        ("Tencent-Hunyuan/HunyuanDiT", "t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.bin", "models/HunyuanDiT/t2i/sdxl-vae-fp16-fix"),
+    ],
+    "stable-video-diffusion-img2vid-xt": [
+        ("stabilityai/stable-video-diffusion-img2vid-xt", "svd_xt.safetensors", "models/stable_video_diffusion"),
+    ],
+    "ExVideo-SVD-128f-v1": [
+        ("ECNU-CILab/ExVideo-SVD-128f-v1", "model.fp16.safetensors", "models/stable_video_diffusion"),
+    ],
+    # Stable Diffusion
+    "StableDiffusion_v15": [
+        ("benjamin-paine/stable-diffusion-v1-5", "v1-5-pruned-emaonly.safetensors", "models/stable_diffusion"),
+    ],
+    "DreamShaper_8": [
+        ("Yntec/Dreamshaper8", "dreamshaper_8.safetensors", "models/stable_diffusion"),
+    ],
+    # Textual Inversion
+    "TextualInversion_VeryBadImageNegative_v1.3": [
+        ("gemasai/verybadimagenegative_v1.3", "verybadimagenegative_v1.3.pt", "models/textual_inversion"),
+    ],
+    # Stable Diffusion XL
+    "StableDiffusionXL_v1": [
+        ("stabilityai/stable-diffusion-xl-base-1.0", "sd_xl_base_1.0.safetensors", "models/stable_diffusion_xl"),
+    ],
+    "BluePencilXL_v200": [
+        ("frankjoshua/bluePencilXL_v200", "bluePencilXL_v200.safetensors", "models/stable_diffusion_xl"),
+    ],
+    "StableDiffusionXL_Turbo": [
+        ("stabilityai/sdxl-turbo", "sd_xl_turbo_1.0_fp16.safetensors", "models/stable_diffusion_xl_turbo"),
+    ],
+    # Stable Diffusion 3
+    "StableDiffusion3": [
+        ("stabilityai/stable-diffusion-3-medium", "sd3_medium_incl_clips_t5xxlfp16.safetensors", "models/stable_diffusion_3"),
+    ],
+    "StableDiffusion3_without_T5": [
+        ("stabilityai/stable-diffusion-3-medium", "sd3_medium_incl_clips.safetensors", "models/stable_diffusion_3"),
+    ],
+    # ControlNet
+    "ControlNet_v11f1p_sd15_depth": [
+        ("lllyasviel/ControlNet-v1-1", "control_v11f1p_sd15_depth.pth", "models/ControlNet"),
+        ("lllyasviel/Annotators", "dpt_hybrid-midas-501f0c75.pt", "models/Annotators")
+    ],
+    "ControlNet_v11p_sd15_softedge": [
+        ("lllyasviel/ControlNet-v1-1", "control_v11p_sd15_softedge.pth", "models/ControlNet"),
+        ("lllyasviel/Annotators", "ControlNetHED.pth", "models/Annotators")
+    ],
+    "ControlNet_v11f1e_sd15_tile": [
+        ("lllyasviel/ControlNet-v1-1", "control_v11f1e_sd15_tile.pth", "models/ControlNet")
+    ],
+    "ControlNet_v11p_sd15_lineart": [
+        ("lllyasviel/ControlNet-v1-1", "control_v11p_sd15_lineart.pth", "models/ControlNet"),
+        ("lllyasviel/Annotators", "sk_model.pth", "models/Annotators"),
+        ("lllyasviel/Annotators", "sk_model2.pth", "models/Annotators")
+    ],
+    "ControlNet_union_sdxl_promax": [
+        ("xinsir/controlnet-union-sdxl-1.0", "diffusion_pytorch_model_promax.safetensors", "models/ControlNet/controlnet_union"),
+        ("lllyasviel/Annotators", "dpt_hybrid-midas-501f0c75.pt", "models/Annotators")
+    ],
+    # AnimateDiff
+    "AnimateDiff_v2": [
+        ("guoyww/animatediff", "mm_sd_v15_v2.ckpt", "models/AnimateDiff"),
+    ],
+    "AnimateDiff_xl_beta": [
+        ("guoyww/animatediff", "mm_sdxl_v10_beta.ckpt", "models/AnimateDiff"),
+    ],
+
+    # Qwen Prompt
+    "QwenPrompt": [
+        ("Qwen/Qwen2-1.5B-Instruct", "config.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "generation_config.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "model.safetensors", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "special_tokens_map.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "tokenizer.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "tokenizer_config.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "merges.txt", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ("Qwen/Qwen2-1.5B-Instruct", "vocab.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+    ],
+    # Beautiful Prompt
+    "BeautifulPrompt": [
+        ("alibaba-pai/pai-bloom-1b1-text2prompt-sd", "config.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+        ("alibaba-pai/pai-bloom-1b1-text2prompt-sd", "generation_config.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+        ("alibaba-pai/pai-bloom-1b1-text2prompt-sd", "model.safetensors", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+        ("alibaba-pai/pai-bloom-1b1-text2prompt-sd", "special_tokens_map.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+        ("alibaba-pai/pai-bloom-1b1-text2prompt-sd", "tokenizer.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+        ("alibaba-pai/pai-bloom-1b1-text2prompt-sd", "tokenizer_config.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+    ],
+    # Omost prompt
+    "OmostPrompt":[
+        ("lllyasviel/omost-llama-3-8b-4bits", "model-00001-of-00002.safetensors", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ("lllyasviel/omost-llama-3-8b-4bits", "model-00002-of-00002.safetensors", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ("lllyasviel/omost-llama-3-8b-4bits", "tokenizer.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ("lllyasviel/omost-llama-3-8b-4bits", "tokenizer_config.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),  
+        ("lllyasviel/omost-llama-3-8b-4bits", "config.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ("lllyasviel/omost-llama-3-8b-4bits", "generation_config.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ("lllyasviel/omost-llama-3-8b-4bits", "model.safetensors.index.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ("lllyasviel/omost-llama-3-8b-4bits", "special_tokens_map.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+    ],
+    # Translator
+    "opus-mt-zh-en": [
+        ("Helsinki-NLP/opus-mt-zh-en", "config.json", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "generation_config.json", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "metadata.json", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "pytorch_model.bin", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "source.spm", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "target.spm", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "tokenizer_config.json", "models/translator/opus-mt-zh-en"),
+        ("Helsinki-NLP/opus-mt-zh-en", "vocab.json", "models/translator/opus-mt-zh-en"),
+    ],
+    # IP-Adapter
+    "IP-Adapter-SD": [
+        ("h94/IP-Adapter", "models/image_encoder/model.safetensors", "models/IpAdapter/stable_diffusion/image_encoder"),
+        ("h94/IP-Adapter", "models/ip-adapter_sd15.bin", "models/IpAdapter/stable_diffusion"),
+    ],
+    "IP-Adapter-SDXL": [
+        ("h94/IP-Adapter", "sdxl_models/image_encoder/model.safetensors", "models/IpAdapter/stable_diffusion_xl/image_encoder"),
+        ("h94/IP-Adapter", "sdxl_models/ip-adapter_sdxl.bin", "models/IpAdapter/stable_diffusion_xl"),
+    ],
+    "SDXL-vae-fp16-fix": [
+        ("madebyollin/sdxl-vae-fp16-fix", "diffusion_pytorch_model.safetensors", "models/sdxl-vae-fp16-fix")
+    ],
+    # Kolors
+    "Kolors": [
+        ("Kwai-Kolors/Kolors", "text_encoder/config.json", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model.bin.index.json", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00001-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00002-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00003-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00004-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00005-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00006-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00007-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+        ("Kwai-Kolors/Kolors", "unet/diffusion_pytorch_model.safetensors", "models/kolors/Kolors/unet"),
+        ("Kwai-Kolors/Kolors", "vae/diffusion_pytorch_model.safetensors", "models/kolors/Kolors/vae"),
+    ],
+    # FLUX
+    "FLUX.1-dev": [
+        ("black-forest-labs/FLUX.1-dev", "text_encoder/model.safetensors", "models/FLUX/FLUX.1-dev/text_encoder"),
+        ("black-forest-labs/FLUX.1-dev", "text_encoder_2/config.json", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+        ("black-forest-labs/FLUX.1-dev", "text_encoder_2/model-00001-of-00002.safetensors", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+        ("black-forest-labs/FLUX.1-dev", "text_encoder_2/model-00002-of-00002.safetensors", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+        ("black-forest-labs/FLUX.1-dev", "text_encoder_2/model.safetensors.index.json", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+        ("black-forest-labs/FLUX.1-dev", "ae.safetensors", "models/FLUX/FLUX.1-dev"),
+        ("black-forest-labs/FLUX.1-dev", "flux1-dev.safetensors", "models/FLUX/FLUX.1-dev"),
+    ],
+    "InstantX/FLUX.1-dev-IP-Adapter": {
+        "file_list": [
+            ("InstantX/FLUX.1-dev-IP-Adapter", "ip-adapter.bin", "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter"),
+            ("google/siglip-so400m-patch14-384", "model.safetensors", "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/image_encoder"),
+            ("google/siglip-so400m-patch14-384", "config.json", "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/image_encoder"),
+        ],
+        "load_path": [
+            "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/ip-adapter.bin",
+            "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/image_encoder",
+        ],
+    },
+    # RIFE
+    "RIFE": [
+        ("AlexWortega/RIFE", "flownet.pkl", "models/RIFE"),
+    ],
+    # CogVideo
+    "CogVideoX-5B": [
+        ("THUDM/CogVideoX-5b", "text_encoder/config.json", "models/CogVideo/CogVideoX-5b/text_encoder"),
+        ("THUDM/CogVideoX-5b", "text_encoder/model.safetensors.index.json", "models/CogVideo/CogVideoX-5b/text_encoder"),
+        ("THUDM/CogVideoX-5b", "text_encoder/model-00001-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/text_encoder"),
+        ("THUDM/CogVideoX-5b", "text_encoder/model-00002-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/text_encoder"),
+        ("THUDM/CogVideoX-5b", "transformer/config.json", "models/CogVideo/CogVideoX-5b/transformer"),
+        ("THUDM/CogVideoX-5b", "transformer/diffusion_pytorch_model.safetensors.index.json", "models/CogVideo/CogVideoX-5b/transformer"),
+        ("THUDM/CogVideoX-5b", "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/transformer"),
+        ("THUDM/CogVideoX-5b", "transformer/diffusion_pytorch_model-00002-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/transformer"),
+        ("THUDM/CogVideoX-5b", "vae/diffusion_pytorch_model.safetensors", "models/CogVideo/CogVideoX-5b/vae"),
+    ],
+    # Stable Diffusion 3.5
+    "StableDiffusion3.5-large": [
+        ("stabilityai/stable-diffusion-3.5-large", "sd3.5_large.safetensors", "models/stable_diffusion_3"),
+        ("stabilityai/stable-diffusion-3.5-large", "text_encoders/clip_l.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("stabilityai/stable-diffusion-3.5-large", "text_encoders/clip_g.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("stabilityai/stable-diffusion-3.5-large", "text_encoders/t5xxl_fp16.safetensors", "models/stable_diffusion_3/text_encoders"),
+    ],
+}
+preset_models_on_modelscope = {
+    # Hunyuan DiT
+    "HunyuanDiT": [
+        ("modelscope/HunyuanDiT", "t2i/clip_text_encoder/pytorch_model.bin", "models/HunyuanDiT/t2i/clip_text_encoder"),
+        ("modelscope/HunyuanDiT", "t2i/mt5/pytorch_model.bin", "models/HunyuanDiT/t2i/mt5"),
+        ("modelscope/HunyuanDiT", "t2i/model/pytorch_model_ema.pt", "models/HunyuanDiT/t2i/model"),
+        ("modelscope/HunyuanDiT", "t2i/sdxl-vae-fp16-fix/diffusion_pytorch_model.bin", "models/HunyuanDiT/t2i/sdxl-vae-fp16-fix"),
+    ],
+    # Stable Video Diffusion
+    "stable-video-diffusion-img2vid-xt": [
+        ("AI-ModelScope/stable-video-diffusion-img2vid-xt", "svd_xt.safetensors", "models/stable_video_diffusion"),
+    ],
+    # ExVideo
+    "ExVideo-SVD-128f-v1": [
+        ("ECNU-CILab/ExVideo-SVD-128f-v1", "model.fp16.safetensors", "models/stable_video_diffusion"),
+    ],
+    "ExVideo-CogVideoX-LoRA-129f-v1": [
+        ("ECNU-CILab/ExVideo-CogVideoX-LoRA-129f-v1", "ExVideo-CogVideoX-LoRA-129f-v1.safetensors", "models/lora"),
+    ],
+    # Stable Diffusion
+    "StableDiffusion_v15": [
+        ("AI-ModelScope/stable-diffusion-v1-5", "v1-5-pruned-emaonly.safetensors", "models/stable_diffusion"),
+    ],
+    "DreamShaper_8": [
+        ("sd_lora/dreamshaper_8", "dreamshaper_8.safetensors", "models/stable_diffusion"),
+    ],
+    "AingDiffusion_v12": [
+        ("sd_lora/aingdiffusion_v12", "aingdiffusion_v12.safetensors", "models/stable_diffusion"),
+    ],
+    "Flat2DAnimerge_v45Sharp": [
+        ("sd_lora/Flat-2D-Animerge", "flat2DAnimerge_v45Sharp.safetensors", "models/stable_diffusion"),
+    ],
+    # Textual Inversion
+    "TextualInversion_VeryBadImageNegative_v1.3": [
+        ("sd_lora/verybadimagenegative_v1.3", "verybadimagenegative_v1.3.pt", "models/textual_inversion"),
+    ],
+    # Stable Diffusion XL
+    "StableDiffusionXL_v1": [
+        ("AI-ModelScope/stable-diffusion-xl-base-1.0", "sd_xl_base_1.0.safetensors", "models/stable_diffusion_xl"),
+    ],
+    "BluePencilXL_v200": [
+        ("sd_lora/bluePencilXL_v200", "bluePencilXL_v200.safetensors", "models/stable_diffusion_xl"),
+    ],
+    "StableDiffusionXL_Turbo": [
+        ("AI-ModelScope/sdxl-turbo", "sd_xl_turbo_1.0_fp16.safetensors", "models/stable_diffusion_xl_turbo"),
+    ],
+    "SDXL_lora_zyd232_ChineseInkStyle_SDXL_v1_0": [
+        ("sd_lora/zyd232_ChineseInkStyle_SDXL_v1_0", "zyd232_ChineseInkStyle_SDXL_v1_0.safetensors", "models/lora"),
+    ],
+    # Stable Diffusion 3
+    "StableDiffusion3": [
+        ("AI-ModelScope/stable-diffusion-3-medium", "sd3_medium_incl_clips_t5xxlfp16.safetensors", "models/stable_diffusion_3"),
+    ],
+    "StableDiffusion3_without_T5": [
+        ("AI-ModelScope/stable-diffusion-3-medium", "sd3_medium_incl_clips.safetensors", "models/stable_diffusion_3"),
+    ],
+    # ControlNet
+    "ControlNet_v11f1p_sd15_depth": [
+        ("AI-ModelScope/ControlNet-v1-1", "control_v11f1p_sd15_depth.pth", "models/ControlNet"),
+        ("sd_lora/Annotators", "dpt_hybrid-midas-501f0c75.pt", "models/Annotators")
+    ],
+    "ControlNet_v11p_sd15_softedge": [
+        ("AI-ModelScope/ControlNet-v1-1", "control_v11p_sd15_softedge.pth", "models/ControlNet"),
+        ("sd_lora/Annotators", "ControlNetHED.pth", "models/Annotators")
+    ],
+    "ControlNet_v11f1e_sd15_tile": [
+        ("AI-ModelScope/ControlNet-v1-1", "control_v11f1e_sd15_tile.pth", "models/ControlNet")
+    ],
+    "ControlNet_v11p_sd15_lineart": [
+        ("AI-ModelScope/ControlNet-v1-1", "control_v11p_sd15_lineart.pth", "models/ControlNet"),
+        ("sd_lora/Annotators", "sk_model.pth", "models/Annotators"),
+        ("sd_lora/Annotators", "sk_model2.pth", "models/Annotators")
+    ],
+    "ControlNet_union_sdxl_promax": [
+        ("AI-ModelScope/controlnet-union-sdxl-1.0", "diffusion_pytorch_model_promax.safetensors", "models/ControlNet/controlnet_union"),
+        ("sd_lora/Annotators", "dpt_hybrid-midas-501f0c75.pt", "models/Annotators")
+    ],
+    "Annotators:Depth": [
+        ("sd_lora/Annotators", "dpt_hybrid-midas-501f0c75.pt", "models/Annotators"),
+    ],
+    "Annotators:Softedge": [
+        ("sd_lora/Annotators", "ControlNetHED.pth", "models/Annotators"),
+    ],
+    "Annotators:Lineart": [
+        ("sd_lora/Annotators", "sk_model.pth", "models/Annotators"),
+        ("sd_lora/Annotators", "sk_model2.pth", "models/Annotators"),
+    ],
+    "Annotators:Normal": [
+        ("sd_lora/Annotators", "scannet.pt", "models/Annotators"),
+    ],
+    "Annotators:Openpose": [
+        ("sd_lora/Annotators", "body_pose_model.pth", "models/Annotators"),
+        ("sd_lora/Annotators", "facenet.pth", "models/Annotators"),
+        ("sd_lora/Annotators", "hand_pose_model.pth", "models/Annotators"),
+    ],
+    # AnimateDiff
+    "AnimateDiff_v2": [
+        ("Shanghai_AI_Laboratory/animatediff", "mm_sd_v15_v2.ckpt", "models/AnimateDiff"),
+    ],
+    "AnimateDiff_xl_beta": [
+        ("Shanghai_AI_Laboratory/animatediff", "mm_sdxl_v10_beta.ckpt", "models/AnimateDiff"),
+    ],
+    # RIFE
+    "RIFE": [
+        ("Damo_XR_Lab/cv_rife_video-frame-interpolation", "flownet.pkl", "models/RIFE"),
+    ],
+    # Qwen Prompt
+    "QwenPrompt": {
+        "file_list": [
+            ("qwen/Qwen2-1.5B-Instruct", "config.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "generation_config.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "model.safetensors", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "special_tokens_map.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "tokenizer.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "tokenizer_config.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "merges.txt", "models/QwenPrompt/qwen2-1.5b-instruct"),
+            ("qwen/Qwen2-1.5B-Instruct", "vocab.json", "models/QwenPrompt/qwen2-1.5b-instruct"),
+        ],
+        "load_path": [
+            "models/QwenPrompt/qwen2-1.5b-instruct",
+        ],
+    },
+    # Beautiful Prompt
+    "BeautifulPrompt": {
+        "file_list": [
+            ("AI-ModelScope/pai-bloom-1b1-text2prompt-sd", "config.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+            ("AI-ModelScope/pai-bloom-1b1-text2prompt-sd", "generation_config.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+            ("AI-ModelScope/pai-bloom-1b1-text2prompt-sd", "model.safetensors", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+            ("AI-ModelScope/pai-bloom-1b1-text2prompt-sd", "special_tokens_map.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+            ("AI-ModelScope/pai-bloom-1b1-text2prompt-sd", "tokenizer.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+            ("AI-ModelScope/pai-bloom-1b1-text2prompt-sd", "tokenizer_config.json", "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd"),
+        ],
+        "load_path": [
+            "models/BeautifulPrompt/pai-bloom-1b1-text2prompt-sd",
+        ],
+    },
+    # Omost prompt
+    "OmostPrompt": {
+        "file_list": [
+            ("Omost/omost-llama-3-8b-4bits", "model-00001-of-00002.safetensors", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+            ("Omost/omost-llama-3-8b-4bits", "model-00002-of-00002.safetensors", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+            ("Omost/omost-llama-3-8b-4bits", "tokenizer.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+            ("Omost/omost-llama-3-8b-4bits", "tokenizer_config.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),  
+            ("Omost/omost-llama-3-8b-4bits", "config.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+            ("Omost/omost-llama-3-8b-4bits", "generation_config.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+            ("Omost/omost-llama-3-8b-4bits", "model.safetensors.index.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+            ("Omost/omost-llama-3-8b-4bits", "special_tokens_map.json", "models/OmostPrompt/omost-llama-3-8b-4bits"),
+        ],
+        "load_path": [
+            "models/OmostPrompt/omost-llama-3-8b-4bits",
+        ],
+    },
+    # Translator
+    "opus-mt-zh-en": {
+        "file_list": [
+            ("moxying/opus-mt-zh-en", "config.json", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "generation_config.json", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "metadata.json", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "pytorch_model.bin", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "source.spm", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "target.spm", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "tokenizer_config.json", "models/translator/opus-mt-zh-en"),
+            ("moxying/opus-mt-zh-en", "vocab.json", "models/translator/opus-mt-zh-en"),
+        ],
+        "load_path": [
+            "models/translator/opus-mt-zh-en",
+        ],
+    },
+    # IP-Adapter
+    "IP-Adapter-SD": [
+        ("AI-ModelScope/IP-Adapter", "models/image_encoder/model.safetensors", "models/IpAdapter/stable_diffusion/image_encoder"),
+        ("AI-ModelScope/IP-Adapter", "models/ip-adapter_sd15.bin", "models/IpAdapter/stable_diffusion"),
+    ],
+    "IP-Adapter-SDXL": [
+        ("AI-ModelScope/IP-Adapter", "sdxl_models/image_encoder/model.safetensors", "models/IpAdapter/stable_diffusion_xl/image_encoder"),
+        ("AI-ModelScope/IP-Adapter", "sdxl_models/ip-adapter_sdxl.bin", "models/IpAdapter/stable_diffusion_xl"),
+    ],
+    # Kolors
+    "Kolors": {
+        "file_list": [
+            ("Kwai-Kolors/Kolors", "text_encoder/config.json", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model.bin.index.json", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00001-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00002-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00003-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00004-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00005-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00006-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "text_encoder/pytorch_model-00007-of-00007.bin", "models/kolors/Kolors/text_encoder"),
+            ("Kwai-Kolors/Kolors", "unet/diffusion_pytorch_model.safetensors", "models/kolors/Kolors/unet"),
+            ("Kwai-Kolors/Kolors", "vae/diffusion_pytorch_model.safetensors", "models/kolors/Kolors/vae"),
+        ],
+        "load_path": [
+            "models/kolors/Kolors/text_encoder",
+            "models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors",
+            "models/kolors/Kolors/vae/diffusion_pytorch_model.safetensors",
+        ],
+    },
+    "SDXL-vae-fp16-fix": [
+        ("AI-ModelScope/sdxl-vae-fp16-fix", "diffusion_pytorch_model.safetensors", "models/sdxl-vae-fp16-fix")
+    ],
+    # FLUX
+    "FLUX.1-dev": {
+        "file_list": [
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder/model.safetensors", "models/FLUX/FLUX.1-dev/text_encoder"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/config.json", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/model-00001-of-00002.safetensors", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/model-00002-of-00002.safetensors", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/model.safetensors.index.json", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "ae.safetensors", "models/FLUX/FLUX.1-dev"),
+            ("AI-ModelScope/FLUX.1-dev", "flux1-dev.safetensors", "models/FLUX/FLUX.1-dev"),
+        ],
+        "load_path": [
+            "models/FLUX/FLUX.1-dev/text_encoder/model.safetensors",
+            "models/FLUX/FLUX.1-dev/text_encoder_2",
+            "models/FLUX/FLUX.1-dev/ae.safetensors",
+            "models/FLUX/FLUX.1-dev/flux1-dev.safetensors"
+        ],
+    },
+    "FLUX.1-schnell": {
+        "file_list": [
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder/model.safetensors", "models/FLUX/FLUX.1-dev/text_encoder"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/config.json", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/model-00001-of-00002.safetensors", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/model-00002-of-00002.safetensors", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "text_encoder_2/model.safetensors.index.json", "models/FLUX/FLUX.1-dev/text_encoder_2"),
+            ("AI-ModelScope/FLUX.1-dev", "ae.safetensors", "models/FLUX/FLUX.1-dev"),
+            ("AI-ModelScope/FLUX.1-schnell", "flux1-schnell.safetensors", "models/FLUX/FLUX.1-schnell"),
+        ],
+        "load_path": [
+            "models/FLUX/FLUX.1-dev/text_encoder/model.safetensors",
+            "models/FLUX/FLUX.1-dev/text_encoder_2",
+            "models/FLUX/FLUX.1-dev/ae.safetensors",
+            "models/FLUX/FLUX.1-schnell/flux1-schnell.safetensors"
+        ],
+    },
+    "InstantX/FLUX.1-dev-Controlnet-Union-alpha": [
+        ("InstantX/FLUX.1-dev-Controlnet-Union-alpha", "diffusion_pytorch_model.safetensors", "models/ControlNet/InstantX/FLUX.1-dev-Controlnet-Union-alpha"),
+    ],
+    "jasperai/Flux.1-dev-Controlnet-Depth": [
+        ("jasperai/Flux.1-dev-Controlnet-Depth", "diffusion_pytorch_model.safetensors", "models/ControlNet/jasperai/Flux.1-dev-Controlnet-Depth"),
+    ],
+    "jasperai/Flux.1-dev-Controlnet-Surface-Normals": [
+        ("jasperai/Flux.1-dev-Controlnet-Surface-Normals", "diffusion_pytorch_model.safetensors", "models/ControlNet/jasperai/Flux.1-dev-Controlnet-Surface-Normals"),
+    ],
+    "jasperai/Flux.1-dev-Controlnet-Upscaler": [
+        ("jasperai/Flux.1-dev-Controlnet-Upscaler", "diffusion_pytorch_model.safetensors", "models/ControlNet/jasperai/Flux.1-dev-Controlnet-Upscaler"),
+    ],
+    "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha": [
+        ("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha", "diffusion_pytorch_model.safetensors", "models/ControlNet/alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha"),
+    ],
+    "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta": [
+        ("alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta", "diffusion_pytorch_model.safetensors", "models/ControlNet/alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta"),
+    ],
+    "Shakker-Labs/FLUX.1-dev-ControlNet-Depth": [
+        ("Shakker-Labs/FLUX.1-dev-ControlNet-Depth", "diffusion_pytorch_model.safetensors", "models/ControlNet/Shakker-Labs/FLUX.1-dev-ControlNet-Depth"),
+    ],
+    "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro": [
+        ("Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro", "diffusion_pytorch_model.safetensors", "models/ControlNet/Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro"),
+    ],
+    "InstantX/FLUX.1-dev-IP-Adapter": {
+        "file_list": [
+            ("InstantX/FLUX.1-dev-IP-Adapter", "ip-adapter.bin", "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter"),
+            ("AI-ModelScope/siglip-so400m-patch14-384", "model.safetensors", "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/image_encoder"),
+            ("AI-ModelScope/siglip-so400m-patch14-384", "config.json", "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/image_encoder"),
+        ],
+        "load_path": [
+            "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/ip-adapter.bin",
+            "models/IpAdapter/InstantX/FLUX.1-dev-IP-Adapter/image_encoder",
+        ],
+    },
+    # ESRGAN
+    "ESRGAN_x4": [
+        ("AI-ModelScope/Real-ESRGAN", "RealESRGAN_x4.pth", "models/ESRGAN"),
+    ],
+    # RIFE
+    "RIFE": [
+        ("AI-ModelScope/RIFE", "flownet.pkl", "models/RIFE"),
+    ],
+    # Omnigen
+    "OmniGen-v1": {
+        "file_list": [
+            ("BAAI/OmniGen-v1", "vae/diffusion_pytorch_model.safetensors", "models/OmniGen/OmniGen-v1/vae"),
+            ("BAAI/OmniGen-v1", "model.safetensors", "models/OmniGen/OmniGen-v1"),
+            ("BAAI/OmniGen-v1", "config.json", "models/OmniGen/OmniGen-v1"),
+            ("BAAI/OmniGen-v1", "special_tokens_map.json", "models/OmniGen/OmniGen-v1"),
+            ("BAAI/OmniGen-v1", "tokenizer_config.json", "models/OmniGen/OmniGen-v1"),
+            ("BAAI/OmniGen-v1", "tokenizer.json", "models/OmniGen/OmniGen-v1"),
+        ],
+        "load_path": [
+            "models/OmniGen/OmniGen-v1/vae/diffusion_pytorch_model.safetensors",
+            "models/OmniGen/OmniGen-v1/model.safetensors",
+        ]
+    },
+    # CogVideo
+    "CogVideoX-5B": {
+        "file_list": [
+            ("ZhipuAI/CogVideoX-5b", "text_encoder/config.json", "models/CogVideo/CogVideoX-5b/text_encoder"),
+            ("ZhipuAI/CogVideoX-5b", "text_encoder/model.safetensors.index.json", "models/CogVideo/CogVideoX-5b/text_encoder"),
+            ("ZhipuAI/CogVideoX-5b", "text_encoder/model-00001-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/text_encoder"),
+            ("ZhipuAI/CogVideoX-5b", "text_encoder/model-00002-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/text_encoder"),
+            ("ZhipuAI/CogVideoX-5b", "transformer/config.json", "models/CogVideo/CogVideoX-5b/transformer"),
+            ("ZhipuAI/CogVideoX-5b", "transformer/diffusion_pytorch_model.safetensors.index.json", "models/CogVideo/CogVideoX-5b/transformer"),
+            ("ZhipuAI/CogVideoX-5b", "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/transformer"),
+            ("ZhipuAI/CogVideoX-5b", "transformer/diffusion_pytorch_model-00002-of-00002.safetensors", "models/CogVideo/CogVideoX-5b/transformer"),
+            ("ZhipuAI/CogVideoX-5b", "vae/diffusion_pytorch_model.safetensors", "models/CogVideo/CogVideoX-5b/vae"),
+        ],
+        "load_path": [
+            "models/CogVideo/CogVideoX-5b/text_encoder",
+            "models/CogVideo/CogVideoX-5b/transformer",
+            "models/CogVideo/CogVideoX-5b/vae/diffusion_pytorch_model.safetensors",
+        ],
+    },
+    # Stable Diffusion 3.5
+    "StableDiffusion3.5-large": [
+        ("AI-ModelScope/stable-diffusion-3.5-large", "sd3.5_large.safetensors", "models/stable_diffusion_3"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/clip_l.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/clip_g.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/t5xxl_fp16.safetensors", "models/stable_diffusion_3/text_encoders"),
+    ],
+    "StableDiffusion3.5-medium": [
+        ("AI-ModelScope/stable-diffusion-3.5-medium", "sd3.5_medium.safetensors", "models/stable_diffusion_3"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/clip_l.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/clip_g.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/t5xxl_fp16.safetensors", "models/stable_diffusion_3/text_encoders"),
+    ],
+    "StableDiffusion3.5-large-turbo": [
+        ("AI-ModelScope/stable-diffusion-3.5-large-turbo", "sd3.5_large_turbo.safetensors", "models/stable_diffusion_3"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/clip_l.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/clip_g.safetensors", "models/stable_diffusion_3/text_encoders"),
+        ("AI-ModelScope/stable-diffusion-3.5-large", "text_encoders/t5xxl_fp16.safetensors", "models/stable_diffusion_3/text_encoders"),
+    ],
+    "HunyuanVideo":{
+        "file_list": [
+            ("AI-ModelScope/clip-vit-large-patch14", "model.safetensors", "models/HunyuanVideo/text_encoder"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00001-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00002-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00003-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00004-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "config.json", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model.safetensors.index.json", "models/HunyuanVideo/text_encoder_2"),
+            ("AI-ModelScope/HunyuanVideo", "hunyuan-video-t2v-720p/vae/pytorch_model.pt", "models/HunyuanVideo/vae"),
+            ("AI-ModelScope/HunyuanVideo", "hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt", "models/HunyuanVideo/transformers")
+        ],
+        "load_path": [
+            "models/HunyuanVideo/text_encoder/model.safetensors",
+            "models/HunyuanVideo/text_encoder_2",
+            "models/HunyuanVideo/vae/pytorch_model.pt",
+            "models/HunyuanVideo/transformers/mp_rank_00_model_states.pt"
+        ],
+    },
+    "HunyuanVideoI2V":{
+        "file_list": [
+            ("AI-ModelScope/clip-vit-large-patch14", "model.safetensors", "models/HunyuanVideoI2V/text_encoder"),
+            ("AI-ModelScope/llava-llama-3-8b-v1_1-transformers", "model-00001-of-00004.safetensors", "models/HunyuanVideoI2V/text_encoder_2"),
+            ("AI-ModelScope/llava-llama-3-8b-v1_1-transformers", "model-00002-of-00004.safetensors", "models/HunyuanVideoI2V/text_encoder_2"),
+            ("AI-ModelScope/llava-llama-3-8b-v1_1-transformers", "model-00003-of-00004.safetensors", "models/HunyuanVideoI2V/text_encoder_2"),
+            ("AI-ModelScope/llava-llama-3-8b-v1_1-transformers", "model-00004-of-00004.safetensors", "models/HunyuanVideoI2V/text_encoder_2"),
+            ("AI-ModelScope/llava-llama-3-8b-v1_1-transformers", "config.json", "models/HunyuanVideoI2V/text_encoder_2"),
+            ("AI-ModelScope/llava-llama-3-8b-v1_1-transformers", "model.safetensors.index.json", "models/HunyuanVideoI2V/text_encoder_2"),
+            ("AI-ModelScope/HunyuanVideo-I2V", "hunyuan-video-i2v-720p/vae/pytorch_model.pt", "models/HunyuanVideoI2V/vae"),
+            ("AI-ModelScope/HunyuanVideo-I2V", "hunyuan-video-i2v-720p/transformers/mp_rank_00_model_states.pt", "models/HunyuanVideoI2V/transformers")
+        ],
+        "load_path": [
+            "models/HunyuanVideoI2V/text_encoder/model.safetensors",
+            "models/HunyuanVideoI2V/text_encoder_2",
+            "models/HunyuanVideoI2V/vae/pytorch_model.pt",
+            "models/HunyuanVideoI2V/transformers/mp_rank_00_model_states.pt"
+        ],
+    },
+    "HunyuanVideo-fp8":{
+        "file_list": [
+            ("AI-ModelScope/clip-vit-large-patch14", "model.safetensors", "models/HunyuanVideo/text_encoder"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00001-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00002-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00003-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model-00004-of-00004.safetensors", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "config.json", "models/HunyuanVideo/text_encoder_2"),
+            ("DiffSynth-Studio/HunyuanVideo_MLLM_text_encoder", "model.safetensors.index.json", "models/HunyuanVideo/text_encoder_2"),
+            ("AI-ModelScope/HunyuanVideo", "hunyuan-video-t2v-720p/vae/pytorch_model.pt", "models/HunyuanVideo/vae"),
+            ("DiffSynth-Studio/HunyuanVideo-safetensors", "model.fp8.safetensors", "models/HunyuanVideo/transformers")
+        ],
+        "load_path": [
+            "models/HunyuanVideo/text_encoder/model.safetensors",
+            "models/HunyuanVideo/text_encoder_2",
+            "models/HunyuanVideo/vae/pytorch_model.pt",
+            "models/HunyuanVideo/transformers/model.fp8.safetensors"
+        ],
+    },
+}
+Preset_model_id: TypeAlias = Literal[
+    "HunyuanDiT",
+    "stable-video-diffusion-img2vid-xt",
+    "ExVideo-SVD-128f-v1",
+    "ExVideo-CogVideoX-LoRA-129f-v1",
+    "StableDiffusion_v15",
+    "DreamShaper_8",
+    "AingDiffusion_v12",
+    "Flat2DAnimerge_v45Sharp",
+    "TextualInversion_VeryBadImageNegative_v1.3",
+    "StableDiffusionXL_v1",
+    "BluePencilXL_v200",
+    "StableDiffusionXL_Turbo",
+    "ControlNet_v11f1p_sd15_depth",
+    "ControlNet_v11p_sd15_softedge",
+    "ControlNet_v11f1e_sd15_tile",
+    "ControlNet_v11p_sd15_lineart",
+    "AnimateDiff_v2",
+    "AnimateDiff_xl_beta",
+    "RIFE",
+    "BeautifulPrompt",
+    "opus-mt-zh-en",
+    "IP-Adapter-SD",
+    "IP-Adapter-SDXL",
+    "StableDiffusion3",
+    "StableDiffusion3_without_T5",
+    "Kolors",
+    "SDXL-vae-fp16-fix",
+    "ControlNet_union_sdxl_promax",
+    "FLUX.1-dev",
+    "FLUX.1-schnell",
+    "InstantX/FLUX.1-dev-Controlnet-Union-alpha",
+    "jasperai/Flux.1-dev-Controlnet-Depth",
+    "jasperai/Flux.1-dev-Controlnet-Surface-Normals",
+    "jasperai/Flux.1-dev-Controlnet-Upscaler",
+    "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Alpha",
+    "alimama-creative/FLUX.1-dev-Controlnet-Inpainting-Beta",
+    "Shakker-Labs/FLUX.1-dev-ControlNet-Depth",
+    "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro",
+    "InstantX/FLUX.1-dev-IP-Adapter",
+    "SDXL_lora_zyd232_ChineseInkStyle_SDXL_v1_0",
+    "QwenPrompt",
+    "OmostPrompt",
+    "ESRGAN_x4",
+    "RIFE",
+    "OmniGen-v1",
+    "CogVideoX-5B",
+    "Annotators:Depth",
+    "Annotators:Softedge",
+    "Annotators:Lineart",
+    "Annotators:Normal",
+    "Annotators:Openpose",
+    "StableDiffusion3.5-large",
+    "StableDiffusion3.5-medium",
+    "HunyuanVideo",
+    "HunyuanVideo-fp8",
+    "HunyuanVideoI2V",
+]

OmniAvatar/models/audio_pack.py

@@ -0,0 +1,40 @@
+import torch
+from typing import Tuple, Union
+import torch
+from einops import rearrange
+from torch import nn
+
+
+def make_triple(value: Union[int, Tuple[int, int, int]]) -> Tuple[int, int, int]:
+    value = (value,) * 3 if isinstance(value, int) else value
+    assert len(value) == 3
+    return value
+
+
+class AudioPack(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            patch_size: Union[int, Tuple[int, int, int]],
+            dim: int,
+            layernorm=False,
+    ):
+        super().__init__()
+        t, h, w = make_triple(patch_size)
+        self.patch_size = t, h, w
+        self.proj = nn.Linear(in_channels * t * h * w, dim)
+        if layernorm:
+            self.norm_out = nn.LayerNorm(dim)
+        else:
+            self.norm_out = None
+
+    def forward(
+            self,
+            vid: torch.Tensor,
+    ) -> torch.Tensor:
+        t, h, w = self.patch_size
+        vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w)
+        vid = self.proj(vid)
+        if self.norm_out is not None:
+            vid = self.norm_out(vid)
+        return vid

OmniAvatar/models/model_manager.py

@@ -0,0 +1,444 @@
+import os, torch, json, importlib
+from typing import List
+import torch.nn as nn
+from ..configs.model_config import model_loader_configs, huggingface_model_loader_configs
+from ..utils.io_utils import load_state_dict, init_weights_on_device, hash_state_dict_keys, split_state_dict_with_prefix, smart_load_weights
+
+
+def load_model_from_single_file(state_dict, model_names, model_classes, model_resource, torch_dtype, device, infer):
+    loaded_model_names, loaded_models = [], []
+    for model_name, model_class in zip(model_names, model_classes):
+        print(f"    model_name: {model_name} model_class: {model_class.__name__}")
+        state_dict_converter = model_class.state_dict_converter()
+        if model_resource == "civitai":
+            state_dict_results = state_dict_converter.from_civitai(state_dict)
+        elif model_resource == "diffusers":
+            state_dict_results = state_dict_converter.from_diffusers(state_dict)
+        if isinstance(state_dict_results, tuple):
+            model_state_dict, extra_kwargs = state_dict_results
+            print(f"        This model is initialized with extra kwargs: {extra_kwargs}")
+        else:
+            model_state_dict, extra_kwargs = state_dict_results, {}
+        torch_dtype = torch.float32 if extra_kwargs.get("upcast_to_float32", False) else torch_dtype
+        with init_weights_on_device():
+            model = model_class(**extra_kwargs)
+        if hasattr(model, "eval"):
+            model = model.eval()
+        if not infer: # 训练才初始化
+            model = model.to_empty(device=torch.device("cpu"))
+            for name, param in model.named_parameters():
+                if param.dim() > 1:  # 通常只对权重矩阵而不是偏置做初始化
+                    nn.init.xavier_uniform_(param, gain=0.05)
+                else:
+                    nn.init.zeros_(param)
+        else:
+            model = model.to_empty(device=device)
+        model, _, _ = smart_load_weights(model, model_state_dict)
+        # model.load_state_dict(model_state_dict, assign=True, strict=False)
+        model = model.to(dtype=torch_dtype, device=device)
+        loaded_model_names.append(model_name)
+        loaded_models.append(model)
+    return loaded_model_names, loaded_models
+
+
+def load_model_from_huggingface_folder(file_path, model_names, model_classes, torch_dtype, device):
+    loaded_model_names, loaded_models = [], []
+    for model_name, model_class in zip(model_names, model_classes):
+        if torch_dtype in [torch.float32, torch.float16, torch.bfloat16]:
+            model = model_class.from_pretrained(file_path, torch_dtype=torch_dtype).eval()
+        else:
+            model = model_class.from_pretrained(file_path).eval().to(dtype=torch_dtype)
+        if torch_dtype == torch.float16 and hasattr(model, "half"):
+            model = model.half()
+        try:
+            model = model.to(device=device)
+        except:
+            pass
+        loaded_model_names.append(model_name)
+        loaded_models.append(model)
+    return loaded_model_names, loaded_models
+
+
+def load_single_patch_model_from_single_file(state_dict, model_name, model_class, base_model, extra_kwargs, torch_dtype, device):
+    print(f"    model_name: {model_name} model_class: {model_class.__name__} extra_kwargs: {extra_kwargs}")
+    base_state_dict = base_model.state_dict()
+    base_model.to("cpu")
+    del base_model
+    model = model_class(**extra_kwargs)
+    model.load_state_dict(base_state_dict, strict=False)
+    model.load_state_dict(state_dict, strict=False)
+    model.to(dtype=torch_dtype, device=device)
+    return model
+
+
+def load_patch_model_from_single_file(state_dict, model_names, model_classes, extra_kwargs, model_manager, torch_dtype, device):
+    loaded_model_names, loaded_models = [], []
+    for model_name, model_class in zip(model_names, model_classes):
+        while True:
+            for model_id in range(len(model_manager.model)):
+                base_model_name = model_manager.model_name[model_id]
+                if base_model_name == model_name:
+                    base_model_path = model_manager.model_path[model_id]
+                    base_model = model_manager.model[model_id]
+                    print(f"    Adding patch model to {base_model_name} ({base_model_path})")
+                    patched_model = load_single_patch_model_from_single_file(
+                        state_dict, model_name, model_class, base_model, extra_kwargs, torch_dtype, device)
+                    loaded_model_names.append(base_model_name)
+                    loaded_models.append(patched_model)
+                    model_manager.model.pop(model_id)
+                    model_manager.model_path.pop(model_id)
+                    model_manager.model_name.pop(model_id)
+                    break
+            else:
+                break
+    return loaded_model_names, loaded_models
+
+
+
+class ModelDetectorTemplate:
+    def __init__(self):
+        pass
+
+    def match(self, file_path="", state_dict={}):
+        return False
+    
+    def load(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, **kwargs):
+        return [], []
+    
+
+
+class ModelDetectorFromSingleFile:
+    def __init__(self, model_loader_configs=[]):
+        self.keys_hash_with_shape_dict = {}
+        self.keys_hash_dict = {}
+        for metadata in model_loader_configs:
+            self.add_model_metadata(*metadata)
+
+
+    def add_model_metadata(self, keys_hash, keys_hash_with_shape, model_names, model_classes, model_resource):
+        self.keys_hash_with_shape_dict[keys_hash_with_shape] = (model_names, model_classes, model_resource)
+        if keys_hash is not None:
+            self.keys_hash_dict[keys_hash] = (model_names, model_classes, model_resource)
+
+
+    def match(self, file_path="", state_dict={}):
+        if isinstance(file_path, str) and os.path.isdir(file_path):
+            return False
+        if len(state_dict) == 0:
+            state_dict = load_state_dict(file_path)
+        keys_hash_with_shape = hash_state_dict_keys(state_dict, with_shape=True)
+        if keys_hash_with_shape in self.keys_hash_with_shape_dict:
+            return True
+        keys_hash = hash_state_dict_keys(state_dict, with_shape=False)
+        if keys_hash in self.keys_hash_dict:
+            return True
+        return False
+
+
+    def load(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, infer=False, **kwargs):
+        if len(state_dict) == 0:
+            state_dict = load_state_dict(file_path)
+
+        # Load models with strict matching
+        keys_hash_with_shape = hash_state_dict_keys(state_dict, with_shape=True)
+        if keys_hash_with_shape in self.keys_hash_with_shape_dict:
+            model_names, model_classes, model_resource = self.keys_hash_with_shape_dict[keys_hash_with_shape]
+            loaded_model_names, loaded_models = load_model_from_single_file(state_dict, model_names, model_classes, model_resource, torch_dtype, device, infer)
+            return loaded_model_names, loaded_models
+
+        # Load models without strict matching
+        # (the shape of parameters may be inconsistent, and the state_dict_converter will modify the model architecture)
+        keys_hash = hash_state_dict_keys(state_dict, with_shape=False)
+        if keys_hash in self.keys_hash_dict:
+            model_names, model_classes, model_resource = self.keys_hash_dict[keys_hash]
+            loaded_model_names, loaded_models = load_model_from_single_file(state_dict, model_names, model_classes, model_resource, torch_dtype, device, infer)
+            return loaded_model_names, loaded_models
+
+        return loaded_model_names, loaded_models
+
+
+
+class ModelDetectorFromSplitedSingleFile(ModelDetectorFromSingleFile):
+    def __init__(self, model_loader_configs=[]):
+        super().__init__(model_loader_configs)
+
+
+    def match(self, file_path="", state_dict={}):
+        if isinstance(file_path, str) and os.path.isdir(file_path):
+            return False
+        if len(state_dict) == 0:
+            state_dict = load_state_dict(file_path)
+        splited_state_dict = split_state_dict_with_prefix(state_dict)
+        for sub_state_dict in splited_state_dict:
+            if super().match(file_path, sub_state_dict):
+                return True
+        return False
+
+
+    def load(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, **kwargs):
+        # Split the state_dict and load from each component
+        splited_state_dict = split_state_dict_with_prefix(state_dict)
+        valid_state_dict = {}
+        for sub_state_dict in splited_state_dict:
+            if super().match(file_path, sub_state_dict):
+                valid_state_dict.update(sub_state_dict)
+        if super().match(file_path, valid_state_dict):
+            loaded_model_names, loaded_models = super().load(file_path, valid_state_dict, device, torch_dtype)
+        else:
+            loaded_model_names, loaded_models = [], []
+            for sub_state_dict in splited_state_dict:
+                if super().match(file_path, sub_state_dict):
+                    loaded_model_names_, loaded_models_ = super().load(file_path, valid_state_dict, device, torch_dtype)
+                    loaded_model_names += loaded_model_names_
+                    loaded_models += loaded_models_
+        return loaded_model_names, loaded_models
+    
+
+
+class ModelDetectorFromHuggingfaceFolder:
+    def __init__(self, model_loader_configs=[]):
+        self.architecture_dict = {}
+        for metadata in model_loader_configs:
+            self.add_model_metadata(*metadata)
+
+
+    def add_model_metadata(self, architecture, huggingface_lib, model_name, redirected_architecture):
+        self.architecture_dict[architecture] = (huggingface_lib, model_name, redirected_architecture)
+
+
+    def match(self, file_path="", state_dict={}):
+        if not isinstance(file_path, str) or os.path.isfile(file_path):
+            return False
+        file_list = os.listdir(file_path)
+        if "config.json" not in file_list:
+            return False
+        with open(os.path.join(file_path, "config.json"), "r") as f:
+            config = json.load(f)
+        if "architectures" not in config and "_class_name" not in config:
+            return False
+        return True
+
+
+    def load(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, **kwargs):
+        with open(os.path.join(file_path, "config.json"), "r") as f:
+            config = json.load(f)
+        loaded_model_names, loaded_models = [], []
+        architectures = config["architectures"] if "architectures" in config else [config["_class_name"]]
+        for architecture in architectures:
+            huggingface_lib, model_name, redirected_architecture = self.architecture_dict[architecture]
+            if redirected_architecture is not None:
+                architecture = redirected_architecture
+            model_class = importlib.import_module(huggingface_lib).__getattribute__(architecture)
+            loaded_model_names_, loaded_models_ = load_model_from_huggingface_folder(file_path, [model_name], [model_class], torch_dtype, device)
+            loaded_model_names += loaded_model_names_
+            loaded_models += loaded_models_
+        return loaded_model_names, loaded_models
+    
+
+
+class ModelDetectorFromPatchedSingleFile:
+    def __init__(self, model_loader_configs=[]):
+        self.keys_hash_with_shape_dict = {}
+        for metadata in model_loader_configs:
+            self.add_model_metadata(*metadata)
+
+
+    def add_model_metadata(self, keys_hash_with_shape, model_name, model_class, extra_kwargs):
+        self.keys_hash_with_shape_dict[keys_hash_with_shape] = (model_name, model_class, extra_kwargs)
+
+
+    def match(self, file_path="", state_dict={}):
+        if not isinstance(file_path, str) or os.path.isdir(file_path):
+            return False
+        if len(state_dict) == 0:
+            state_dict = load_state_dict(file_path)
+        keys_hash_with_shape = hash_state_dict_keys(state_dict, with_shape=True)
+        if keys_hash_with_shape in self.keys_hash_with_shape_dict:
+            return True
+        return False
+
+
+    def load(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, model_manager=None, **kwargs):
+        if len(state_dict) == 0:
+            state_dict = load_state_dict(file_path)
+
+        # Load models with strict matching
+        loaded_model_names, loaded_models = [], []
+        keys_hash_with_shape = hash_state_dict_keys(state_dict, with_shape=True)
+        if keys_hash_with_shape in self.keys_hash_with_shape_dict:
+            model_names, model_classes, extra_kwargs = self.keys_hash_with_shape_dict[keys_hash_with_shape]
+            loaded_model_names_, loaded_models_ = load_patch_model_from_single_file(
+                state_dict, model_names, model_classes, extra_kwargs, model_manager, torch_dtype, device)
+            loaded_model_names += loaded_model_names_
+            loaded_models += loaded_models_
+        return loaded_model_names, loaded_models
+
+
+
+class ModelManager:
+    def __init__(
+        self,
+        torch_dtype=torch.float16,
+        device="cuda",
+        model_id_list: List = [],
+        downloading_priority: List = ["ModelScope", "HuggingFace"],
+        file_path_list: List[str] = [],
+        infer: bool = False
+    ):
+        self.torch_dtype = torch_dtype
+        self.device = device
+        self.model = []
+        self.model_path = []
+        self.model_name = []
+        self.infer = infer
+        downloaded_files = []
+        self.model_detector = [
+            ModelDetectorFromSingleFile(model_loader_configs),
+            ModelDetectorFromSplitedSingleFile(model_loader_configs),
+            ModelDetectorFromHuggingfaceFolder(huggingface_model_loader_configs),
+        ]
+        self.load_models(downloaded_files + file_path_list)
+
+
+    def load_model_from_single_file(self, file_path="", state_dict={}, model_names=[], model_classes=[], model_resource=None):
+        print(f"Loading models from file: {file_path}")
+        if len(state_dict) == 0:
+            state_dict = load_state_dict(file_path)
+        model_names, models = load_model_from_single_file(state_dict, model_names, model_classes, model_resource, self.torch_dtype, self.device, self.infer)
+        for model_name, model in zip(model_names, models):
+            self.model.append(model)
+            self.model_path.append(file_path)
+            self.model_name.append(model_name)
+        print(f"    The following models are loaded: {model_names}.")
+
+
+    def load_model_from_huggingface_folder(self, file_path="", model_names=[], model_classes=[]):
+        print(f"Loading models from folder: {file_path}")
+        model_names, models = load_model_from_huggingface_folder(file_path, model_names, model_classes, self.torch_dtype, self.device)
+        for model_name, model in zip(model_names, models):
+            self.model.append(model)
+            self.model_path.append(file_path)
+            self.model_name.append(model_name)
+        print(f"    The following models are loaded: {model_names}.")
+
+
+    def load_patch_model_from_single_file(self, file_path="", state_dict={}, model_names=[], model_classes=[], extra_kwargs={}):
+        print(f"Loading patch models from file: {file_path}")
+        model_names, models = load_patch_model_from_single_file(
+            state_dict, model_names, model_classes, extra_kwargs, self, self.torch_dtype, self.device)
+        for model_name, model in zip(model_names, models):
+            self.model.append(model)
+            self.model_path.append(file_path)
+            self.model_name.append(model_name)
+        print(f"    The following patched models are loaded: {model_names}.")
+
+    def load_model(self, file_path, model_names=None, device=None, torch_dtype=None):
+        print(f"Loading models from: {file_path}")
+        if device is None: device = self.device             ## cpu
+        if torch_dtype is None: torch_dtype = self.torch_dtype
+        if isinstance(file_path, list):         ##  <--------   file_path = ["pretrained_models/Wan2.1-T2V-14B/diffusion_pytorch_model-00001-of-00006.safetensors", ..., "...-00006-of-00006.safetensors"
+            state_dict = {}
+            """
+            state_dict:
+            {
+                "patch_embedding.weight": tensor(...),       # from shard 1
+                "patch_embedding.bias": tensor(...),         # from shard 1
+                "blocks.0.self_attn.q.weight": tensor(...),  # from shard 1
+                "blocks.0.self_attn.k.weight": tensor(...),  # from shard 1
+                ...
+                "blocks.20.ffn.2.weight": tensor(...),       # from shard 3 or 4
+                ...
+                "blocks.39.ffn.2.weight": tensor(...),       # from shard 6
+                "head.head.weight": tensor(...),             # from shard 6
+            }
+            """
+            for path in file_path:
+                state_dict.update(load_state_dict(path))
+        elif os.path.isfile(file_path):
+            """
+            For T5 (models_t5_umt5-xxl-enc-bf16.pth):
+                {
+                    "encoder.embed_tokens.weight": tensor(...),
+                    "encoder.block.0.layer.0.SelfAttention.q.weight": tensor(...),
+                    ...
+                    "encoder.block.23.layer.1.DenseReluDense.wi_1.weight": tensor(...),
+                    "encoder.final_layer_norm.weight": tensor(...),
+                }
+
+                For VAE (Wan2.1_VAE.pth):
+                {
+                    "encoder.conv_in.conv.weight": tensor(...),
+                    "encoder.mid.attn_1.norm.weight": tensor(...),
+                    ...
+                    "decoder.conv_out.conv.weight": tensor(...),
+                }
+            """
+            state_dict = load_state_dict(file_path)
+        else:
+            state_dict = None
+
+        """
+        It calls into ModelDetectorFromSingleFile.load() (line 138), which looks up the matched model class from the hash table, then calls load_model_from_single_file() (line 8). That function
+        does:
+
+        1. Gets the model class — e.g. WanVideoModel for DiT, WanVideoVAE for VAE
+        2. state_dict_converter — converts key names from the checkpoint format to the model's internal format (e.g. HuggingFace naming → DiffSynth naming)
+        3. model_class(**extra_kwargs) — constructs the model with empty weights. For the DiT, extra_kwargs includes in_dim from args.model_config (which is 49 for V2V)
+        4. Since infer=False (lines 27-33):
+            - Moves to CPU with to_empty(device="cpu")
+            - Xavier-inits all weight matrices (gain=0.05), zeros all biases
+        5. smart_load_weights(model, model_state_dict) — overlays the Wan 2.1 weights onto the xavier-initialized model. For mismatched shapes (like 16ch checkpoint into 49ch patch_embedding), it
+        copies what fits and leaves the rest xavier-initialized
+        6. Casts to target dtype/device — model.to(dtype=torch_dtype, device="cpu")
+
+        So the model starts fully xavier-initialized, then gets the base Wan 2.1 weights overlaid on top. New parameters (extra patch_embedding channels, audio module placeholders) keep their
+        xavier values.
+        """
+        for model_detector in self.model_detector:
+            if model_detector.match(file_path, state_dict):
+                model_names, models = model_detector.load(
+                    file_path, state_dict,
+                    device=device, torch_dtype=torch_dtype,
+                    allowed_model_names=model_names, model_manager=self, infer=self.infer
+                )
+                for model_name, model in zip(model_names, models):
+                    self.model.append(model)
+                    self.model_path.append(file_path)
+                    self.model_name.append(model_name)
+                print(f"    The following models are loaded: {model_names}.")
+                break
+        else:
+            print(f"    We cannot detect the model type. No models are loaded.")
+        
+
+    def load_models(self, file_path_list, model_names=None, device=None, torch_dtype=None):
+        for file_path in file_path_list:
+            self.load_model(file_path, model_names, device=device, torch_dtype=torch_dtype)
+
+    
+    def fetch_model(self, model_name, file_path=None, require_model_path=False):
+        fetched_models = []
+        fetched_model_paths = []
+        for model, model_path, model_name_ in zip(self.model, self.model_path, self.model_name):
+            if file_path is not None and file_path != model_path:
+                continue
+            if model_name == model_name_:
+                fetched_models.append(model)
+                fetched_model_paths.append(model_path)
+        if len(fetched_models) == 0:
+            print(f"No {model_name} models available.")
+            return None
+        if len(fetched_models) == 1:
+            print(f"Using {model_name} from {fetched_model_paths[0]}.")
+        else:
+            print(f"More than one {model_name} models are loaded in model manager: {fetched_model_paths}. Using {model_name} from {fetched_model_paths[0]}.")
+        if require_model_path:
+            return fetched_models[0], fetched_model_paths[0]
+        else:
+            return fetched_models[0]
+        
+
+    def to(self, device):
+        for model in self.model:
+            model.to(device)
+

OmniAvatar/models/wan_video_dit.py

@@ -0,0 +1,611 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Tuple, Optional
+from einops import rearrange
+from ..utils.io_utils import hash_state_dict_keys
+from .audio_pack import AudioPack
+from ..utils.args_config import args
+
+# IMPORTANT: Import flash_attn_interface BEFORE xfuser to avoid operator conflicts
+# xfuser registers simplified flash_attn_3 operators that would override the real ones
+try:
+    import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    flash_attn_interface = None
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+try:
+    from sageattention import sageattn
+    SAGE_ATTN_AVAILABLE = True
+except ModuleNotFoundError:
+    SAGE_ATTN_AVAILABLE = False
+
+# Lazy import for xfuser - import AFTER flash_attn to avoid operator conflicts
+_xfuser_imported = False
+_xfuser_get_sequence_parallel_rank = None
+_xfuser_get_sequence_parallel_world_size = None
+_xfuser_get_sp_group = None
+
+def _lazy_import_xfuser():
+    """Lazily import xfuser only when context parallel is actually used."""
+    global _xfuser_imported, _xfuser_get_sequence_parallel_rank
+    global _xfuser_get_sequence_parallel_world_size, _xfuser_get_sp_group
+    if not _xfuser_imported:
+        from xfuser.core.distributed import (
+            get_sequence_parallel_rank,
+            get_sequence_parallel_world_size,
+            get_sp_group
+        )
+        _xfuser_get_sequence_parallel_rank = get_sequence_parallel_rank
+        _xfuser_get_sequence_parallel_world_size = get_sequence_parallel_world_size
+        _xfuser_get_sp_group = get_sp_group
+        _xfuser_imported = True
+
+def get_sequence_parallel_rank():
+    _lazy_import_xfuser()
+    return _xfuser_get_sequence_parallel_rank()
+
+def get_sequence_parallel_world_size():
+    _lazy_import_xfuser()
+    return _xfuser_get_sequence_parallel_world_size()
+
+def get_sp_group():
+    _lazy_import_xfuser()
+    return _xfuser_get_sp_group()
+
+
+def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False):
+    # When context parallel is enabled (sp_size > 1), xfuser conflicts with flash_attn_3
+    # so we must fall back to flash_attn_2 or SDPA
+    use_flash_attn_3 = FLASH_ATTN_3_AVAILABLE and (args is None or getattr(args, 'sp_size', 1) <= 1)
+    
+    if compatibility_mode:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    elif use_flash_attn_3:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = flash_attn_interface.flash_attn_func(q, k, v)
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    elif FLASH_ATTN_2_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b s n d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b s n d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b s n d", n=num_heads)
+        x = flash_attn.flash_attn_func(q, k, v)
+        x = rearrange(x, "b s n d -> b s (n d)", n=num_heads)
+    elif SAGE_ATTN_AVAILABLE:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = sageattn(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    else:
+        q = rearrange(q, "b s (n d) -> b n s d", n=num_heads)
+        k = rearrange(k, "b s (n d) -> b n s d", n=num_heads)
+        v = rearrange(v, "b s (n d) -> b n s d", n=num_heads)
+        x = F.scaled_dot_product_attention(q, k, v)
+        x = rearrange(x, "b n s d -> b s (n d)", n=num_heads)
+    return x
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return (x * (1 + scale) + shift)
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(position.type(torch.float64), torch.pow(
+        10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return f_freqs_cis, h_freqs_cis, w_freqs_cis
+
+
+def precompute_freqs_cis(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 1d rope precompute
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)
+                   [: (dim // 2)].double() / dim))
+    freqs = torch.outer(torch.arange(end, device=freqs.device), freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def rope_apply(x, freqs, num_heads):
+    x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
+    x_out = torch.view_as_complex(x.to(torch.float64).reshape(
+        x.shape[0], x.shape[1], x.shape[2], -1, 2))
+    x_out = torch.view_as_real(x_out * freqs).flatten(2)
+    return x_out.to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        dtype = x.dtype
+        return self.norm(x.float()).to(dtype) * self.weight
+
+
+class AttentionModule(nn.Module):
+    def __init__(self, num_heads):
+        super().__init__()
+        self.num_heads = num_heads
+        
+    def forward(self, q, k, v):
+        x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads)
+        return x
+
+
+class SelfAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+        
+        self.attn = AttentionModule(self.num_heads)
+
+    def forward(self, x, freqs):
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(x))
+        v = self.v(x)
+        q = rope_apply(q, freqs, self.num_heads)
+        k = rope_apply(k, freqs, self.num_heads)
+        x = self.attn(q, k, v)
+        return self.o(x)
+
+
+class CrossAttention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.q = nn.Linear(dim, dim)
+        self.k = nn.Linear(dim, dim)
+        self.v = nn.Linear(dim, dim)
+        self.o = nn.Linear(dim, dim)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+        self.has_image_input = has_image_input
+        if has_image_input:
+            self.k_img = nn.Linear(dim, dim)
+            self.v_img = nn.Linear(dim, dim)
+            self.norm_k_img = RMSNorm(dim, eps=eps)
+            
+        self.attn = AttentionModule(self.num_heads)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        if self.has_image_input:
+            img = y[:, :257]
+            ctx = y[:, 257:]
+        else:
+            ctx = y
+        q = self.norm_q(self.q(x))
+        k = self.norm_k(self.k(ctx))
+        v = self.v(ctx)
+        x = self.attn(q, k, v)
+        if self.has_image_input:
+            k_img = self.norm_k_img(self.k_img(img))
+            v_img = self.v_img(img)
+            y = flash_attention(q, k_img, v_img, num_heads=self.num_heads)
+            x = x + y
+        return self.o(x)
+
+
+class GateModule(nn.Module):
+    def __init__(self,):
+        super().__init__()
+
+    def forward(self, x, gate, residual):
+        return x + gate * residual
+
+class DiTBlock(nn.Module):
+    def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+
+        self.self_attn = SelfAttention(dim, num_heads, eps)
+        self.cross_attn = CrossAttention(
+            dim, num_heads, eps, has_image_input=has_image_input)
+        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3 = nn.LayerNorm(dim, eps=eps)
+        self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(
+            approximate='tanh'), nn.Linear(ffn_dim, dim))
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.gate = GateModule()
+
+    def forward(self, x, context, t_mod, freqs):
+        # msa: multi-head self-attention  mlp: multi-layer perceptron
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1)
+        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
+        x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
+        x = x + self.cross_attn(self.norm3(x), context)
+        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = self.gate(x, gate_mlp, self.ffn(input_x))
+        return x
+
+
+class MLP(torch.nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+        self.proj = torch.nn.Sequential(
+            nn.LayerNorm(in_dim),
+            nn.Linear(in_dim, in_dim),
+            nn.GELU(),
+            nn.Linear(in_dim, out_dim),
+            nn.LayerNorm(out_dim)
+        )
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class Head(nn.Module):
+    def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.head = nn.Linear(dim, out_dim * math.prod(patch_size))
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        # t_mod: [B, dim] -> [B, 1, dim] for broadcasting with modulation [1, 2, dim]
+        if t_mod.dim() == 2:
+            t_mod = t_mod.unsqueeze(1)
+        shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1)
+        x = (self.head(self.norm(x) * (1 + scale) + shift))
+        return x
+
+
+
+class WanModel(torch.nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        in_dim: int,
+        ffn_dim: int,
+        out_dim: int,
+        text_dim: int,
+        freq_dim: int,
+        eps: float,
+        patch_size: Tuple[int, int, int],
+        num_heads: int,
+        num_layers: int,
+        has_image_input: bool,
+        audio_hidden_size: int=32,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.has_image_input = has_image_input
+        self.patch_size = patch_size
+
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+            # nn.LayerNorm(dim)
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim),
+            nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim)
+        )
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim)
+        )
+        self.time_projection = nn.Sequential(
+            nn.SiLU(), nn.Linear(dim, dim * 6))
+        self.blocks = nn.ModuleList([
+            DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps)
+            for _ in range(num_layers)
+        ])
+        self.head = Head(dim, out_dim, patch_size, eps)
+        head_dim = dim // num_heads
+        self.freqs = precompute_freqs_cis_3d(head_dim)
+
+        if has_image_input:
+            self.img_emb = MLP(1280, dim)  # clip_feature_dim = 1280
+
+        if 'use_audio' in args:
+            self.use_audio = args.use_audio
+        else:
+            self.use_audio = False
+        if self.use_audio:
+            audio_input_dim = 10752
+            audio_out_dim = dim
+            self.audio_proj = AudioPack(audio_input_dim, [4, 1, 1], audio_hidden_size, layernorm=True)
+            self.audio_cond_projs = nn.ModuleList()
+            for d in range(num_layers // 2 - 1):
+                l = nn.Linear(audio_hidden_size, audio_out_dim)
+                self.audio_cond_projs.append(l)      
+
+    def patchify(self, x: torch.Tensor):
+        grid_size = x.shape[2:]
+        x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous()
+        return x, grid_size  # x, grid_size: (f, h, w)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor):
+        return rearrange(
+            x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)',
+            f=grid_size[0], h=grid_size[1], w=grid_size[2], 
+            x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2]
+        )
+
+    def forward(self,
+                x: torch.Tensor,
+                timestep: torch.Tensor,
+                context: torch.Tensor,
+                clip_feature: Optional[torch.Tensor] = None,
+                y: Optional[torch.Tensor] = None,
+                use_gradient_checkpointing: bool = False,
+                audio_emb: Optional[torch.Tensor] = None,
+                use_gradient_checkpointing_offload: bool = False,
+                tea_cache = None,
+                **kwargs,
+                ):
+        t = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim, timestep))
+        t_mod = self.time_projection(t).unflatten(1, (6, self.dim))
+        context = self.text_embedding(context)
+        lat_h, lat_w = x.shape[-2], x.shape[-1]
+
+        if audio_emb != None and self.use_audio: # TODO  cache
+            audio_emb = audio_emb.permute(0, 2, 1)[:, :, :, None, None]
+            audio_emb = torch.cat([audio_emb[:, :, :1].repeat(1, 1, 3, 1, 1), audio_emb], 2) # 1, 768, 44, 1, 1
+            audio_emb = self.audio_proj(audio_emb)
+
+            audio_emb = torch.stack([audio_cond_proj(audio_emb) for audio_cond_proj in self.audio_cond_projs], dim=1)
+
+        x = torch.cat([x, y], dim=1)
+        x = self.patch_embedding(x)
+        x, (f, h, w) = self.patchify(x)
+        
+        freqs = torch.cat([
+            self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ], dim=-1).reshape(f * h * w, 1, -1).to(x.device)
+        
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+            return custom_forward
+        
+        if tea_cache is not None:
+            tea_cache_update = tea_cache.check(self, x, t_mod)
+        else:
+            tea_cache_update = False
+        ori_x_len = x.shape[1]
+        if tea_cache_update:
+            x = tea_cache.update(x)
+        else:
+            if args.sp_size > 1:
+                # Context Parallel
+                sp_size = get_sequence_parallel_world_size()
+                pad_size = 0
+                if ori_x_len % sp_size != 0:
+                    pad_size = sp_size - ori_x_len % sp_size
+                    x = torch.cat([x, torch.zeros_like(x[:, -1:]).repeat(1, pad_size, 1)], 1)
+                x = torch.chunk(x, sp_size, dim=1)[get_sequence_parallel_rank()]
+
+            # audio_emb is already [B, num_projs, C, T, 1, 1] from torch.stack(..., dim=1)
+                
+            for layer_i, block in enumerate(self.blocks):
+                # audio cond
+                if self.use_audio:
+                    au_idx = None
+                    if (layer_i <= len(self.blocks) // 2 and layer_i > 1): # < len(self.blocks) - 1:
+                        au_idx = layer_i - 2
+                        audio_emb_tmp = audio_emb[:, au_idx].repeat(1, 1, lat_h // 2, lat_w // 2, 1) # 1, 11, 45, 25, 128
+                        audio_cond_tmp = self.patchify(audio_emb_tmp.permute(0, 4, 1, 2, 3))[0]
+                        if args.sp_size > 1:
+                            if pad_size > 0:    
+                                audio_cond_tmp = torch.cat([audio_cond_tmp, torch.zeros_like(audio_cond_tmp[:, -1:]).repeat(1, pad_size, 1)], 1)
+                            audio_cond_tmp = torch.chunk(audio_cond_tmp, sp_size, dim=1)[get_sequence_parallel_rank()]
+                        x = audio_cond_tmp + x
+
+                if self.training and use_gradient_checkpointing:
+                    if use_gradient_checkpointing_offload:
+                        with torch.autograd.graph.save_on_cpu():
+                            x = torch.utils.checkpoint.checkpoint(
+                                create_custom_forward(block),
+                                x, context, t_mod, freqs,
+                                use_reentrant=False,
+                            )
+                    else:
+                        x = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            x, context, t_mod, freqs,
+                            use_reentrant=False,
+                        )
+                else:
+                    x = block(x, context, t_mod, freqs)
+            if tea_cache is not None:
+                x_cache = get_sp_group().all_gather(x, dim=1) # TODO: the size should be devided by sp_size
+                x_cache = x_cache[:, :ori_x_len]
+                tea_cache.store(x_cache)
+
+        x = self.head(x, t)
+        if args.sp_size > 1:
+            # Context Parallel
+            x = get_sp_group().all_gather(x, dim=1) # TODO: the size should be devided by sp_size
+            x = x[:, :ori_x_len]
+
+        x = self.unpatchify(x, (f, h, w))
+        return x
+
+    @staticmethod
+    def state_dict_converter():
+        return WanModelStateDictConverter()
+    
+    
+class WanModelStateDictConverter:
+    def __init__(self):
+        pass
+
+    def from_diffusers(self, state_dict):
+        rename_dict = {
+            "blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight",
+            "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight",
+            "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias",
+            "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight",
+            "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias",
+            "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight",
+            "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias",
+            "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight",
+            "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias",
+            "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight",
+            "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight",
+            "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight",
+            "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias",
+            "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight",
+            "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias",
+            "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight",
+            "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias",
+            "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight",
+            "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias",
+            "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight",
+            "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias",
+            "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight",
+            "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias",
+            "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight",
+            "blocks.0.norm2.bias": "blocks.0.norm3.bias",
+            "blocks.0.norm2.weight": "blocks.0.norm3.weight",
+            "blocks.0.scale_shift_table": "blocks.0.modulation",
+            "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias",
+            "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight",
+            "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias",
+            "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight",
+            "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias",
+            "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight",
+            "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias",
+            "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight",
+            "condition_embedder.time_proj.bias": "time_projection.1.bias",
+            "condition_embedder.time_proj.weight": "time_projection.1.weight",
+            "patch_embedding.bias": "patch_embedding.bias",
+            "patch_embedding.weight": "patch_embedding.weight",
+            "scale_shift_table": "head.modulation",
+            "proj_out.bias": "head.head.bias",
+            "proj_out.weight": "head.head.weight",
+        }
+        state_dict_ = {}
+        for name, param in state_dict.items():
+            if name in rename_dict:
+                state_dict_[rename_dict[name]] = param
+            else:
+                name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:])
+                if name_ in rename_dict:
+                    name_ = rename_dict[name_]
+                    name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:])
+                    state_dict_[name_] = param
+        if hash_state_dict_keys(state_dict) == "cb104773c6c2cb6df4f9529ad5c60d0b":
+            config = {
+                "model_type": "t2v",
+                "patch_size": (1, 2, 2),
+                "text_len": 512,
+                "in_dim": 16,
+                "dim": 5120,
+                "ffn_dim": 13824,
+                "freq_dim": 256,
+                "text_dim": 4096,
+                "out_dim": 16,
+                "num_heads": 40,
+                "num_layers": 40,
+                "window_size": (-1, -1),
+                "qk_norm": True,
+                "cross_attn_norm": True,
+                "eps": 1e-6,
+            }
+        else:
+            config = {}
+        return state_dict_, config
+    
+    def from_civitai(self, state_dict):
+        if hash_state_dict_keys(state_dict) == "9269f8db9040a9d860eaca435be61814":
+            config = {
+                "has_image_input": False,
+                "patch_size": [1, 2, 2],
+                "in_dim": 16,
+                "dim": 1536,
+                "ffn_dim": 8960,
+                "freq_dim": 256,
+                "text_dim": 4096,
+                "out_dim": 16,
+                "num_heads": 12,
+                "num_layers": 30,
+                "eps": 1e-6
+            }
+        elif hash_state_dict_keys(state_dict) == "aafcfd9672c3a2456dc46e1cb6e52c70":
+            config = {
+                "has_image_input": False,
+                "patch_size": [1, 2, 2],
+                "in_dim": 16,
+                "dim": 5120,
+                "ffn_dim": 13824,
+                "freq_dim": 256,
+                "text_dim": 4096,
+                "out_dim": 16,
+                "num_heads": 40,
+                "num_layers": 40,
+                "eps": 1e-6
+            }
+        elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e":
+            config = {
+                "has_image_input": True,
+                "patch_size": [1, 2, 2],
+                "in_dim": 36,
+                "dim": 5120,
+                "ffn_dim": 13824,
+                "freq_dim": 256,
+                "text_dim": 4096,
+                "out_dim": 16,
+                "num_heads": 40,
+                "num_layers": 40,
+                "eps": 1e-6
+            }
+        else:
+            config = {}
+        if hasattr(args, "model_config"):
+            model_config = args.model_config
+            if model_config is not None:
+                config.update(model_config)        
+        return state_dict, config

OmniAvatar/models/wan_video_text_encoder.py

@@ -0,0 +1,269 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def fp16_clamp(x):
+    if x.dtype == torch.float16 and torch.isinf(x).any():
+        clamp = torch.finfo(x.dtype).max - 1000
+        x = torch.clamp(x, min=-clamp, max=clamp)
+    return x
+
+
+class GELU(nn.Module):
+
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.tanh(
+            math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+
+    def __init__(self, dim, eps=1e-6):
+        super(T5LayerNorm, self).__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+                            self.eps)
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            x = x.type_as(self.weight)
+        return self.weight * x
+
+
+class T5Attention(nn.Module):
+
+    def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+        assert dim_attn % num_heads == 0
+        super(T5Attention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.num_heads = num_heads
+        self.head_dim = dim_attn // num_heads
+
+        # layers
+        self.q = nn.Linear(dim, dim_attn, bias=False)
+        self.k = nn.Linear(dim, dim_attn, bias=False)
+        self.v = nn.Linear(dim, dim_attn, bias=False)
+        self.o = nn.Linear(dim_attn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, context=None, mask=None, pos_bias=None):
+        """
+        x:          [B, L1, C].
+        context:    [B, L2, C] or None.
+        mask:       [B, L2] or [B, L1, L2] or None.
+        """
+        # check inputs
+        context = x if context is None else context
+        b, n, c = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.q(x).view(b, -1, n, c)
+        k = self.k(context).view(b, -1, n, c)
+        v = self.v(context).view(b, -1, n, c)
+
+        # attention bias
+        attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+        if pos_bias is not None:
+            attn_bias += pos_bias
+        if mask is not None:
+            assert mask.ndim in [2, 3]
+            mask = mask.view(b, 1, 1,
+                             -1) if mask.ndim == 2 else mask.unsqueeze(1)
+            attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+        # compute attention (T5 does not use scaling)
+        attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+        attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+        x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+        # output
+        x = x.reshape(b, -1, n * c)
+        x = self.o(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5FeedForward(nn.Module):
+
+    def __init__(self, dim, dim_ffn, dropout=0.1):
+        super(T5FeedForward, self).__init__()
+        self.dim = dim
+        self.dim_ffn = dim_ffn
+
+        # layers
+        self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+        self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+        self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        x = self.fc1(x) * self.gate(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class T5SelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 dim_attn,
+                 dim_ffn,
+                 num_heads,
+                 num_buckets,
+                 shared_pos=True,
+                 dropout=0.1):
+        super(T5SelfAttention, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.norm1 = T5LayerNorm(dim)
+        self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+        self.norm2 = T5LayerNorm(dim)
+        self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+        self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True)
+
+    def forward(self, x, mask=None, pos_bias=None):
+        e = pos_bias if self.shared_pos else self.pos_embedding(
+            x.size(1), x.size(1))
+        x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+        x = fp16_clamp(x + self.ffn(self.norm2(x)))
+        return x
+
+
+class T5RelativeEmbedding(nn.Module):
+
+    def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+        super(T5RelativeEmbedding, self).__init__()
+        self.num_buckets = num_buckets
+        self.num_heads = num_heads
+        self.bidirectional = bidirectional
+        self.max_dist = max_dist
+
+        # layers
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, lq, lk):
+        device = self.embedding.weight.device
+        # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+        #     torch.arange(lq).unsqueeze(1).to(device)
+        rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+            torch.arange(lq, device=device).unsqueeze(1)
+        rel_pos = self._relative_position_bucket(rel_pos)
+        rel_pos_embeds = self.embedding(rel_pos)
+        rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+            0)  # [1, N, Lq, Lk]
+        return rel_pos_embeds.contiguous()
+
+    def _relative_position_bucket(self, rel_pos):
+        # preprocess
+        if self.bidirectional:
+            num_buckets = self.num_buckets // 2
+            rel_buckets = (rel_pos > 0).long() * num_buckets
+            rel_pos = torch.abs(rel_pos)
+        else:
+            num_buckets = self.num_buckets
+            rel_buckets = 0
+            rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+        # embeddings for small and large positions
+        max_exact = num_buckets // 2
+        rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+                                     math.log(self.max_dist / max_exact) *
+                                     (num_buckets - max_exact)).long()
+        rel_pos_large = torch.min(
+            rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+        rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+        return rel_buckets
+
+def init_weights(m):
+    if isinstance(m, T5LayerNorm):
+        nn.init.ones_(m.weight)
+    elif isinstance(m, T5FeedForward):
+        nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+    elif isinstance(m, T5Attention):
+        nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+        nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+        nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+    elif isinstance(m, T5RelativeEmbedding):
+        nn.init.normal_(
+            m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class WanTextEncoder(torch.nn.Module):
+
+    def __init__(self,
+                 vocab=256384,
+                 dim=4096,
+                 dim_attn=4096,
+                 dim_ffn=10240,
+                 num_heads=64,
+                 num_layers=24,
+                 num_buckets=32,
+                 shared_pos=False,
+                 dropout=0.1):
+        super(WanTextEncoder, self).__init__()
+        self.dim = dim
+        self.dim_attn = dim_attn
+        self.dim_ffn = dim_ffn
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_buckets = num_buckets
+        self.shared_pos = shared_pos
+
+        # layers
+        self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+            else nn.Embedding(vocab, dim)
+        self.pos_embedding = T5RelativeEmbedding(
+            num_buckets, num_heads, bidirectional=True) if shared_pos else None
+        self.dropout = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([
+            T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+                            shared_pos, dropout) for _ in range(num_layers)
+        ])
+        self.norm = T5LayerNorm(dim)
+
+        # initialize weights
+        self.apply(init_weights)
+
+    def forward(self, ids, mask=None):
+        x = self.token_embedding(ids)
+        x = self.dropout(x)
+        e = self.pos_embedding(x.size(1),
+                               x.size(1)) if self.shared_pos else None
+        for block in self.blocks:
+            x = block(x, mask, pos_bias=e)
+        x = self.norm(x)
+        x = self.dropout(x)
+        return x
+    
+    @staticmethod
+    def state_dict_converter():
+        return WanTextEncoderStateDictConverter()
+    
+    
+class WanTextEncoderStateDictConverter:
+    def __init__(self):
+        pass
+
+    def from_diffusers(self, state_dict):
+        return state_dict
+    
+    def from_civitai(self, state_dict):
+        return state_dict

OmniAvatar/models/wan_video_vae.py

@@ -0,0 +1,938 @@
+from einops import rearrange, repeat
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+CACHE_T = 2
+
+
+def check_is_instance(model, module_class):
+    if isinstance(model, module_class):
+        return True
+    if hasattr(model, "module") and isinstance(model.module, module_class):
+        return True
+    return False
+
+
+def block_causal_mask(x, block_size):
+    # params
+    b, n, s, _, device = *x.size(), x.device
+    assert s % block_size == 0
+    num_blocks = s // block_size
+
+    # build mask
+    mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=device)
+    for i in range(num_blocks):
+        mask[:, :,
+             i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1
+    return mask
+
+
+class CausalConv3d(nn.Conv3d):
+    """
+    Causal 3d convolusion.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._padding = (self.padding[2], self.padding[2], self.padding[1],
+                         self.padding[1], 2 * self.padding[0], 0)
+        self.padding = (0, 0, 0)
+
+    def forward(self, x, cache_x=None):
+        padding = list(self._padding)
+        if cache_x is not None and self._padding[4] > 0:
+            cache_x = cache_x.to(x.device)
+            x = torch.cat([cache_x, x], dim=2)
+            padding[4] -= cache_x.shape[2]
+        x = F.pad(x, padding)
+
+        return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+    def __init__(self, dim, channel_first=True, images=True, bias=False):
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
+
+    def forward(self, x):
+        return F.normalize(
+            x, dim=(1 if self.channel_first else
+                    -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+    def forward(self, x):
+        """
+        Fix bfloat16 support for nearest neighbor interpolation.
+        """
+        return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+    def __init__(self, dim, mode):
+        assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+                        'downsample3d')
+        super().__init__()
+        self.dim = dim
+        self.mode = mode
+
+        # layers
+        if mode == 'upsample2d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+        elif mode == 'upsample3d':
+            self.resample = nn.Sequential(
+                Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+                nn.Conv2d(dim, dim // 2, 3, padding=1))
+            self.time_conv = CausalConv3d(dim,
+                                          dim * 2, (3, 1, 1),
+                                          padding=(1, 0, 0))
+
+        elif mode == 'downsample2d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+        elif mode == 'downsample3d':
+            self.resample = nn.Sequential(
+                nn.ZeroPad2d((0, 1, 0, 1)),
+                nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+            self.time_conv = CausalConv3d(dim,
+                                          dim, (3, 1, 1),
+                                          stride=(2, 1, 1),
+                                          padding=(0, 0, 0))
+
+        else:
+            self.resample = nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        b, c, t, h, w = x.size()
+        if self.mode == 'upsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = 'Rep'
+                    feat_idx[0] += 1
+                else:
+
+                    cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] != 'Rep':
+                        # cache last frame of last two chunk
+                        cache_x = torch.cat([
+                            feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                                cache_x.device), cache_x
+                        ],
+                                            dim=2)
+                    if cache_x.shape[2] < 2 and feat_cache[
+                            idx] is not None and feat_cache[idx] == 'Rep':
+                        cache_x = torch.cat([
+                            torch.zeros_like(cache_x).to(cache_x.device),
+                            cache_x
+                        ],
+                                            dim=2)
+                    if feat_cache[idx] == 'Rep':
+                        x = self.time_conv(x)
+                    else:
+                        x = self.time_conv(x, feat_cache[idx])
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+
+                    x = x.reshape(b, 2, c, t, h, w)
+                    x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+                                    3)
+                    x = x.reshape(b, c, t * 2, h, w)
+        t = x.shape[2]
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.resample(x)
+        x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+        if self.mode == 'downsample3d':
+            if feat_cache is not None:
+                idx = feat_idx[0]
+                if feat_cache[idx] is None:
+                    feat_cache[idx] = x.clone()
+                    feat_idx[0] += 1
+                else:
+                    cache_x = x[:, :, -1:, :, :].clone()
+                    x = self.time_conv(
+                        torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+                    feat_cache[idx] = cache_x
+                    feat_idx[0] += 1
+        return x
+
+    def init_weight(self, conv):
+        conv_weight = conv.weight
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        one_matrix = torch.eye(c1, c2)
+        init_matrix = one_matrix
+        nn.init.zeros_(conv_weight)
+        conv_weight.data[:, :, 1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+    def init_weight2(self, conv):
+        conv_weight = conv.weight.data
+        nn.init.zeros_(conv_weight)
+        c1, c2, t, h, w = conv_weight.size()
+        init_matrix = torch.eye(c1 // 2, c2)
+        conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+        conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+        conv.weight.data.copy_(conv_weight)
+        nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(self, in_dim, out_dim, dropout=0.0):
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+
+        # layers
+        self.residual = nn.Sequential(
+            RMS_norm(in_dim, images=False), nn.SiLU(),
+            CausalConv3d(in_dim, out_dim, 3, padding=1),
+            RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+            CausalConv3d(out_dim, out_dim, 3, padding=1))
+        self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+            if in_dim != out_dim else nn.Identity()
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        h = self.shortcut(x)
+        for layer in self.residual:
+            if check_is_instance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x + h
+
+
+class AttentionBlock(nn.Module):
+    """
+    Causal self-attention with a single head.
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.dim = dim
+
+        # layers
+        self.norm = RMS_norm(dim)
+        self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+        self.proj = nn.Conv2d(dim, dim, 1)
+
+        # zero out the last layer params
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, x):
+        identity = x
+        b, c, t, h, w = x.size()
+        x = rearrange(x, 'b c t h w -> (b t) c h w')
+        x = self.norm(x)
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute(
+            0, 1, 3, 2).contiguous().chunk(3, dim=-1)
+
+        # apply attention
+        x = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            #attn_mask=block_causal_mask(q, block_size=h * w)
+        )
+        x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+        # output
+        x = self.proj(x)
+        x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+        return x + identity
+
+
+class Encoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[True, True, False],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+
+        # dimensions
+        dims = [dim * u for u in [1] + dim_mult]
+        scale = 1.0
+
+        # init block
+        self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+        # downsample blocks
+        downsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            for _ in range(num_res_blocks):
+                downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    downsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # downsample block
+            if i != len(dim_mult) - 1:
+                mode = 'downsample3d' if temperal_downsample[
+                    i] else 'downsample2d'
+                downsamples.append(Resample(out_dim, mode=mode))
+                scale /= 2.0
+        self.downsamples = nn.Sequential(*downsamples)
+
+        # middle blocks
+        self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout),
+                                    AttentionBlock(out_dim),
+                                    ResidualBlock(out_dim, out_dim, dropout))
+
+        # output blocks
+        self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(),
+                                  CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## downsamples
+        for layer in self.downsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## middle
+        for layer in self.middle:
+            if check_is_instance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if check_is_instance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+        return x
+
+
+class Decoder3d(nn.Module):
+
+    def __init__(self,
+                 dim=128,
+                 z_dim=4,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_upsample=[False, True, True],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_upsample = temperal_upsample
+
+        # dimensions
+        dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+        scale = 1.0 / 2**(len(dim_mult) - 2)
+
+        # init block
+        self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+        # middle blocks
+        self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout),
+                                    AttentionBlock(dims[0]),
+                                    ResidualBlock(dims[0], dims[0], dropout))
+
+        # upsample blocks
+        upsamples = []
+        for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+            # residual (+attention) blocks
+            if i == 1 or i == 2 or i == 3:
+                in_dim = in_dim // 2
+            for _ in range(num_res_blocks + 1):
+                upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+                if scale in attn_scales:
+                    upsamples.append(AttentionBlock(out_dim))
+                in_dim = out_dim
+
+            # upsample block
+            if i != len(dim_mult) - 1:
+                mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+                upsamples.append(Resample(out_dim, mode=mode))
+                scale *= 2.0
+        self.upsamples = nn.Sequential(*upsamples)
+
+        # output blocks
+        self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(),
+                                  CausalConv3d(out_dim, 3, 3, padding=1))
+
+        # Gradient checkpointing support (default off, enable for aux loss training)
+        self.gradient_checkpointing = False
+
+    def forward(self, x, feat_cache=None, feat_idx=[0]):
+        if torch.is_grad_enabled() and self.gradient_checkpointing:
+            def create_custom_forward(module):
+                def custom_forward(*inputs):
+                    return module._decode(*inputs)
+                return custom_forward
+            return torch.utils.checkpoint.checkpoint(
+                create_custom_forward(self),
+                x, feat_cache, feat_idx,
+                use_reentrant=False
+            )
+        else:
+            return self._decode(x, feat_cache, feat_idx)
+
+    def _decode(self, x, in_cache=None, feat_idx=[0]):
+        # Copy input cache to prevent mutations during gradient checkpoint recomputation
+        feat_cache = in_cache.copy() if in_cache is not None else None
+
+        ## conv1
+        if feat_cache is not None:
+            idx = feat_idx[0]
+            cache_x = x[:, :, -CACHE_T:, :, :].clone()
+            if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                # cache last frame of last two chunk
+                cache_x = torch.cat([
+                    feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                        cache_x.device), cache_x
+                ],
+                                    dim=2)
+            x = self.conv1(x, feat_cache[idx])
+            feat_cache[idx] = cache_x
+            feat_idx[0] += 1
+        else:
+            x = self.conv1(x)
+
+        ## middle
+        for layer in self.middle:
+            if check_is_instance(layer, ResidualBlock) and feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## upsamples
+        for layer in self.upsamples:
+            if feat_cache is not None:
+                x = layer(x, feat_cache, feat_idx)
+            else:
+                x = layer(x)
+
+        ## head
+        for layer in self.head:
+            if check_is_instance(layer, CausalConv3d) and feat_cache is not None:
+                idx = feat_idx[0]
+                cache_x = x[:, :, -CACHE_T:, :, :].clone()
+                if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+                    # cache last frame of last two chunk
+                    cache_x = torch.cat([
+                        feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+                            cache_x.device), cache_x
+                    ],
+                                        dim=2)
+                x = layer(x, feat_cache[idx])
+                feat_cache[idx] = cache_x
+                feat_idx[0] += 1
+            else:
+                x = layer(x)
+
+        # Reset index counter for consistent state across checkpoint recomputations
+        feat_idx[0] = 0
+        return x, feat_cache
+
+
+def count_conv3d(model):
+    count = 0
+    for m in model.modules():
+        if check_is_instance(m, CausalConv3d):
+            count += 1
+    return count
+
+
+class VideoVAE_(nn.Module):
+
+    def __init__(self,
+                 dim=96,
+                 z_dim=16,
+                 dim_mult=[1, 2, 4, 4],
+                 num_res_blocks=2,
+                 attn_scales=[],
+                 temperal_downsample=[False, True, True],
+                 dropout=0.0):
+        super().__init__()
+        self.dim = dim
+        self.z_dim = z_dim
+        self.dim_mult = dim_mult
+        self.num_res_blocks = num_res_blocks
+        self.attn_scales = attn_scales
+        self.temperal_downsample = temperal_downsample
+        self.temperal_upsample = temperal_downsample[::-1]
+
+        # modules
+        self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_downsample, dropout)
+        self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+        self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+        self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+                                 attn_scales, self.temperal_upsample, dropout)
+
+    def forward(self, x):
+        mu, log_var = self.encode(x)
+        z = self.reparameterize(mu, log_var)
+        x_recon = self.decode(z)
+        return x_recon, mu, log_var
+
+    def encode(self, x, scale, keep_cache: bool = False):
+        """Encode video frames to latents.
+
+        Args:
+            keep_cache: if True, do not call clear_cache() at start. The caller
+                must drive cache state across calls. When True, x is treated as
+                a single 4-frame chunk (or 1-frame if cache is empty).
+        """
+        if not keep_cache:
+            self.clear_cache()
+        ## cache
+        t = x.shape[2]
+        iter_ = 1 + (t - 1) // 4
+
+        for i in range(iter_):
+            self._enc_conv_idx = [0]
+            if i == 0:
+                out = self.encoder(x[:, :, :1, :, :],
+                                   feat_cache=self._enc_feat_map,
+                                   feat_idx=self._enc_conv_idx)
+            else:
+                out_ = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+                                    feat_cache=self._enc_feat_map,
+                                    feat_idx=self._enc_conv_idx)
+                out = torch.cat([out, out_], 2)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale]
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=mu.dtype, device=mu.device)
+            mu = (mu - scale[0]) * scale[1]
+        return mu
+
+    def streaming_encode_step(self, x_chunk, scale):
+        """Encode one chunk of video frames using the live feat_cache state.
+
+        Caller invariant:
+          - First call after reset_encode_cache(): x_chunk has 1 frame
+          - Subsequent calls: x_chunk has 4 frames each
+        Output: 1 latent frame per call.
+
+        Caller is responsible for invoking clear_cache() (or
+        clear_encode_cache()) once before the first chunk; cache is preserved
+        across calls.
+        """
+        self._enc_conv_idx = [0]
+        out = self.encoder(x_chunk, feat_cache=self._enc_feat_map,
+                           feat_idx=self._enc_conv_idx)
+        mu, log_var = self.conv1(out).chunk(2, dim=1)
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale]
+            mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=mu.dtype, device=mu.device)
+            mu = (mu - scale[0]) * scale[1]
+        return mu
+
+    def decode(self, z, scale, keep_cache: bool = False):
+        """Decode latents to video frames.
+
+        Args:
+            keep_cache: if True, do not call clear_cache() at start. Use for
+                streaming inference where the caller drives state across calls.
+                The caller is responsible for invoking clear_cache() before the
+                first chunk and not in between.
+        """
+        if not keep_cache:
+            self.clear_cache()
+        # z: [b,c,t,h,w]
+        if isinstance(scale[0], torch.Tensor):
+            scale = [s.to(dtype=z.dtype, device=z.device) for s in scale]
+            z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+                1, self.z_dim, 1, 1, 1)
+        else:
+            scale = scale.to(dtype=z.dtype, device=z.device)
+            z = z / scale[1] + scale[0]
+        iter_ = z.shape[2]
+        x = self.conv2(z)
+        for i in range(iter_):
+            self._conv_idx = [0]
+            if i == 0:
+                out, self._feat_map = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+            else:
+                out_, self._feat_map = self.decoder(
+                    x[:, :, i:i + 1, :, :],
+                    feat_cache=self._feat_map,
+                    feat_idx=self._conv_idx)
+                out = torch.cat([out, out_], 2)
+        return out
+
+    def reparameterize(self, mu, log_var):
+        std = torch.exp(0.5 * log_var)
+        eps = torch.randn_like(std)
+        return eps * std + mu
+
+    def sample(self, imgs, deterministic=False):
+        mu, log_var = self.encode(imgs)
+        if deterministic:
+            return mu
+        std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+        return mu + std * torch.randn_like(std)
+
+    def clear_cache(self):
+        self._conv_num = count_conv3d(self.decoder)
+        self._conv_idx = [0]
+        self._feat_map = [None] * self._conv_num
+        # cache encode
+        self._enc_conv_num = count_conv3d(self.encoder)
+        self._enc_conv_idx = [0]
+        self._enc_feat_map = [None] * self._enc_conv_num
+
+
+class WanVideoVAE(nn.Module):
+
+    def __init__(self, z_dim=16):
+        super().__init__()
+
+        mean = [
+            -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+            0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+        ]
+        std = [
+            2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+            3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+        ]
+        self.mean = torch.tensor(mean)
+        self.std = torch.tensor(std)
+        self.scale = [self.mean, 1.0 / self.std]
+
+        # init model
+        self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False)
+        self.upsampling_factor = 8
+
+
+    def build_1d_mask(self, length, left_bound, right_bound, border_width):
+        x = torch.ones((length,))
+        if not left_bound:
+            x[:border_width] = (torch.arange(border_width) + 1) / border_width
+        if not right_bound:
+            x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / border_width, dims=(0,))
+        return x
+
+
+    def build_mask(self, data, is_bound, border_width):
+        _, _, _, H, W = data.shape
+        h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0])
+        w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1])
+
+        h = repeat(h, "H -> H W", H=H, W=W)
+        w = repeat(w, "W -> H W", H=H, W=W)
+
+        mask = torch.stack([h, w]).min(dim=0).values
+        mask = rearrange(mask, "H W -> 1 1 1 H W")
+        return mask
+
+
+    def tiled_decode(self, hidden_states, device, tile_size, tile_stride):
+        _, _, T, H, W = hidden_states.shape
+        size_h, size_w = tile_size
+        stride_h, stride_w = tile_stride
+
+        # Split tasks
+        tasks = []
+        for h in range(0, H, stride_h):
+            if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue
+            for w in range(0, W, stride_w):
+                if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue
+                h_, w_ = h + size_h, w + size_w
+                tasks.append((h, h_, w, w_))
+
+        data_device = "cpu"
+        computation_device = device
+
+        out_T = T * 4 - 3
+        weight = torch.zeros((1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device)
+        values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device)
+
+        for h, h_, w, w_ in tasks:
+            hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device)
+            hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device)
+
+            mask = self.build_mask(
+                hidden_states_batch,
+                is_bound=(h==0, h_>=H, w==0, w_>=W),
+                border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor)
+            ).to(dtype=hidden_states.dtype, device=data_device)
+
+            target_h = h * self.upsampling_factor
+            target_w = w * self.upsampling_factor
+            values[
+                :,
+                :,
+                :,
+                target_h:target_h + hidden_states_batch.shape[3],
+                target_w:target_w + hidden_states_batch.shape[4],
+            ] += hidden_states_batch * mask
+            weight[
+                :,
+                :,
+                :,
+                target_h: target_h + hidden_states_batch.shape[3],
+                target_w: target_w + hidden_states_batch.shape[4],
+            ] += mask
+        values = values / weight
+        values = values.clamp_(-1, 1)
+        return values
+
+
+    def tiled_encode(self, video, device, tile_size, tile_stride):
+        _, _, T, H, W = video.shape
+        size_h, size_w = tile_size
+        stride_h, stride_w = tile_stride
+
+        # Split tasks
+        tasks = []
+        for h in range(0, H, stride_h):
+            if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue
+            for w in range(0, W, stride_w):
+                if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue
+                h_, w_ = h + size_h, w + size_w
+                tasks.append((h, h_, w, w_))
+
+        data_device = "cpu"
+        computation_device = device
+
+        out_T = (T + 3) // 4
+        weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device)
+        values = torch.zeros((1, 16, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device)
+
+        for h, h_, w, w_ in tasks:
+            hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device)
+            hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device)
+
+            mask = self.build_mask(
+                hidden_states_batch,
+                is_bound=(h==0, h_>=H, w==0, w_>=W),
+                border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor)
+            ).to(dtype=video.dtype, device=data_device)
+
+            target_h = h // self.upsampling_factor
+            target_w = w // self.upsampling_factor
+            values[
+                :,
+                :,
+                :,
+                target_h:target_h + hidden_states_batch.shape[3],
+                target_w:target_w + hidden_states_batch.shape[4],
+            ] += hidden_states_batch * mask
+            weight[
+                :,
+                :,
+                :,
+                target_h: target_h + hidden_states_batch.shape[3],
+                target_w: target_w + hidden_states_batch.shape[4],
+            ] += mask
+        values = values / weight
+        return values
+
+
+    def single_encode(self, video, device):
+        video = video.to(device)
+        x = self.model.encode(video, self.scale)
+        return x
+
+
+    def single_decode(self, hidden_state, device):
+        hidden_state = hidden_state.to(device)
+        video = self.model.decode(hidden_state, self.scale)
+        return video.clamp_(-1, 1)
+
+    # ------------------------------------------------------------------
+    # Streaming encode/decode API
+    # ------------------------------------------------------------------
+    def reset_encode_cache(self):
+        """Reset encoder feat_cache before starting a new streaming encode."""
+        self.model.clear_cache()
+
+    def save_encode_cache_state(self):
+        """Snapshot the current encoder feat_cache (shallow list of tensors)."""
+        return list(self.model._enc_feat_map)
+
+    def load_encode_cache_state(self, state):
+        """Restore a previously snapshotted encoder feat_cache."""
+        if state is None:
+            self.reset_encode_cache()
+        else:
+            self.model._enc_feat_map = list(state)
+
+    def streaming_encode_chunk(self, video_chunk, device):
+        """Encode one chunk of video frames while preserving feat_cache state.
+
+        Caller must invoke reset_encode_cache() (or reset_decode_cache(), they
+        share state) once before the first chunk. The chunking convention
+        matches Wan VAE's internal pattern:
+          - First call after reset: video_chunk = [c, 1, h, w]   (1 frame)
+          - Subsequent calls:        video_chunk = [c, 4, h, w]   (4 frames)
+        Each call returns 1 latent frame.
+
+        Args:
+            video_chunk: [c, t_chunk, h, w]  (t_chunk == 1 first, 4 after)
+            device: target compute device
+
+        Returns:
+            latent: [1, z_dim, 1, h_lat, w_lat]
+        """
+        x = video_chunk.unsqueeze(0).to(device)
+        return self.model.streaming_encode_step(x, self.scale)
+
+    def reset_decode_cache(self):
+        """Reset the decoder feat_cache before starting a new streaming decode."""
+        self.model.clear_cache()
+
+    def streaming_decode_chunk(self, hidden_state_chunk, device):
+        """Decode one chunk of latents while preserving feat_cache state.
+
+        The caller must invoke reset_decode_cache() once before the first chunk.
+        Subsequent calls reuse the cached temporal context, so the result is
+        bit-identical (modulo float order) to a single full-length decode of
+        the concatenated chunks.
+
+        Args:
+            hidden_state_chunk: [c, t_chunk, h, w] latent tensor
+            device: target device for compute
+
+        Returns:
+            video: [1, 3, t_video, H, W] tensor in [-1, 1]; t_video equals the
+            number of video frames produced by this chunk under Wan VAE 4x
+            temporal upsampling (1 latent → 1 video frame on the very first
+            call, then 4 per latent thereafter).
+        """
+        hidden_state = hidden_state_chunk.unsqueeze(0).to(device)
+        video = self.model.decode(hidden_state, self.scale, keep_cache=True)
+        return video.clamp_(-1, 1)
+
+
+    def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)):
+
+        videos = [video.to("cpu") for video in videos]
+        hidden_states = []
+        for video in videos:
+            video = video.unsqueeze(0)
+            if tiled:
+                tile_size = (tile_size[0] * 8, tile_size[1] * 8)
+                tile_stride = (tile_stride[0] * 8, tile_stride[1] * 8)
+                hidden_state = self.tiled_encode(video, device, tile_size, tile_stride)
+            else:
+                hidden_state = self.single_encode(video, device)
+            hidden_state = hidden_state.squeeze(0)
+            hidden_states.append(hidden_state)
+        hidden_states = torch.stack(hidden_states)
+        return hidden_states
+
+
+    def decode(self, hidden_states, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)):
+        hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states]
+        videos = []
+        for hidden_state in hidden_states:
+            hidden_state = hidden_state.unsqueeze(0)
+            if tiled:
+                video = self.tiled_decode(hidden_state, device, tile_size, tile_stride)
+            else:
+                video = self.single_decode(hidden_state, device)
+            video = video.squeeze(0)
+            videos.append(video)
+        videos = torch.stack(videos)
+        return videos
+
+
+    @staticmethod
+    def state_dict_converter():
+        return WanVideoVAEStateDictConverter()
+
+
+class WanVideoVAEStateDictConverter:
+
+    def __init__(self):
+        pass
+
+    def from_civitai(self, state_dict):
+        state_dict_ = {}
+        if 'model_state' in state_dict:
+            state_dict = state_dict['model_state']
+        for name in state_dict:
+            state_dict_['model.' + name] = state_dict[name]
+        return state_dict_

OmniAvatar/models/wav2vec.py

@@ -0,0 +1,209 @@
+# pylint: disable=R0901
+# src/models/wav2vec.py
+
+"""
+This module defines the Wav2Vec model, which is a pre-trained model for speech recognition and understanding.
+It inherits from the Wav2Vec2Model class in the transformers library and provides additional functionalities
+such as feature extraction and encoding.
+
+Classes:
+    Wav2VecModel: Inherits from Wav2Vec2Model and adds additional methods for feature extraction and encoding.
+
+Functions:
+    linear_interpolation: Interpolates the features based on the sequence length.
+"""
+
+import torch.nn.functional as F
+from transformers import Wav2Vec2Model
+from transformers.modeling_outputs import BaseModelOutput
+
+
+class Wav2VecModel(Wav2Vec2Model):
+    """
+    Wav2VecModel is a custom model class that extends the Wav2Vec2Model class from the transformers library. 
+    It inherits all the functionality of the Wav2Vec2Model and adds additional methods for feature extraction and encoding.
+    ...
+
+    Attributes:
+        base_model (Wav2Vec2Model): The base Wav2Vec2Model object.
+
+    Methods:
+        forward(input_values, seq_len, attention_mask=None, mask_time_indices=None
+        , output_attentions=None, output_hidden_states=None, return_dict=None):
+            Forward pass of the Wav2VecModel. 
+            It takes input_values, seq_len, and other optional parameters as input and returns the output of the base model.
+
+        feature_extract(input_values, seq_len):
+            Extracts features from the input_values using the base model.
+
+        encode(extract_features, attention_mask=None, mask_time_indices=None, output_attentions=None, output_hidden_states=None, return_dict=None):
+            Encodes the extracted features using the base model and returns the encoded features.
+    """
+    def forward(
+        self,
+        input_values,
+        seq_len,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        """
+        Forward pass of the Wav2Vec model.
+
+        Args:
+            self: The instance of the model.
+            input_values: The input values (waveform) to the model.
+            seq_len: The sequence length of the input values.
+            attention_mask: Attention mask to be used for the model.
+            mask_time_indices: Mask indices to be used for the model.
+            output_attentions: If set to True, returns attentions.
+            output_hidden_states: If set to True, returns hidden states.
+            return_dict: If set to True, returns a BaseModelOutput instead of a tuple.
+
+        Returns:
+            The output of the Wav2Vec model.
+        """
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+    def feature_extract(
+        self,
+        input_values,
+        seq_len,
+    ):
+        """
+        Extracts features from the input values and returns the extracted features.
+
+        Parameters:
+        input_values (torch.Tensor): The input values to be processed.
+        seq_len (torch.Tensor): The sequence lengths of the input values.
+
+        Returns:
+        extracted_features (torch.Tensor): The extracted features from the input values.
+        """
+        extract_features = self.feature_extractor(input_values)
+        extract_features = extract_features.transpose(1, 2)
+        extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+        return extract_features
+
+    def encode(
+        self,
+        extract_features,
+        attention_mask=None,
+        mask_time_indices=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        """
+        Encodes the input features into the output space.
+
+        Args:
+            extract_features (torch.Tensor): The extracted features from the audio signal.
+            attention_mask (torch.Tensor, optional): Attention mask to be used for padding.
+            mask_time_indices (torch.Tensor, optional): Masked indices for the time dimension.
+            output_attentions (bool, optional): If set to True, returns the attention weights.
+            output_hidden_states (bool, optional): If set to True, returns all hidden states.
+            return_dict (bool, optional): If set to True, returns a BaseModelOutput instead of the tuple.
+
+        Returns:
+            The encoded output features.
+        """
+        self.config.output_attentions = True
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is not None:
+            # compute reduced attention_mask corresponding to feature vectors
+            attention_mask = self._get_feature_vector_attention_mask(
+                extract_features.shape[1], attention_mask, add_adapter=False
+            )
+
+        hidden_states, extract_features = self.feature_projection(extract_features)
+        hidden_states = self._mask_hidden_states(
+            hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+        )
+
+        encoder_outputs = self.encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = encoder_outputs[0]
+
+        if self.adapter is not None:
+            hidden_states = self.adapter(hidden_states)
+
+        if not return_dict:
+            return (hidden_states, ) + encoder_outputs[1:]
+        return BaseModelOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+def linear_interpolation(features, seq_len):
+    """
+    Transpose the features to interpolate linearly.
+
+    Args:
+        features (torch.Tensor): The extracted features to be interpolated.
+        seq_len (torch.Tensor): The sequence lengths of the features.
+
+    Returns:
+        torch.Tensor: The interpolated features.
+    """
+    features = features.transpose(1, 2)
+    output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+    return output_features.transpose(1, 2)

OmniAvatar/prompters/__init__.py

@@ -0,0 +1 @@
+from .wan_prompter import WanPrompter

OmniAvatar/prompters/base_prompter.py

@@ -0,0 +1,70 @@
+from ..models.model_manager import ModelManager
+import torch
+
+
+
+def tokenize_long_prompt(tokenizer, prompt, max_length=None):
+    # Get model_max_length from self.tokenizer
+    length = tokenizer.model_max_length if max_length is None else max_length
+
+    # To avoid the warning. set self.tokenizer.model_max_length to +oo.
+    tokenizer.model_max_length = 99999999
+
+    # Tokenize it!
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+
+    # Determine the real length.
+    max_length = (input_ids.shape[1] + length - 1) // length * length
+
+    # Restore tokenizer.model_max_length
+    tokenizer.model_max_length = length
+    
+    # Tokenize it again with fixed length.
+    input_ids = tokenizer(
+        prompt,
+        return_tensors="pt",
+        padding="max_length",
+        max_length=max_length,
+        truncation=True
+    ).input_ids
+
+    # Reshape input_ids to fit the text encoder.
+    num_sentence = input_ids.shape[1] // length
+    input_ids = input_ids.reshape((num_sentence, length))
+    
+    return input_ids
+
+
+
+class BasePrompter:
+    def __init__(self):
+        self.refiners = []
+        self.extenders = []
+
+
+    def load_prompt_refiners(self, model_manager: ModelManager, refiner_classes=[]):
+        for refiner_class in refiner_classes:
+            refiner = refiner_class.from_model_manager(model_manager)
+            self.refiners.append(refiner)
+    
+    def load_prompt_extenders(self,model_manager:ModelManager,extender_classes=[]):
+        for extender_class in extender_classes:
+            extender = extender_class.from_model_manager(model_manager)
+            self.extenders.append(extender)
+
+
+    @torch.no_grad()
+    def process_prompt(self, prompt, positive=True):
+        if isinstance(prompt, list):
+            prompt = [self.process_prompt(prompt_, positive=positive) for prompt_ in prompt]
+        else:
+            for refiner in self.refiners:
+                prompt = refiner(prompt, positive=positive)
+        return prompt
+
+    @torch.no_grad()
+    def extend_prompt(self, prompt:str, positive=True):
+        extended_prompt = dict(prompt=prompt)
+        for extender in self.extenders:
+            extended_prompt = extender(extended_prompt)
+        return extended_prompt

OmniAvatar/prompters/wan_prompter.py

@@ -0,0 +1,109 @@
+from .base_prompter import BasePrompter
+from ..models.wan_video_text_encoder import WanTextEncoder
+from transformers import AutoTokenizer
+import os, torch
+import ftfy
+import html
+import string
+import regex as re
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+def canonicalize(text, keep_punctuation_exact_string=None):
+    text = text.replace('_', ' ')
+    if keep_punctuation_exact_string:
+        text = keep_punctuation_exact_string.join(
+            part.translate(str.maketrans('', '', string.punctuation))
+            for part in text.split(keep_punctuation_exact_string))
+    else:
+        text = text.translate(str.maketrans('', '', string.punctuation))
+    text = text.lower()
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
+
+
+class HuggingfaceTokenizer:
+
+    def __init__(self, name, seq_len=None, clean=None, **kwargs):
+        assert clean in (None, 'whitespace', 'lower', 'canonicalize')
+        self.name = name
+        self.seq_len = seq_len
+        self.clean = clean
+
+        # init tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs)
+        self.vocab_size = self.tokenizer.vocab_size
+
+    def __call__(self, sequence, **kwargs):
+        return_mask = kwargs.pop('return_mask', False)
+
+        # arguments
+        _kwargs = {'return_tensors': 'pt'}
+        if self.seq_len is not None:
+            _kwargs.update({
+                'padding': 'max_length',
+                'truncation': True,
+                'max_length': self.seq_len
+            })
+        _kwargs.update(**kwargs)
+
+        # tokenization
+        if isinstance(sequence, str):
+            sequence = [sequence]
+        if self.clean:
+            sequence = [self._clean(u) for u in sequence]
+        ids = self.tokenizer(sequence, **_kwargs)
+
+        # output
+        if return_mask:
+            return ids.input_ids, ids.attention_mask
+        else:
+            return ids.input_ids
+
+    def _clean(self, text):
+        if self.clean == 'whitespace':
+            text = whitespace_clean(basic_clean(text))
+        elif self.clean == 'lower':
+            text = whitespace_clean(basic_clean(text)).lower()
+        elif self.clean == 'canonicalize':
+            text = canonicalize(basic_clean(text))
+        return text
+
+
+class WanPrompter(BasePrompter):
+
+    def __init__(self, tokenizer_path=None, text_len=512):
+        super().__init__()
+        self.text_len = text_len
+        self.text_encoder = None
+        self.fetch_tokenizer(tokenizer_path)
+        
+    def fetch_tokenizer(self, tokenizer_path=None):
+        if tokenizer_path is not None:
+            self.tokenizer = HuggingfaceTokenizer(name=tokenizer_path, seq_len=self.text_len, clean='whitespace')
+
+    def fetch_models(self, text_encoder: WanTextEncoder = None):
+        self.text_encoder = text_encoder
+
+    def encode_prompt(self, prompt, positive=True, device="cuda"):
+        prompt = self.process_prompt(prompt, positive=positive)
+        
+        ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True)
+        ids = ids.to(device)
+        mask = mask.to(device)
+        seq_lens = mask.gt(0).sum(dim=1).long()
+        prompt_emb = self.text_encoder(ids, mask)
+        for i, v in enumerate(seq_lens):
+            prompt_emb[:, v:] = 0
+        return prompt_emb

OmniAvatar/utils/args_config.py

@@ -0,0 +1,122 @@
+import json
+import os
+import argparse
+import yaml
+args = None
+
+def parse_hp_string(hp_string):
+    result = {}
+    for pair in hp_string.split(','):
+        if not pair:
+            continue
+        key, value = pair.split('=')
+        try:
+            # 自动转换为 int / float / str
+            ori_value = value
+            value = float(value)
+            if '.' not in str(ori_value):
+                value = int(value)
+        except ValueError:
+            pass
+
+        if value in ['true', 'True']:
+            value = True
+        if value in ['false', 'False']:
+            value = False
+        if '.' in key:
+            keys = key.split('.')
+            keys = keys
+            current = result
+            for key in keys[:-1]:
+                if key not in current or not isinstance(current[key], dict):
+                    current[key] = {}
+                current = current[key]
+            current[keys[-1]] = value
+        else:
+            result[key.strip()] = value
+    return result
+
+def parse_args():
+    global args
+    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument("--config", type=str, required=True, help="Path to YAML config file.")
+    
+    # 定义 argparse 参数
+    parser.add_argument("--exp_path", type=str, help="Path to save the model.")
+    parser.add_argument("--input_file", type=str, help="Path to inference txt.")
+    parser.add_argument("--debug", action='store_true', default=None)
+    parser.add_argument("--infer", action='store_true')
+    parser.add_argument("-hp", "--hparams", type=str, default="")
+
+    args = parser.parse_args()
+
+    # 读取 YAML 配置（如果提供了 --config 参数）
+    if args.config:
+        with open(args.config, "r") as f:
+            yaml_config = yaml.safe_load(f)
+        
+        # 遍历 YAML 配置，将其添加到 args（如果 argparse 里没有定义）
+        for key, value in yaml_config.items():
+            if not hasattr(args, key):  # argparse 没有的参数
+                setattr(args, key, value)
+            elif getattr(args, key) is None:  # argparse 有但值为空
+                setattr(args, key, value)
+
+    args.rank = int(os.getenv("RANK", "0"))
+    args.world_size = int(os.getenv("WORLD_SIZE", "1"))
+    args.local_rank = int(os.getenv("LOCAL_RANK", "0"))  # torchrun
+    args.device = f'cuda:{args.local_rank}'
+    args.num_nodes = int(os.getenv("NNODES", "1"))
+    debug = args.debug
+
+    # Apply hparams BEFORE reload_cfg so -hp exp_path=... takes effect first
+    if len(args.hparams) > 0:
+        hp_dict = parse_hp_string(args.hparams)
+        for key, value in hp_dict.items():
+            if not hasattr(args, key):
+                setattr(args, key, value)
+            else:
+                if isinstance(value, dict):
+                    ori_v = getattr(args, key)
+                    ori_v.update(value)
+                    setattr(args, key, ori_v)
+                else:
+                    setattr(args, key, value)
+
+    if not os.path.exists(args.exp_path):
+        args.exp_path = f'checkpoints/{args.exp_path}'
+
+    if hasattr(args, 'reload_cfg') and args.reload_cfg:
+        # 重新加载配置文件
+        conf_path = os.path.join(args.exp_path, "config.json")
+        if os.path.exists(conf_path):
+            print('| Reloading config from:', conf_path)
+            args = reload(args, conf_path)
+    args.debug = debug
+    dict_args = convert_namespace_to_dict(args)
+    if args.local_rank == 0:
+        print(dict_args)
+    return args
+
+def reload(args, conf_path):
+    """重新加载配置文件,不覆盖已有的参数"""
+    with open(conf_path, "r") as f:
+        yaml_config = yaml.safe_load(f)
+    # 遍历 YAML 配置，将其添加到 args（如果 argparse 里没有定义）
+    for key, value in yaml_config.items():
+        if not hasattr(args, key):  # argparse 没有的参数
+            setattr(args, key, value)
+        elif getattr(args, key) is None:  # argparse 有但值为空
+            setattr(args, key, value)
+    return args
+
+def convert_namespace_to_dict(namespace):
+    """将 argparse.Namespace 转为字典，并处理不可序列化对象"""
+    result = {}
+    for key, value in vars(namespace).items():
+        try:
+            json.dumps(value)  # 检查是否可序列化
+            result[key] = value
+        except (TypeError, OverflowError):
+            result[key] = str(value)  # 将不可序列化的对象转为字符串表示
+    return result

OmniAvatar/utils/io_utils.py

@@ -0,0 +1,245 @@
+import subprocess
+import torch, os
+from safetensors import safe_open
+from OmniAvatar.utils.args_config import args
+from contextlib import contextmanager
+
+import re
+import tempfile
+import numpy as np
+import imageio
+from glob import glob 
+import soundfile as sf
+from einops import rearrange
+import hashlib
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+@contextmanager
+def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False):
+    
+    old_register_parameter = torch.nn.Module.register_parameter
+    if include_buffers:
+        old_register_buffer = torch.nn.Module.register_buffer
+    
+    def register_empty_parameter(module, name, param):
+        old_register_parameter(module, name, param)
+        if param is not None:
+            param_cls = type(module._parameters[name])
+            kwargs = module._parameters[name].__dict__
+            kwargs["requires_grad"] = param.requires_grad
+            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)
+
+    def register_empty_buffer(module, name, buffer, persistent=True):
+        old_register_buffer(module, name, buffer, persistent=persistent)
+        if buffer is not None:
+            module._buffers[name] = module._buffers[name].to(device)
+            
+    def patch_tensor_constructor(fn):
+        def wrapper(*args, **kwargs):
+            kwargs["device"] = device
+            return fn(*args, **kwargs)
+
+        return wrapper
+    
+    if include_buffers:
+        tensor_constructors_to_patch = {
+            torch_function_name: getattr(torch, torch_function_name)
+            for torch_function_name in ["empty", "zeros", "ones", "full"]
+        }
+    else:
+        tensor_constructors_to_patch = {}
+    
+    try:
+        torch.nn.Module.register_parameter = register_empty_parameter
+        if include_buffers:
+            torch.nn.Module.register_buffer = register_empty_buffer
+        for torch_function_name in tensor_constructors_to_patch.keys():
+            setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name)))
+        yield
+    finally:
+        torch.nn.Module.register_parameter = old_register_parameter
+        if include_buffers:
+            torch.nn.Module.register_buffer = old_register_buffer
+        for torch_function_name, old_torch_function in tensor_constructors_to_patch.items():
+            setattr(torch, torch_function_name, old_torch_function)
+
+def load_state_dict_from_folder(file_path, torch_dtype=None):
+    state_dict = {}
+    for file_name in os.listdir(file_path):
+        if "." in file_name and file_name.split(".")[-1] in [
+            "safetensors", "bin", "ckpt", "pth", "pt"
+        ]:
+            state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype))
+    return state_dict
+
+
+def load_state_dict(file_path, torch_dtype=None):
+    if file_path.endswith(".safetensors"):
+        return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype)
+    else:
+        return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype)
+
+
+def load_state_dict_from_safetensors(file_path, torch_dtype=None):
+    state_dict = {}
+    with safe_open(file_path, framework="pt", device="cpu") as f:
+        for k in f.keys():
+            state_dict[k] = f.get_tensor(k)
+            if torch_dtype is not None:
+                state_dict[k] = state_dict[k].to(torch_dtype)
+    return state_dict
+
+
+def load_state_dict_from_bin(file_path, torch_dtype=None):
+    state_dict = torch.load(file_path, map_location="cpu", weights_only=True)
+    if torch_dtype is not None:
+        for i in state_dict:
+            if isinstance(state_dict[i], torch.Tensor):
+                state_dict[i] = state_dict[i].to(torch_dtype)
+    return state_dict
+
+def smart_load_weights(model, ckpt_state_dict):
+    model_state_dict = model.state_dict()
+    new_state_dict = {}
+
+    for name, param in model_state_dict.items():
+        if name in ckpt_state_dict:
+            ckpt_param = ckpt_state_dict[name]
+            if param.shape == ckpt_param.shape:
+                new_state_dict[name] = ckpt_param
+            else:
+                # 自动修剪维度以匹配
+                if all(p >= c for p, c in zip(param.shape, ckpt_param.shape)):
+                    print(f"[Truncate] {name}: ckpt {ckpt_param.shape} -> model {param.shape}")
+                    # 创建新张量，拷贝旧数据
+                    new_param = param.clone()
+                    slices = tuple(slice(0, s) for s in ckpt_param.shape)
+                    new_param[slices] = ckpt_param
+                    new_state_dict[name] = new_param
+                else:
+                    print(f"[Skip] {name}: ckpt {ckpt_param.shape} is larger than model {param.shape}")
+
+    # 更新 state_dict，只更新那些匹配的
+    missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, assign=True, strict=False)
+    return model, missing_keys, unexpected_keys
+
+def save_wav(audio, audio_path):
+    if isinstance(audio, torch.Tensor):
+        audio = audio.float().detach().cpu().numpy()
+    
+    if audio.ndim == 1:
+        audio = np.expand_dims(audio, axis=0)  # (1, samples)
+
+    sf.write(audio_path, audio.T, 16000)
+
+    return True
+
+def save_video_as_grid_and_mp4(video_batch: torch.Tensor, save_path: str, fps: float = 5,prompt=None, prompt_path=None, audio=None, audio_path=None, prefix=None):
+    os.makedirs(save_path, exist_ok=True)
+    out_videos = []
+    with tempfile.TemporaryDirectory() as tmp_path:
+        for i, vid in enumerate(video_batch):
+            gif_frames = []
+            for frame in vid:
+                frame = rearrange(frame, "c h w -> h w c")
+                frame = (255.0 * frame).cpu().numpy().astype(np.uint8)
+                gif_frames.append(frame)
+            if prefix is not None:
+                now_save_path = os.path.join(save_path, f"{prefix}_{i:03d}.mp4")
+                tmp_save_path = os.path.join(tmp_path, f"{prefix}_{i:03d}.mp4")
+            else:
+                now_save_path = os.path.join(save_path, f"{i:03d}.mp4")
+                tmp_save_path = os.path.join(tmp_path, f"{i:03d}.mp4")
+            with imageio.get_writer(tmp_save_path, fps=fps) as writer:
+                for frame in gif_frames:
+                    writer.append_data(frame)
+            subprocess.run([f"cp {tmp_save_path} {now_save_path}"], check=True, shell=True)
+            print(f'save res video to : {now_save_path}')
+            if audio is not None or audio_path is not None:
+                if audio is not None:
+                    audio_path = os.path.join(tmp_path, f"{i:06d}.mp3")
+                    save_wav(audio[i], audio_path)
+                # cmd = f'/usr/bin/ffmpeg -i {tmp_save_path} -i {audio_path} -v quiet -c:v copy -c:a libmp3lame -strict experimental {tmp_save_path[:-4]}_wav.mp4 -y'
+                cmd = f'/usr/bin/ffmpeg -i {tmp_save_path} -i {audio_path} -v quiet -map 0:v:0 -map 1:a:0 -c:v copy -c:a aac {tmp_save_path[:-4]}_wav.mp4 -y'
+                subprocess.check_call(cmd, stdout=None, stdin=subprocess.PIPE, shell=True)
+                subprocess.run([f"cp {tmp_save_path[:-4]}_wav.mp4 {now_save_path[:-4]}_wav.mp4"], check=True, shell=True)
+                os.remove(now_save_path)
+            if prompt is not None and prompt_path is not None:
+                with open(prompt_path, "w") as f:
+                    f.write(prompt)
+            out_videos.append(now_save_path)
+    return out_videos
+
+def is_zero_stage_3(trainer):
+    strategy = getattr(trainer, "strategy", None)
+    if strategy and hasattr(strategy, "model"):
+        ds_engine = strategy.model
+        stage = ds_engine.config.get("zero_optimization", {}).get("stage", 0)
+        return stage == 3
+    return False
+
+def hash_state_dict_keys(state_dict, with_shape=True):
+    keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape)
+    keys_str = keys_str.encode(encoding="UTF-8")
+    return hashlib.md5(keys_str).hexdigest()
+
+def split_state_dict_with_prefix(state_dict):
+    keys = sorted([key for key in state_dict if isinstance(key, str)])
+    prefix_dict = {}
+    for key in  keys:
+        prefix = key if "." not in key else key.split(".")[0]
+        if prefix not in prefix_dict:
+            prefix_dict[prefix] = []
+        prefix_dict[prefix].append(key)
+    state_dicts = []
+    for prefix, keys in prefix_dict.items():
+        sub_state_dict = {key: state_dict[key] for key in keys}
+        state_dicts.append(sub_state_dict)
+    return state_dicts
+
+def hash_state_dict_keys(state_dict, with_shape=True):
+    keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape)
+    keys_str = keys_str.encode(encoding="UTF-8")
+    return hashlib.md5(keys_str).hexdigest()
+
+def split_state_dict_with_prefix(state_dict):
+    keys = sorted([key for key in state_dict if isinstance(key, str)])
+    prefix_dict = {}
+    for key in  keys:
+        prefix = key if "." not in key else key.split(".")[0]
+        if prefix not in prefix_dict:
+            prefix_dict[prefix] = []
+        prefix_dict[prefix].append(key)
+    state_dicts = []
+    for prefix, keys in prefix_dict.items():
+        sub_state_dict = {key: state_dict[key] for key in keys}
+        state_dicts.append(sub_state_dict)
+    return state_dicts
+
+def search_for_files(folder, extensions):
+    files = []
+    if os.path.isdir(folder):
+        for file in sorted(os.listdir(folder)):
+            files += search_for_files(os.path.join(folder, file), extensions)
+    elif os.path.isfile(folder):
+        for extension in extensions:
+            if folder.endswith(extension):
+                files.append(folder)
+                break
+    return files
+
+def convert_state_dict_keys_to_single_str(state_dict, with_shape=True):
+    keys = []
+    for key, value in state_dict.items():
+        if isinstance(key, str):
+            if isinstance(value, torch.Tensor):
+                if with_shape:
+                    shape = "_".join(map(str, list(value.shape)))
+                    keys.append(key + ":" + shape)
+                keys.append(key)
+            elif isinstance(value, dict):
+                keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape))
+    keys.sort()
+    keys_str = ",".join(keys)
+    return keys_str

OmniAvatar/utils/latentsync/__init__.py

@@ -0,0 +1,19 @@
+"""
+LatentSync face detection, affine alignment, and compositing utilities.
+
+Provides InsightFace-based face detection, Procrustes affine alignment to 512x512,
+and inverse affine compositing with soft blending for video-to-video lip sync.
+
+Adapted from LatentSync.
+"""
+
+from .face_detector import FaceDetector
+from .affine_transform import AlignRestore
+from .image_processor import ImageProcessor, load_fixed_mask
+
+__all__ = [
+    "FaceDetector",
+    "AlignRestore",
+    "ImageProcessor",
+    "load_fixed_mask",
+]

OmniAvatar/utils/latentsync/affine_transform.py

@@ -0,0 +1,151 @@
+# Adapted from https://github.com/guanjz20/StyleSync/blob/main/utils.py
+# Adapted from LatentSync
+
+import numpy as np
+import cv2
+import torch
+from einops import rearrange
+import kornia
+
+
+class AlignRestore(object):
+    def __init__(self, align_points=3, resolution=256, device="cpu", dtype=torch.float16):
+        if align_points == 3:
+            self.upscale_factor = 1
+            ratio = resolution / 256 * 2.8
+            self.crop_ratio = (ratio, ratio)
+            self.face_template = np.array([[19 - 2, 30 - 10], [56 + 2, 30 - 10], [37.5, 45 - 5]])
+            self.face_template = self.face_template * ratio
+            self.face_size = (int(75 * self.crop_ratio[0]), int(100 * self.crop_ratio[1]))
+            self.p_bias = None
+            self.device = device
+            self.dtype = dtype
+            self.fill_value = torch.tensor([127, 127, 127], device=device, dtype=dtype)
+            self.mask = torch.ones((1, 1, self.face_size[1], self.face_size[0]), device=device, dtype=dtype)
+
+    def align_warp_face(self, img, landmarks3, smooth=True):
+        affine_matrix, self.p_bias = self.transformation_from_points(
+            landmarks3, self.face_template, smooth, self.p_bias
+        )
+
+        img = rearrange(torch.from_numpy(img).to(device=self.device, dtype=self.dtype), "h w c -> c h w").unsqueeze(0)
+        affine_matrix = torch.from_numpy(affine_matrix).to(device=self.device, dtype=self.dtype).unsqueeze(0)
+
+        cropped_face = kornia.geometry.transform.warp_affine(
+            img,
+            affine_matrix,
+            (self.face_size[1], self.face_size[0]),
+            mode="bilinear",
+            padding_mode="fill",
+            fill_value=self.fill_value,
+        )
+        cropped_face = rearrange(cropped_face.squeeze(0), "c h w -> h w c").cpu().numpy().astype(np.uint8)
+        return cropped_face, affine_matrix
+
+    def restore_img(self, input_img, face, affine_matrix):
+        """Inverse-warp composited face back onto original frame with soft blending.
+
+        Uses float32 for all intermediate computation to avoid color shifts from
+        float16 rounding (float16 cannot exactly represent all uint8 values 0-255).
+        """
+        h, w, _ = input_img.shape
+        # Use float32 for compositing precision regardless of self.dtype
+        work_dtype = torch.float32
+
+        if isinstance(affine_matrix, np.ndarray):
+            affine_matrix = torch.from_numpy(affine_matrix).to(device=self.device, dtype=work_dtype).unsqueeze(0)
+        else:
+            affine_matrix = affine_matrix.to(dtype=work_dtype)
+
+        inv_affine_matrix = kornia.geometry.transform.invert_affine_transform(affine_matrix)
+        face = face.to(device=self.device, dtype=work_dtype).unsqueeze(0)
+        fill_value = self.fill_value.to(dtype=work_dtype)
+
+        inv_face = kornia.geometry.transform.warp_affine(
+            face, inv_affine_matrix, (h, w), mode="bilinear", padding_mode="fill", fill_value=fill_value
+        ).squeeze(0)
+        inv_face = (inv_face / 2 + 0.5).clamp(0, 1) * 255
+
+        input_img = rearrange(torch.from_numpy(input_img).to(device=self.device, dtype=work_dtype), "h w c -> c h w")
+        mask_f32 = self.mask.to(dtype=work_dtype)
+        inv_mask = kornia.geometry.transform.warp_affine(
+            mask_f32, inv_affine_matrix, (h, w), padding_mode="zeros"
+        )  # (1, 1, h_up, w_up)
+
+        erosion_kernel = torch.ones(
+            (int(2 * self.upscale_factor), int(2 * self.upscale_factor)),
+            device=self.device, dtype=work_dtype,
+        )
+        inv_mask_erosion = kornia.morphology.erosion(inv_mask, erosion_kernel)
+
+        inv_mask_erosion_t = inv_mask_erosion.squeeze(0).expand_as(inv_face)
+        pasted_face = inv_mask_erosion_t * inv_face
+        total_face_area = torch.sum(inv_mask_erosion)
+        w_edge = int(total_face_area**0.5) // 20
+        erosion_radius = w_edge * 2
+
+        # Run on CPU to avoid consuming a large amount of GPU memory.
+        inv_mask_erosion = inv_mask_erosion.squeeze().cpu().numpy().astype(np.float32)
+        inv_mask_center = cv2.erode(inv_mask_erosion, np.ones((erosion_radius, erosion_radius), np.uint8))
+        inv_mask_center = torch.from_numpy(inv_mask_center).to(device=self.device, dtype=work_dtype)[None, None, ...]
+
+        blur_size = w_edge * 2 + 1
+        sigma = 0.3 * ((blur_size - 1) * 0.5 - 1) + 0.8
+        inv_soft_mask = kornia.filters.gaussian_blur2d(
+            inv_mask_center, (blur_size, blur_size), (sigma, sigma)
+        ).squeeze(0)
+        inv_soft_mask_3d = inv_soft_mask.expand_as(inv_face)
+        img_back = inv_soft_mask_3d * pasted_face + (1 - inv_soft_mask_3d) * input_img
+
+        img_back = rearrange(img_back, "c h w -> h w c").contiguous().to(dtype=torch.uint8)
+        img_back = img_back.cpu().numpy()
+        return img_back
+
+    def transformation_from_points(self, points1: torch.Tensor, points0: torch.Tensor, smooth=True, p_bias=None):
+        if isinstance(points0, np.ndarray):
+            points2 = torch.tensor(points0, device=self.device, dtype=torch.float32)
+        else:
+            points2 = points0.clone()
+
+        if isinstance(points1, np.ndarray):
+            points1_tensor = torch.tensor(points1, device=self.device, dtype=torch.float32)
+        else:
+            points1_tensor = points1.clone()
+
+        c1 = torch.mean(points1_tensor, dim=0)
+        c2 = torch.mean(points2, dim=0)
+
+        points1_centered = points1_tensor - c1
+        points2_centered = points2 - c2
+
+        s1 = torch.std(points1_centered)
+        s2 = torch.std(points2_centered)
+
+        points1_normalized = points1_centered / s1
+        points2_normalized = points2_centered / s2
+
+        covariance = torch.matmul(points1_normalized.T, points2_normalized)
+        U, S, V = torch.svd(covariance.float())
+
+        R = torch.matmul(V, U.T)
+
+        det = torch.det(R.float())
+        if det < 0:
+            V[:, -1] = -V[:, -1]
+            R = torch.matmul(V, U.T)
+
+        sR = (s2 / s1) * R
+        T = c2.reshape(2, 1) - (s2 / s1) * torch.matmul(R, c1.reshape(2, 1))
+
+        M = torch.cat((sR, T), dim=1)
+
+        if smooth:
+            bias = points2_normalized[2] - points1_normalized[2]
+            if p_bias is None:
+                p_bias = bias
+            else:
+                bias = p_bias * 0.2 + bias * 0.8
+            p_bias = bias
+            M[:, 2] = M[:, 2] + bias
+
+        return M.cpu().numpy(), p_bias

OmniAvatar/utils/latentsync/face_detector.py

@@ -0,0 +1,93 @@
+# Adapted from LatentSync
+
+import os
+
+# Disable ONNX Runtime thread affinity to suppress pthread warnings
+# MUST be set before importing insightface/onnxruntime
+os.environ['ORT_DISABLE_THREAD_AFFINITY'] = '1'
+
+from insightface.app import FaceAnalysis
+import numpy as np
+import torch
+
+INSIGHTFACE_DETECT_SIZE = 512
+
+
+class FaceDetector:
+    def __init__(self, device="cuda", insightface_root="checkpoints/auxiliary"):
+        self.app = FaceAnalysis(
+            allowed_modules=["detection", "landmark_2d_106"],
+            root=insightface_root,
+            providers=["CUDAExecutionProvider"],
+        )
+        self.app.prepare(ctx_id=cuda_to_int(device), det_size=(INSIGHTFACE_DETECT_SIZE, INSIGHTFACE_DETECT_SIZE))
+
+    def __call__(self, frame, threshold=0.5):
+        f_h, f_w, _ = frame.shape
+
+        faces = self.app.get(frame)
+
+        get_face_store = None
+        max_size = 0
+
+        if len(faces) == 0:
+            return None, None
+        else:
+            for face in faces:
+                bbox = face.bbox.astype(np.int_).tolist()
+                w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
+                if w < 50 or h < 80:
+                    continue
+                if w / h > 1.5 or w / h < 0.2:
+                    continue
+                if face.det_score < threshold:
+                    continue
+                size_now = w * h
+
+                if size_now > max_size:
+                    max_size = size_now
+                    get_face_store = face
+
+        if get_face_store is None:
+            return None, None
+        else:
+            face = get_face_store
+            lmk = np.round(face.landmark_2d_106).astype(np.int_)
+
+            halk_face_coord = np.mean([lmk[74], lmk[73]], axis=0)
+
+            sub_lmk = lmk[LMK_ADAPT_ORIGIN_ORDER]
+            halk_face_dist = np.max(sub_lmk[:, 1]) - halk_face_coord[1]
+            upper_bond = halk_face_coord[1] - halk_face_dist
+
+            x1, y1, x2, y2 = (np.min(sub_lmk[:, 0]), int(upper_bond), np.max(sub_lmk[:, 0]), np.max(sub_lmk[:, 1]))
+
+            if y2 - y1 <= 0 or x2 - x1 <= 0 or x1 < 0:
+                x1, y1, x2, y2 = face.bbox.astype(np.int_).tolist()
+
+            y2 += int((x2 - x1) * 0.1)
+            x1 -= int((x2 - x1) * 0.05)
+            x2 += int((x2 - x1) * 0.05)
+
+            x1 = max(0, x1)
+            y1 = max(0, y1)
+            x2 = min(f_w, x2)
+            y2 = min(f_h, y2)
+
+            return (x1, y1, x2, y2), lmk
+
+
+def cuda_to_int(cuda_str: str) -> int:
+    """Convert the string with format "cuda:X" to integer X."""
+    if cuda_str == "cuda":
+        return 0
+    device = torch.device(cuda_str)
+    if device.type != "cuda":
+        raise ValueError(f"Device type must be 'cuda', got: {device.type}")
+    return device.index
+
+
+LMK_ADAPT_ORIGIN_ORDER = [
+    1, 10, 12, 14, 16, 3, 5, 7, 0, 23, 21, 19, 32, 30, 28, 26,
+    17, 43, 48, 49, 51, 50, 102, 103, 104, 105, 101, 73, 74, 86,
+]

OmniAvatar/utils/latentsync/image_processor.py

@@ -0,0 +1,109 @@
+# Adapted from LatentSync
+# Copyright (c) 2024 Bytedance Ltd. and/or its affiliates
+# Licensed under the Apache License, Version 2.0
+
+from torchvision import transforms
+import cv2
+from einops import rearrange
+import torch
+import numpy as np
+from typing import Union
+import os
+from .affine_transform import AlignRestore
+from .face_detector import FaceDetector
+
+
+def load_fixed_mask(resolution: int, mask_image_path=None) -> torch.Tensor:
+    """Load the LatentSync mouth-shaped mask from PNG.
+
+    Returns [3, resolution, resolution] tensor. Values in [0, 1].
+    Semantics: 1.0 = keep (upper face), 0.0 = mask out (mouth region).
+    """
+    if mask_image_path is None:
+        # Default to mask.png in the same directory as this module
+        mask_image_path = os.path.join(os.path.dirname(__file__), "mask.png")
+    mask_image = cv2.imread(mask_image_path)
+    mask_image = cv2.cvtColor(mask_image, cv2.COLOR_BGR2RGB)
+    mask_image = cv2.resize(mask_image, (resolution, resolution), interpolation=cv2.INTER_LANCZOS4) / 255.0
+    mask_image = rearrange(torch.from_numpy(mask_image), "h w c -> c h w")
+    return mask_image
+
+
+class ImageProcessor:
+    def __init__(self, resolution: int = 512, device: str = "cpu",
+                 mask_image=None, insightface_root="checkpoints/auxiliary"):
+        self.resolution = resolution
+        self.resize = transforms.Resize(
+            (resolution, resolution), interpolation=transforms.InterpolationMode.BICUBIC, antialias=True
+        )
+        self.normalize = transforms.Normalize([0.5], [0.5], inplace=True)
+
+        self.restorer = AlignRestore(resolution=resolution, device=device)
+
+        if mask_image is None:
+            self.mask_image = load_fixed_mask(resolution)
+        else:
+            self.mask_image = mask_image
+
+        if device == "cpu":
+            self.face_detector = None
+        else:
+            self.face_detector = FaceDetector(device=device, insightface_root=insightface_root)
+
+    def affine_transform(self, image: np.ndarray):
+        """Detect face and align to resolution x resolution via affine transform.
+
+        Args:
+            image: np.ndarray [H, W, 3] uint8 RGB frame.
+
+        Returns:
+            face: torch.Tensor [C, resolution, resolution] uint8
+            box: list [x1, y1, x2, y2] in aligned coordinate space
+            affine_matrix: torch.Tensor [1, 2, 3] affine transformation matrix
+        """
+        if self.face_detector is None:
+            raise NotImplementedError("Using the CPU for face detection is not supported")
+        bbox, landmark_2d_106 = self.face_detector(image)
+        if bbox is None:
+            raise RuntimeError("Face not detected")
+
+        pt_left_eye = np.mean(landmark_2d_106[[43, 48, 49, 51, 50]], axis=0)  # left eyebrow center
+        pt_right_eye = np.mean(landmark_2d_106[101:106], axis=0)  # right eyebrow center
+        pt_nose = np.mean(landmark_2d_106[[74, 77, 83, 86]], axis=0)  # nose center
+
+        landmarks3 = np.round([pt_left_eye, pt_right_eye, pt_nose])
+
+        face, affine_matrix = self.restorer.align_warp_face(image.copy(), landmarks3=landmarks3, smooth=True)
+        box = [0, 0, face.shape[1], face.shape[0]]  # x1, y1, x2, y2
+        face = cv2.resize(face, (self.resolution, self.resolution), interpolation=cv2.INTER_LANCZOS4)
+        face = rearrange(torch.from_numpy(face), "h w c -> c h w")
+        return face, box, affine_matrix
+
+    def preprocess_fixed_mask_image(self, image: torch.Tensor, affine_transform=False):
+        if affine_transform:
+            image, _, _ = self.affine_transform(image)
+        else:
+            image = self.resize(image)
+        pixel_values = self.normalize(image / 255.0)
+        masked_pixel_values = pixel_values * self.mask_image
+        return pixel_values, masked_pixel_values, self.mask_image[0:1]
+
+    def prepare_masks_and_masked_images(self, images: Union[torch.Tensor, np.ndarray], affine_transform=False):
+        if isinstance(images, np.ndarray):
+            images = torch.from_numpy(images)
+        if images.shape[3] == 3:
+            images = rearrange(images, "f h w c -> f c h w")
+
+        results = [self.preprocess_fixed_mask_image(image, affine_transform=affine_transform) for image in images]
+
+        pixel_values_list, masked_pixel_values_list, masks_list = list(zip(*results))
+        return torch.stack(pixel_values_list), torch.stack(masked_pixel_values_list), torch.stack(masks_list)
+
+    def process_images(self, images: Union[torch.Tensor, np.ndarray]):
+        if isinstance(images, np.ndarray):
+            images = torch.from_numpy(images)
+        if images.shape[3] == 3:
+            images = rearrange(images, "f h w c -> f c h w")
+        images = self.resize(images)
+        pixel_values = self.normalize(images / 255.0)
+        return pixel_values

README.md

@@ -1,13 +1,33 @@
 ---
 title: Lip Forcing
-emoji: 🏢
-colorFrom: gray
+emoji: 🗣️
+colorFrom: indigo
 colorTo: green
 sdk: gradio
-sdk_version: 6.20.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+python_version: "3.10"
+short_description: Few-step autoregressive diffusion for real-time lip sync
+startup_duration_timeout: 60m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Lip Forcing
+
+Gradio demo for **Lip Forcing: Few-Step Autoregressive Diffusion for Real-time
+Lip Synchronization** (KAIST AI · AIPARK), running the released self-contained
+**14B student** checkpoint.
+
+Given a talking-head reference video and a driving audio clip, the model detects
+and aligns the face to 512×512, then regenerates the mouth region to match the
+audio using a causal 2-step autoregressive diffusion student, and composites the
+result back into the original frames.
+
+- Paper: https://arxiv.org/abs/2606.11180
+- Project page: https://cvlab-kaist.github.io/LipForcing/
+- Code: https://github.com/cvlab-kaist/LipForcing
+- Weights: https://huggingface.co/JinhyukJang/lipforcing
+
+This demo reproduces the official streaming inference pipeline
+(`scripts/inference/inference_streaming.py`) 1:1 on ZeroGPU. Driving audio is
+capped to the first few seconds per request to keep a single call within the GPU
+time budget.

app.py

@@ -0,0 +1,373 @@
+"""Lip Forcing — few-step autoregressive diffusion for real-time lip synchronization.
+
+ZeroGPU Gradio demo for the released 14B student
+(https://huggingface.co/JinhyukJang/lipforcing). Given a talking-head reference
+video and a driving audio clip, it re-synchronizes the mouth to the audio using
+the streaming per-chunk AR pipeline from the official repo
+(scripts/inference/inference_streaming.py), reproduced 1:1 here.
+"""
+
+import os
+
+# Allocator: the streaming AR loop has transient spikes (VAE encode/decode of
+# 512x512 chunks + KV cache). expandable segments avoids fragmentation OOMs.
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("ORT_DISABLE_THREAD_AFFINITY", "1")
+os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+import spaces  # noqa: E402  — must precede torch / CUDA-touching imports
+
+import sys
+import types
+import tempfile
+import traceback
+
+import numpy as np
+import torch
+import gradio as gr
+from PIL import Image
+from huggingface_hub import hf_hub_download, snapshot_download
+
+# The inference scripts import their helpers as top-level modules
+# (`from _common import ...`), so make scripts/inference importable that way.
+REPO_ROOT = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, REPO_ROOT)
+sys.path.insert(0, os.path.join(REPO_ROOT, "scripts", "inference"))
+
+# ---------------------------------------------------------------------------
+# Weights (downloaded once at startup into the HF cache)
+# ---------------------------------------------------------------------------
+print("Downloading weights ...", flush=True)
+
+CKPT_PATH = hf_hub_download("JinhyukJang/lipforcing", "lipforcing_14b.pth")
+
+WAN_REPO = "Wan-AI/Wan2.1-T2V-14B"
+VAE_PATH = hf_hub_download(WAN_REPO, "Wan2.1_VAE.pth")
+T5_PATH = hf_hub_download(WAN_REPO, "models_t5_umt5-xxl-enc-bf16.pth")
+# UMT5 tokenizer (lives under google/umt5-xxl/ inside the Wan repo)
+for _f in (
+    "google/umt5-xxl/special_tokens_map.json",
+    "google/umt5-xxl/spiece.model",
+    "google/umt5-xxl/tokenizer.json",
+    "google/umt5-xxl/tokenizer_config.json",
+):
+    hf_hub_download(WAN_REPO, _f)
+# T5_PATH's parent dir now also holds google/umt5-xxl/* (same snapshot dir).
+
+WAV2VEC_DIR = snapshot_download("facebook/wav2vec2-base-960h")
+
+# TAEW tiny streaming decoder + LatentSync mouth mask.
+# taew2_1.pth lives in the taehv GitHub repo; mask.png in the LatentSync repo.
+import urllib.request  # noqa: E402
+_cache = os.path.join(tempfile.gettempdir(), "lipforcing_assets")
+os.makedirs(_cache, exist_ok=True)
+
+TAEHV_CKPT = os.path.join(_cache, "taew2_1.pth")
+if not os.path.exists(TAEHV_CKPT):
+    urllib.request.urlretrieve(
+        "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth",
+        TAEHV_CKPT,
+    )
+
+MASK_PATH = os.path.join(_cache, "mask.png")
+if not os.path.exists(MASK_PATH):
+    urllib.request.urlretrieve(
+        "https://raw.githubusercontent.com/bytedance/LatentSync/main/latentsync/utils/mask.png",
+        MASK_PATH,
+    )
+
+print("Weights downloaded.", flush=True)
+
+DTYPE = torch.bfloat16
+DEVICE = "cuda"
+
+# ---------------------------------------------------------------------------
+# Args shim — the loaders/helpers read attributes off an argparse-like object.
+# We build one with the released 14B student's default (2-step t769) schedule.
+# ---------------------------------------------------------------------------
+def _make_args():
+    a = types.SimpleNamespace()
+    a.ckpt_path = CKPT_PATH
+    a.vae_path = VAE_PATH
+    a.wav2vec_path = WAV2VEC_DIR
+    a.mask_path = MASK_PATH
+    a.taehv_ckpt = TAEHV_CKPT
+    a.base_model_paths = None
+    a.omniavatar_ckpt_path = None
+    a.model_size = "14B"
+    a.merge_lora_post_load = True
+    a.text_embeds_path = None
+    a.text_encoder_path = None           # text encoded once at startup (below)
+    a.prompt = "a person talking"
+    a.streaming_decoder = "streaming_taehv"
+    a.t_list = [0.999, 0.769, 0.0]        # released 14B 2-step schedule
+    a.chunk_size = 3
+    a.num_latent_frames = None
+    a.min_latent_frames = 0
+    a.context_noise = 0.0
+    a.seed = 42
+    a.fps = 25.0
+    a.dtype = "bf16"
+    a.device = DEVICE
+    a.local_attn_size = 7
+    a.sink_size = 1
+    a.use_dynamic_rope = True
+    a.skip_preprocessing = False
+    a.face_cache_dir = None
+    a.composite_full_face = False
+    a.streamwise_encode = True
+    a.defer_composite = False
+    a.compile = False
+    a.input_dir = None
+    a.output_dir = None
+    a.video_path = None
+    a.audio_path = None
+    a.output_path = None
+    return a
+
+
+ARGS = _make_args()
+
+# ---------------------------------------------------------------------------
+# Text embedding: encode the default prompt ONCE on CPU, then free the 11 GB
+# UMT5-XXL encoder. This keeps the encoder off the GPU so peak VRAM stays low
+# (~37 GB), per the model card's "48 GB cards work with precomputed embeddings".
+# ---------------------------------------------------------------------------
+def _precompute_text_embeds(prompt: str) -> torch.Tensor:
+    from OmniAvatar.models.wan_video_text_encoder import WanTextEncoder
+    from OmniAvatar.prompters.wan_prompter import WanPrompter
+    from lipforcing import preprocess as pp
+
+    print(f"Encoding text prompt on CPU: {prompt!r} ...", flush=True)
+    text_encoder = WanTextEncoder()
+    te_state = torch.load(T5_PATH, map_location="cpu", weights_only=False)
+    converter = WanTextEncoder.state_dict_converter()
+    te_state = converter.from_civitai(te_state)
+    text_encoder.load_state_dict(te_state, strict=True)
+    text_encoder = text_encoder.to("cpu").eval()
+
+    tokenizer_path = pp._resolve_tokenizer_path(T5_PATH)
+    prompter = WanPrompter(tokenizer_path=tokenizer_path, text_len=512)
+    prompter.fetch_models(text_encoder=text_encoder)
+    with torch.no_grad():
+        emb = prompter.encode_prompt(prompt, positive=True, device="cpu")
+    if emb.dim() == 2:
+        emb = emb.unsqueeze(0)
+    emb = emb.to(dtype=DTYPE).contiguous()
+    del text_encoder, prompter, te_state
+    import gc
+    gc.collect()
+    print(f"Text embeds: {tuple(emb.shape)}", flush=True)
+    return emb
+
+
+TEXT_EMBEDS_CPU = _precompute_text_embeds(ARGS.prompt)
+
+# ---------------------------------------------------------------------------
+# Models — loaded at module scope, .to("cuda") intercepted by ZeroGPU.
+# ---------------------------------------------------------------------------
+print("Loading diffusion model (14B student) ...", flush=True)
+from _loader import load_diffusion_model            # noqa: E402
+from _common import (                               # noqa: E402
+    TAEHVDecoderWrapper, load_vae, load_wav2vec,
+    resolve_audio, compute_generation_length,
+    load_image_processor, preprocess_with_latentsync,
+    build_condition_streamwise,
+)
+from inference_streaming import run_streaming_pipeline  # noqa: E402
+
+MODEL = load_diffusion_model(ARGS, DEVICE, DTYPE)
+
+print("Loading Wan VAE ...", flush=True)
+VAE = load_vae(ARGS.vae_path, DEVICE)
+
+print("Loading TAEHV decoder ...", flush=True)
+DECODER_VAE = TAEHVDecoderWrapper(ARGS.taehv_ckpt, DEVICE)
+
+print("Loading Wav2Vec2 ...", flush=True)
+WAV2VEC_MODEL, WAV2VEC_EXTRACTOR = load_wav2vec(ARGS.wav2vec_path, DEVICE)
+
+# LatentSync face detector / aligner uses insightface + onnxruntime; those need
+# a live GPU context, so it is initialized lazily inside the GPU call.
+IMAGE_PROCESSOR = None
+
+
+def _get_image_processor():
+    global IMAGE_PROCESSOR
+    if IMAGE_PROCESSOR is None:
+        IMAGE_PROCESSOR = load_image_processor(ARGS.mask_path, DEVICE)
+    return IMAGE_PROCESSOR
+
+
+# ---------------------------------------------------------------------------
+# Inference
+# ---------------------------------------------------------------------------
+MAX_SECONDS = 8.0  # cap driving audio so a single call stays within GPU budget
+
+
+def _estimate_duration(video_path, audio_path, *a, **k):
+    # 14B student: streaming AR + face detect/composite. Budget generously per
+    # second of (capped) audio, plus fixed preprocessing/warmup overhead.
+    return 300
+
+
+@spaces.GPU(duration=_estimate_duration)
+def lip_sync(video_path: str, audio_path: str,
+             seed: int = 42) -> str:
+    """Lip-sync a talking-head video to a driving audio clip.
+
+    Args:
+        video_path: reference talking-head video (any resolution; a single
+            clear front-facing face is detected, aligned to 512x512, and the
+            mouth region is regenerated to match the audio).
+        audio_path: driving speech audio; the output length follows the audio
+            (capped to keep a single request within the GPU budget).
+        seed: RNG seed for reproducibility.
+
+    Returns:
+        Path to the generated lip-synced mp4 (muxed with the driving audio).
+    """
+    if not video_path:
+        raise gr.Error("Please provide a reference talking-head video.")
+    if not audio_path:
+        raise gr.Error("Please provide a driving audio clip.")
+
+    import imageio_ffmpeg
+    import subprocess
+
+    args = _make_args()
+    args.seed = int(seed)
+    args.video_path = video_path
+    args.audio_path = audio_path
+
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+
+    image_processor = _get_image_processor()
+
+    # Cap audio length so runtime stays bounded.
+    ff = imageio_ffmpeg.get_ffmpeg_exe()
+    capped_audio = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
+    subprocess.run(
+        [ff, "-y", "-loglevel", "error", "-nostdin", "-i", audio_path,
+         "-t", str(MAX_SECONDS), "-ar", "16000", "-ac", "1", capped_audio],
+        check=True,
+    )
+
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+
+    tmp_audio = None
+    try:
+        used_audio, tmp_audio = resolve_audio(audio_path=capped_audio)
+
+        num_latent_frames, num_video_frames = compute_generation_length(
+            used_audio, args.num_latent_frames, args.chunk_size, args.fps,
+            min_latent_frames=args.min_latent_frames,
+        )
+
+        print("Face detection + 512x512 alignment ...", flush=True)
+        meta = preprocess_with_latentsync(
+            args.video_path, image_processor, args.face_cache_dir,
+            num_frames=num_video_frames,
+        )
+        if meta is None:
+            raise gr.Error(
+                "Face detection failed — please provide a video with a single, "
+                "clear, front-facing talking head."
+            )
+
+        aligned_faces = meta["aligned_faces"]
+        ref_frames_np = np.stack([
+            f.permute(1, 2, 0).numpy() if isinstance(f, torch.Tensor) else f
+            for f in aligned_faces[:num_video_frames]
+        ], axis=0)
+
+        text_embeds = TEXT_EMBEDS_CPU.to(device=DEVICE, dtype=DTYPE)
+
+        condition, video_tensor, masked_video_tensor = build_condition_streamwise(
+            VAE, WAV2VEC_MODEL, WAV2VEC_EXTRACTOR,
+            ref_frames_np, used_audio, text_embeds, args.mask_path,
+            num_video_frames, num_latent_frames, DEVICE, DTYPE,
+        )
+
+        print("Running streaming pipeline ...", flush=True)
+        run_streaming_pipeline(
+            MODEL, DECODER_VAE, VAE, condition,
+            num_latent_frames, num_video_frames,
+            args, meta, image_processor,
+            used_audio, out_path, DEVICE, DTYPE,
+            video_tensor=video_tensor,
+            masked_video_tensor=masked_video_tensor,
+        )
+    except gr.Error:
+        raise
+    except Exception as e:
+        traceback.print_exc()
+        raise gr.Error(f"Inference failed: {e}")
+    finally:
+        MODEL.clear_caches()
+        torch.cuda.empty_cache()
+        if tmp_audio and os.path.exists(tmp_audio):
+            os.remove(tmp_audio)
+        if os.path.exists(capped_audio):
+            os.remove(capped_audio)
+
+    return out_path
+
+
+# ---------------------------------------------------------------------------
+# UI
+# ---------------------------------------------------------------------------
+CSS = """
+#col-container { max-width: 1100px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+DESCRIPTION = """
+# Lip Forcing 🗣️
+**Few-Step Autoregressive Diffusion for Real-time Lip Synchronization**  · 
+14B student  · 
+[Paper](https://arxiv.org/abs/2606.11180)  · 
+[Project](https://cvlab-kaist.github.io/LipForcing/)  · 
+[Code](https://github.com/cvlab-kaist/LipForcing)  · 
+[Weights](https://huggingface.co/JinhyukJang/lipforcing)
+
+Give it a **talking-head video** and a **driving audio** clip — it detects and aligns
+the face, then regenerates the mouth to match the audio with a 2-step causal diffusion
+student. Audio is capped to the first few seconds per run.
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(DESCRIPTION)
+        with gr.Row():
+            with gr.Column():
+                video_in = gr.Video(label="Reference talking-head video", height=340)
+                audio_in = gr.Audio(label="Driving audio", type="filepath")
+                run_btn = gr.Button("Lip-sync", variant="primary")
+            with gr.Column():
+                video_out = gr.Video(label="Lip-synced result", height=340)
+        with gr.Accordion("Advanced settings", open=False):
+            seed = gr.Number(label="Seed", value=42, precision=0)
+
+        run_btn.click(
+            fn=lip_sync,
+            inputs=[video_in, audio_in, seed],
+            outputs=video_out,
+            api_name="lip_sync",
+        )
+
+        gr.Examples(
+            examples=[
+                ["examples/example1_video.mp4", "examples/example1_audio.wav"],
+                ["examples/example2_video.mp4", "examples/example2_audio.wav"],
+            ],
+            inputs=[video_in, audio_in],
+            outputs=video_out,
+            fn=lip_sync,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

examples/example1_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe2f4abfdfa42bb474880cd265071d10c5c743f2dbd4f48abbc06a361ede731f
+size 128078

examples/example1_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bbdab9fe0231c662b44897146e52a4a604db27f3f55e81a8c1e3a2a664989a2
+size 896877

examples/example2_audio.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dfb6202c7411ac57d16084ac2722c8fb88e1976eef544f169eb6e6e9c5286aec
+size 128078

examples/example2_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13ac3d03a4f9f2a81f2d337f6f07283f58f5f19ecb8901dc786a5575f36fd97a
+size 634917

lipforcing/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

lipforcing/callbacks/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

lipforcing/callbacks/callback.py

@@ -0,0 +1,183 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+from typing import Callable, Any, TYPE_CHECKING
+import torch
+from torch.utils.data import DataLoader
+
+from lipforcing.utils import instantiate
+import lipforcing.utils.logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.configs.config import BaseConfig
+    from lipforcing.trainer import Trainer
+    from lipforcing.methods import FastGenModel
+
+
+class CallbackDict:
+    def __init__(self, config: BaseConfig, trainer: Trainer):
+        self._callbacks = {}
+        callback_configs = config.trainer.callbacks
+        if callback_configs:
+            if isinstance(callback_configs, list):
+                logger.warning(msg="The 'config.trainer.callbacks' parameter should be a dict instead of a list. ")
+                callback_configs = {f"callback_{k}": v for k, v in enumerate(callback_configs)}
+            for callback_name, current_callback_cfg in callback_configs.items():
+                if "_target_" not in current_callback_cfg:
+                    logger.critical(
+                        f"Callback {callback_name} is missing the '_target_' field. \n Skip {current_callback_cfg}"
+                    )
+                    continue
+                logger.critical(f"Instantiating callback {callback_name}: {current_callback_cfg}")
+                _callback = instantiate(current_callback_cfg)
+                assert isinstance(_callback, Callback), f"{current_callback_cfg} is not a valid callback."
+                _callback.config = config
+                _callback.trainer = trainer
+                _callback.on_app_begin()
+                self._callbacks[callback_name] = _callback
+
+    def __getattr__(self, method_name: str) -> Callable:
+        def load_state_dict(state_dict: dict[str, Any]) -> None:
+            for name, callback in self._callbacks.items():
+                if name in state_dict:
+                    callback.load_state_dict(state_dict[name])
+                else:
+                    logger.warning(f"Callback {name} not found in checkpoint.")
+
+        def state_dict() -> dict[str, Any]:
+            return {name: self._callbacks[name].state_dict() for name in self._callbacks}
+
+        def callbacks_wrapper(*args, **kwargs):
+            for callback in self._callbacks.values():
+                assert hasattr(callback, method_name)
+                method = getattr(callback, method_name)
+                assert callable(method), f"{method_name} is not callable."
+                method(*args, **kwargs)
+
+        if method_name == "state_dict":
+            return state_dict
+        if method_name == "load_state_dict":
+            return load_state_dict
+        return callbacks_wrapper
+
+
+class Callback:
+    config: "BaseConfig"
+    trainer: "Trainer"
+
+    def on_app_begin(self) -> None:
+        pass
+
+    def on_model_init_start(self, model: FastGenModel) -> None:
+        pass
+
+    def on_model_init_end(self, model: FastGenModel | torch.nn.parallel.DistributedDataParallel) -> None:
+        pass
+
+    def on_optimizer_init_start(self, model: FastGenModel) -> None:
+        pass
+
+    def on_optimizer_init_end(self, model: FastGenModel) -> None:
+        pass
+
+    def on_load_checkpoint_start(self, model: FastGenModel) -> None:
+        pass
+
+    def on_load_checkpoint_end(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_dataloader_init_start(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_dataloader_init_end(
+        self, model: FastGenModel, dataloader_train: DataLoader, dataloader_val: DataLoader, iteration: int = 0
+    ) -> None:
+        pass
+
+    def on_train_begin(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_training_step_begin(
+        self,
+        model: FastGenModel,
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_training_accum_step_begin(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        iteration: int = 0,
+        accum_iter: int = 0,
+    ) -> None:
+        pass
+
+    def on_backward_begin(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+        accum_iter: int = 0,
+    ) -> None:
+        pass
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_optimizer_step_begin(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_train_end(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_validation_begin(self, model: FastGenModel, iteration: int = 0, idx: int = 0) -> None:
+        pass
+
+    def on_validation_step_begin(
+        self, model: FastGenModel, data_batch: dict[str, torch.Tensor], step: int = 0, iteration: int = 0, idx: int = 0
+    ) -> None:
+        pass
+
+    def on_validation_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        step: int = 0,
+        iteration: int = 0,
+        idx: int = 0,
+    ) -> None:
+        pass
+
+    def on_validation_end(self, model: FastGenModel, iteration: int = 0, idx: int = 0) -> None:
+        pass
+
+    def on_save_checkpoint_start(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_save_checkpoint_success(self, model: FastGenModel, iteration: int = 0, path: str = None) -> None:
+        pass
+
+    def on_save_checkpoint_end(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def on_app_end(self, model: FastGenModel, iteration: int = 0) -> None:
+        pass
+
+    def state_dict(self) -> dict[str, Any]:
+        return {}
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        pass

lipforcing/callbacks/ct_schedule.py

@@ -0,0 +1,83 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+from typing import Callable, TYPE_CHECKING
+import wandb
+
+import torch
+
+from lipforcing.callbacks.callback import Callback
+import lipforcing.utils.logging_utils as logger
+from lipforcing.utils.basic_utils import get_batch_size_total
+from lipforcing.utils.distributed import is_rank0
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+    from lipforcing.configs.config import BaseConfig
+
+
+class CTScheduleCallback(Callback):
+    config: "BaseConfig"
+
+    def __init__(
+        self,
+        q: float = 2.0,
+        ratio_limit: float = 0.999,
+        kimg_per_stage: int = 12500,
+        batch_size: int = 1,
+    ):
+        self.q = q
+        self.ratio_limit = ratio_limit
+        self.kimg_per_stage = kimg_per_stage
+        self.batch_size = batch_size
+
+        self.stage = 0
+        self.ratio = 0.0
+
+    def _get_cur_stage(self, model, iteration):
+        # Start from the saved iteration of the first-stage model in TCM
+        if hasattr(model, "resume_iter"):
+            assert isinstance(model.resume_iter, int)
+            iteration = iteration + model.resume_iter
+
+        batch_size = self.batch_size
+        if hasattr(self, "config"):
+            # override the batch_size using self.config
+            batch_size = get_batch_size_total(self.config)
+
+        cur_nimg = iteration * batch_size
+        stage = cur_nimg // (self.kimg_per_stage * 1000)
+        return stage, cur_nimg
+
+    def _update_schedule(self, stage):
+        self.stage = stage
+        self.ratio = 1 - 1 / self.q ** (stage + 1)
+        if self.ratio > self.ratio_limit:
+            logger.info(f"Clipping ratio from {self.ratio} -> {self.ratio_limit}")
+            self.ratio = self.ratio_limit
+
+    def on_train_begin(self, model: FastGenModel, iteration: int = 0) -> None:
+        stage, _ = self._get_cur_stage(model, iteration)
+        self._update_schedule(stage)
+        setattr(model, "ratio", self.ratio)
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        del data_batch, output_batch, loss_dict
+        new_stage, cur_nimg = self._get_cur_stage(model, iteration)
+        if new_stage > self.stage:
+            self._update_schedule(new_stage)
+            setattr(model, "ratio", self.ratio)
+
+        if hasattr(self, "config"):
+            # only wandb log when config exists
+            if iteration % self.config.trainer.logging_iter == 0 and is_rank0():
+                if wandb.run:
+                    wandb.log({"ct_schedule/kimg": cur_nimg / 1e3, "ct_schedule/ratio": self.ratio}, step=iteration)

lipforcing/callbacks/ema.py

@@ -0,0 +1,169 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Callable, TYPE_CHECKING, Optional
+
+import torch
+import wandb
+
+from lipforcing.callbacks.callback import Callback
+from lipforcing.utils.basic_utils import get_batch_size_total
+from lipforcing.utils.distributed import synchronize, is_rank0
+import lipforcing.utils.logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+class EMACallback(Callback):
+    def __init__(
+        self,
+        type: str = "constant",
+        # params for type=constant
+        beta: float = 0.9999,
+        # params for type=power
+        gamma: float = 16.97,
+        # params for type=halflife
+        ema_halflife_kimg: float = 500,
+        ema_rampup_ratio: Optional[float] = 0.05,
+        start_iter: int = 0,
+        ema_name: str = "ema",
+        batch_size: int = 1,  # overwritten by self.config if it exists
+        fsdp: bool = False,  # overwritten by self.config if it exists
+    ):
+        self.type = type
+        self.beta = beta
+        self.gamma = gamma
+        self.ema_halflife_kimg = ema_halflife_kimg
+        self.ema_rampup_ratio = ema_rampup_ratio
+        self.start_iter = start_iter
+        self.ema_name = ema_name
+        self.batch_size = batch_size
+        self._is_fsdp = fsdp
+        self._enabled = True
+
+    def on_app_begin(self) -> None:
+        if hasattr(self, "config"):
+            # override using config
+            self._is_fsdp = self.config.trainer.fsdp
+            self.batch_size = get_batch_size_total(self.config)
+
+    def on_model_init_end(
+        self, model: FastGenModel | torch.nn.parallel.DistributedDataParallel, iteration: int = 0
+    ) -> None:
+        # Unwrap DDP if needed to access the original model's attributes
+        if hasattr(model, "module"):
+            model = model.module
+
+        # check ema initialization
+        ema = getattr(model, self.ema_name, None)
+        if ema is None:
+            self._enabled = False
+            logger.info(f"EMA {self.ema_name} is not enabled, skipping callback.")
+            return
+
+        assert ema.training is False, f"EMA {self.ema_name} should be in eval mode"
+        for name, p_net in ema.named_parameters():
+            assert not p_net.requires_grad, f"EMA parameter {name} should not require gradients"
+
+    def _total_iteration(self, model: FastGenModel, iteration: int) -> int:
+        if hasattr(model, "resume_iter"):
+            assert isinstance(model.resume_iter, int)
+            iteration = iteration + model.resume_iter
+        return iteration
+
+    def _power_function_beta(self, iteration):
+        beta = (1 - 1 / iteration) ** (self.gamma + 1)
+        return beta
+
+    def _get_cur_nimg(self, iteration):
+        cur_nimg = iteration * self.batch_size
+        return self.batch_size, cur_nimg
+
+    def _halflife_beta(self, iteration):
+        ema_halflife_nimg = self.ema_halflife_kimg * 1000
+        batch_size, cur_nimg = self._get_cur_nimg(iteration)
+        if self.ema_rampup_ratio is not None:
+            ema_halflife_nimg = min(ema_halflife_nimg, cur_nimg * self.ema_rampup_ratio)
+        ema_beta = 0.5 ** (batch_size / max(ema_halflife_nimg, 1e-8))
+        return ema_beta
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        del data_batch, output_batch, loss_dict
+
+        # Check if EMA is enabled
+        if not self._enabled:
+            return
+
+        # Get total iteration and skip if before start_iter
+        total_iteration = self._total_iteration(model, iteration)
+        if total_iteration < self.start_iter:
+            return
+        elif total_iteration == self.start_iter:
+            logger.info(f"Starting to update {self.ema_name} at iteration {total_iteration}.")
+
+        if self.type == "constant":
+            beta = self.beta
+        elif self.type == "power":
+            beta = self._power_function_beta(total_iteration)
+        elif self.type == "halflife":
+            beta = self._halflife_beta(total_iteration)
+        else:
+            raise ValueError(f"Invalid {self.ema_name} type: {self.type}")
+
+        with torch.no_grad():
+            ema = getattr(model, self.ema_name)
+            ema_state_dict = ema.state_dict()
+
+            for name, p_net in model.net.named_parameters():
+                if self._is_fsdp and hasattr(p_net, "full_tensor"):
+                    # Gather the full tensor from all ranks if using FSDP with DTensor
+                    # When CPU offloading is enabled, we need to move to CUDA first because
+                    # full_tensor() performs an all_gather which requires a CUDA backend
+                    if p_net.device.type == "cpu":
+                        # Move local shard to CUDA, gather, then the result stays on CUDA
+                        # which is fine since we'll copy to EMA (which handles device placement)
+                        full_tensor = p_net.to("cuda").full_tensor()
+                    else:
+                        full_tensor = p_net.full_tensor()
+                else:
+                    full_tensor = p_net
+                # Strip checkpoint wrapper prefix if present (EMA doesn't have checkpointing)
+                ema_name = name.replace("_checkpoint_wrapped_module.", "")
+                # Cast to EMA dtype and device for lerp_ compatibility
+                if ema_name in ema_state_dict:
+                    ema_param = ema_state_dict[ema_name]
+                    if total_iteration == self.start_iter:
+                        # re-initialize EMA parameter
+                        ema_param.copy_(full_tensor.to(device=ema_param.device, dtype=ema_param.dtype))
+                    else:
+                        # interpolate EMA parameter
+                        ema_param.lerp_(full_tensor.to(device=ema_param.device, dtype=ema_param.dtype), 1.0 - beta)
+                elif iteration == 1:
+                    # only warn on first iteration if parameter is not found
+                    logger.warning(f"EMA parameter {ema_name} not found in EMA state dict, skipping update.")
+
+            # FSDP2 doesn't shard buffers, so we can just copy them
+            for name, p_net in model.net.named_buffers():
+                if name in ema_state_dict:
+                    ema_param = ema_state_dict[name]
+                    ema_param.copy_(p_net.to(device=ema_param.device, dtype=ema_param.dtype))
+                elif iteration == 1:
+                    # only warn on first iteration if buffer is not found
+                    logger.warning(f"EMA buffer {name} not found in EMA state dict, skipping update.")
+
+            if hasattr(self, "config"):
+                # only wandb log when config exists
+                if iteration % self.config.trainer.logging_iter == 0 and is_rank0():
+                    if wandb.run:
+                        wandb.log({f"ema/{self.ema_name}_beta": beta}, step=iteration)
+        synchronize()

lipforcing/callbacks/forced_weight_norm.py

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+
+from lipforcing.callbacks.callback import Callback
+import lipforcing.utils.logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+class ForcedWeightNormCallback(Callback):
+    def on_training_accum_step_begin(
+        self,
+        model: FastGenModel,
+        *args,
+        **kwargs,
+    ) -> None:
+        if hasattr(model.net, "forced_weight_normalization"):
+            model.net.forced_weight_normalization()
+        else:
+            logger.warning(
+                "Enabled ForcedWeightNormCallback but model.net does not have the forced_weight_normalization method."
+            )

lipforcing/callbacks/gpu_mem_profiler.py

@@ -0,0 +1,134 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+import os
+from lipforcing.utils import logging_utils as logger
+from lipforcing.callbacks.callback import Callback
+import atexit
+import pickle
+from typing import Callable, Optional, TYPE_CHECKING
+import base64
+import json
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+def create_dump(dump_path):
+    logger.critical(f"Creating {dump_path}")
+    if not dump_path.endswith("html"):
+        print(f"[{__file__}] create_dump produces an HTML file but was called with {dump_path=}")
+    torch.cuda.memory._dump_snapshot(dump_path + ".pickle")
+    with open(dump_path + ".pickle", "rb") as f:
+        data = pickle.load(f)
+    _memory_viz_template = r"""
+    <!DOCTYPE html>
+    <html>
+    <head>
+    </head>
+    <body>
+    <script type="module">
+    import {add_local_files} from "https://cdn.jsdelivr.net/gh/pytorch/pytorch@main/torch/utils/viz/MemoryViz.js"
+    const local_files = $SNAPSHOT
+    add_local_files(local_files, $VIZ_KIND)
+    </script>
+    </body>
+    """
+
+    # find which GPU was active
+    idx_device = -1
+    for i in range(8):
+        if data["device_traces"][i]:
+            idx_device = i
+            break
+
+    traces = data["device_traces"][idx_device]  # create an aliasing variable for convenience
+    traces = [
+        d for d in traces if d["action"] == "alloc" or d["action"] == "free_completed"
+    ]  # only the `alloc` and `free_completed` events matter for our visualization
+
+    for d in traces:
+        d["fastgen_frames"] = [
+            f for f in d["frames"] if "lipforcing" in f["filename"]
+        ]  # get the callstack frames from lipforcing code (e.g. ignore frames in pytorch/other libraries)
+        if not d["fastgen_frames"]:
+            d["fastgen_frames"] = d["frames"]
+
+    # run through the trace and find allocations that were allocated but never freed
+    set_alloced_addrs: dict = {}
+    for d in traces:
+        if d["action"] == "alloc":
+            set_alloced_addrs[d["addr"]] = d
+        elif d["action"] == "free_completed":
+            if d["addr"] in set_alloced_addrs:
+                del set_alloced_addrs[d["addr"]]
+        else:
+            raise NotImplementedError(f"{d['action']}")
+
+    never_freed_traces = list(set_alloced_addrs.values())
+    KB = 1 << 10
+    never_freed_traces = [t for t in never_freed_traces if t["size"] > KB]  # get rid of allocations below 1 KB
+
+    # now proceed through the trace (guarenteed to be all `alloc` events as we removed all free events).
+    # for each pair of alloc events, merge them iff they share a common lipforcing ancestor.
+    # Merging events is useful as it both speeds up the visualization rendering and also makes it more understandable.
+    i = 0
+    while i < len(never_freed_traces) - 1:
+        curr_frames = never_freed_traces[i]["fastgen_frames"]
+        next_frames = never_freed_traces[i + 1]["fastgen_frames"]
+        if (
+            curr_frames and next_frames and curr_frames[0] == next_frames[0]
+        ):  # note: compares only the innermost frame, not the full callstack
+            # same ancestor, delete next event and add its size to current event
+            never_freed_traces[i]["size"] += never_freed_traces[i + 1]["size"]
+            never_freed_traces.pop(i + 1)
+        else:
+            i += 1  # different ancestor, do not combine and move on
+
+    data["device_traces"][idx_device] = never_freed_traces  # update the trace to only be the merged-alloc events
+    data["segments"] = []  # shrink the trace, unused in memory timeline
+    data["external_annotations"] = []  # shrink the trace, unused in memory timeline
+    buffer = pickle.dumps(data)
+    buffer += b"\x00" * (3 - len(buffer) % 3)
+    encoded_buffer = base64.b64encode(buffer).decode("utf-8")
+    json_format = json.dumps([{"name": "snapshot.pickle", "base64": encoded_buffer}])
+    html_src = _memory_viz_template.replace("$VIZ_KIND", repr("Active Memory Timeline")).replace(
+        "$SNAPSHOT", json_format
+    )
+    with open(dump_path, "w") as f:
+        f.write(html_src)
+
+
+class MemTrackerCallback(Callback):
+    def __init__(self, save_every_n_iters: Optional[int] = None, deactivate_after_n_iters: int = 100):
+        def close_and_save():
+            create_dump(
+                f"{os.environ.get('LIPFORCING_OUTPUT_ROOT', 'LIPFORCING_OUTPUT')}/crash_rank{os.environ.get('RANK', '0')}.html"
+            )
+
+        self.deactivate_after_n_iters = deactivate_after_n_iters  # Deactivate eventually to prevent leaking host memory
+        self.save_every_n_iters = save_every_n_iters
+        self.atexit_fn = close_and_save
+        atexit.register(self.atexit_fn)
+
+    def on_app_begin(self):
+        logger.info("[MemTrackerCallback] Tracking peak memory usage")
+        torch.cuda.memory._record_memory_history(stacks="python")
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        if iteration > self.deactivate_after_n_iters:
+            torch.cuda.memory._record_memory_history(enabled=None)  # frees pytorch tracking datastructures
+        if self.save_every_n_iters is not None and (iteration % self.save_every_n_iters) == 0:
+            create_dump(
+                f"{os.environ.get('LIPFORCING_OUTPUT_ROOT', 'LIPFORCING_OUTPUT')}/step{iteration}_rank{os.environ.get('RANK', '0')}.html"
+            )

lipforcing/callbacks/gpu_stats.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING, Callable, Any, Dict, List
+
+import pandas as pd
+import psutil
+import torch
+
+from lipforcing.callbacks.callback import Callback
+from lipforcing.utils.distributed import world_size, is_rank0, synchronize
+import lipforcing.utils.logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+def log_prof_data(data_list: List[Dict[str, Any]]):
+    # Create a table to log data with rank information
+    metrics = list(data_list[0].keys())
+
+    # Initialize dictionaries to store min and max values for each metric
+    min_values = {key: float("inf") for key in metrics}
+    max_values = {key: float("-inf") for key in metrics}
+    sum_values = {key: 0.0 for key in metrics}
+
+    count = 0
+
+    for _rank, prof_data in enumerate(data_list):
+        count += 1
+
+        # Update min, max, and sum values
+        for key in metrics:
+            min_values[key] = min(min_values[key], prof_data[key])
+            max_values[key] = max(max_values[key], prof_data[key])
+            sum_values[key] += prof_data[key]
+
+    # Calculate average values
+    avg_values = {key: sum_values[key] / count for key in metrics}
+    summary_df = pd.DataFrame({"Avg": avg_values, "Max": max_values, "Min": min_values})
+
+    logger.info(f"GPU stats:\n{summary_df.to_string()}")
+
+
+class GPUStatsCallback(Callback):
+    def __init__(self, every_n: int = 100):
+        self.every_n = every_n
+
+    def on_train_begin(self, model: FastGenModel, iteration: int = 0):
+        torch.cuda.reset_peak_memory_stats()
+        if hasattr(self, "config"):
+            # overwritten by logging_iter if self.config exists
+            self.every_n = self.config.trainer.logging_iter
+        logger.info(f"every_n to measure gpus stats: {self.every_n}")
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        del data_batch, output_batch, loss_dict
+        if iteration % self.every_n == 0:
+            cur_process = psutil.Process(os.getpid())
+            cpu_memory_usage = sum(p.memory_info().rss for p in [cur_process] + cur_process.children(recursive=True))
+            cpu_mem_gb = cpu_memory_usage / (1024**3)
+
+            peak_gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
+            peak_gpu_mem_reserved_gb = torch.cuda.max_memory_reserved() / (1024**3)
+            util = torch.cuda.utilization()
+
+            prof_data = {
+                "cpu_mem_gb": float(cpu_mem_gb),
+                "peak_gpu_mem_gb": float(peak_gpu_mem_gb),
+                "peak_gpu_mem_reserved_gb": float(peak_gpu_mem_reserved_gb),
+                "util": float(util),
+            }
+
+            synchronize()
+            data_list = [prof_data] * world_size()
+            # this is blocking by default
+            if world_size() > 1:
+                torch.distributed.all_gather_object(data_list, prof_data)
+
+            if is_rank0():
+                log_prof_data(data_list)
+            synchronize()

lipforcing/callbacks/grad_clip.py

@@ -0,0 +1,219 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Optional
+import wandb
+
+import torch
+from torch.distributed.tensor import DTensor
+
+from lipforcing.callbacks.callback import Callback
+from lipforcing.utils.distributed import is_rank0, world_size
+import lipforcing.utils.logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+@contextmanager
+def cast_gradients_dtype(model, dtype=torch.float32, enabled=True):
+    if enabled:
+        try:
+            # Cast gradients to the desired dtype
+            for param in model.parameters():
+                if param.grad is not None and param.grad.dtype != dtype:
+                    param.grad.data = param.grad.data.to(dtype)
+            yield
+        finally:
+            # Restore original gradient dtypes
+            for param in model.parameters():
+                if param.grad is not None and param.grad.dtype != param.dtype:
+                    param.grad.data = param.grad.data.to(param.dtype)
+    else:
+        yield
+
+
+def clip_grad_norm_fsdp(
+    parameters,
+    max_norm: float,
+    norm_type: float = 2.0,
+    device: Optional[torch.device] = None,
+) -> torch.Tensor:
+    """
+    Clip gradients for FSDP2 models with CPU offloading.
+
+    The standard torch.nn.utils.clip_grad_norm_ fails with FSDP2 CPU offloading because
+    DTensor operations (like division) trigger all_reduce on CPU, which has no backend.
+
+    This implementation:
+    1. Extracts local tensors from DTensors
+    2. Computes local norms on native device (CPU or GPU)
+    3. All-reduces the scalar norm to get global norm
+    4. Clips gradients in-place using the global norm
+
+    Args:
+        parameters: Iterable of parameters with gradients
+        max_norm: Maximum norm value
+        norm_type: Type of norm (default: L2)
+        device: Device for all-reduce tensor. If None, inferred from gradients or defaults to cuda.
+
+    Returns:
+        Total gradient norm (global across all ranks) as a regular tensor
+    """
+    if isinstance(parameters, torch.Tensor):
+        parameters = [parameters]
+    parameters = list(p for p in parameters if p.grad is not None)
+
+    if len(parameters) == 0:
+        return torch.tensor(0.0)
+
+    # Compute per-parameter norms on their native device (CPU or GPU)
+    # We compute norm^norm_type to allow proper aggregation across ranks
+    local_norm_sum = 0.0
+    inferred_device = None
+    for p in parameters:
+        if isinstance(p.grad, DTensor):
+            grad = p.grad._local_tensor
+        else:
+            grad = p.grad
+
+        # Infer CUDA device from gradients (use first CUDA device found)
+        if inferred_device is None and grad.device.type == "cuda":
+            inferred_device = grad.device
+
+        # Compute norm on the gradient's native device, accumulate as Python float
+        local_norm_sum += torch.norm(grad.detach().float(), norm_type).item() ** norm_type
+
+    # Use provided device, or inferred device, or fall back to current CUDA device
+    if device is None:
+        device = inferred_device if inferred_device is not None else torch.device("cuda")
+    local_norm_sum = torch.tensor(local_norm_sum, device=device)
+
+    # All-reduce to get global norm across all ranks
+    if world_size() > 1:
+        torch.distributed.all_reduce(local_norm_sum, op=torch.distributed.ReduceOp.SUM)
+
+    # Compute final global norm
+    total_norm = local_norm_sum ** (1.0 / norm_type)
+
+    # Compute clip coefficient (regular tensor division, no DTensor ops)
+    clip_coef = max_norm / (total_norm + 1e-6)
+    clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
+
+    # Apply clipping to gradients
+    for p in parameters:
+        if isinstance(p.grad, DTensor):
+            # For DTensor, scale the local tensor directly
+            local_grad = p.grad._local_tensor
+            local_grad.mul_(clip_coef_clamped.to(local_grad.device))
+        else:
+            p.grad.detach().mul_(clip_coef_clamped.to(p.grad.device))
+
+    return total_norm
+
+
+class GradClipCallback(Callback):
+    def __init__(
+        self,
+        grad_norm: float | None = 1.0,
+        model_key: str = "net",
+        posinf: float | None = None,
+        neginf: float | None = None,
+        precision_grad_clip: Optional[torch.dtype] = None,
+    ) -> None:
+        self.grad_norm = grad_norm
+        self.model_key = model_key
+        self.posinf = posinf
+        self.neginf = neginf
+        self.precision_grad_clip = precision_grad_clip
+
+    def nan_to_num(self, module: torch.nn.Module) -> tuple[int, torch.dtype | None]:
+        grad_dtype = None
+        non_finite_grads_count = 0
+
+        for name, param in module.named_parameters():
+            if param.grad is not None:
+                grad_dtype = param.grad.dtype
+
+                # Extract local tensor for DTensor (avoids triggering distributed ops on CPU)
+                if isinstance(param.grad, DTensor):
+                    grad = param.grad._local_tensor
+                else:
+                    grad = param.grad
+
+                non_finite_grads = grad.numel() - grad.isfinite().sum().item()
+                if non_finite_grads:
+                    non_finite_grads_count += non_finite_grads
+                    logger.debug(
+                        f"Gradient of {name} (dtype {grad_dtype}) is not finite: "
+                        f"Setting {grad.isnan().sum().item()} NaNs to 0 and {grad.isinf().sum().item()} Infs "
+                        f"to {self.posinf} or {self.neginf}."
+                    )
+                    torch.nan_to_num(grad, nan=0.0, posinf=self.posinf, neginf=self.neginf, out=grad)
+
+        return non_finite_grads_count, grad_dtype
+
+    def on_optimizer_step_begin(self, model: FastGenModel, iteration: int = 0) -> None:
+        # unscale the optimizer related to the `model_key`
+        assert (
+            self.model_key in model.optimizer_dict.keys()
+        ), f"Keys in optimizer_dict: {list(model.optimizer_dict.keys())}."
+        optimizer = model.optimizer_dict[self.model_key]
+        # Only unscale if grad_scaler should be used (checks enabled + float32 grads)
+        if model.should_use_grad_scaler(optimizer):
+            model.grad_scaler.unscale_(optimizer)
+
+        # Save model device before selecting subnet (subnet may not have .device)
+        model_device = model.device
+
+        # select subnet if specified (by default, we only perform gradient clips on model.net)
+        subnets = self.model_key.split(".")
+        for subnet in subnets:
+            model = getattr(model, subnet)
+
+        # set nan to num for each parameter
+        non_finite_grads_count, grad_dtype = self.nan_to_num(model)
+        logger.debug(f"Gradient dtype of {self.model_key}: {grad_dtype}")
+        if non_finite_grads_count > 0:
+            logger.info(
+                f"Number of parameters with non-finite gradients (of dtype {grad_dtype}): {non_finite_grads_count}"
+            )
+        log_dict = {f"optimizer/non_finite_grads_count (model_key {self.model_key})": non_finite_grads_count}
+
+        if self.grad_norm is not None:
+            # Cast all gradients to precision_grad_clip for numerical stability during clipping
+            cast_grads = (
+                self.precision_grad_clip is not None
+                and grad_dtype is not None
+                and grad_dtype != self.precision_grad_clip
+            )
+
+            # log value at first iteration
+            if iteration == 1 and cast_grads:
+                logger.info(f"Casting gradients from {grad_dtype} to {self.precision_grad_clip} before clipping.")
+
+            # Check if CPU offloading is enabled by looking for DTensor grads on CPU
+            # CPU offloading = DTensor local tensors are on CPU
+            # Check if any gradients are DTensors (FSDP2 sharded params).
+            # Standard clip_grad_norm_ with foreach=True can't mix DTensor and regular Tensor.
+            has_dtensor_grads = any(
+                isinstance(p.grad, DTensor) for p in model.parameters() if p.grad is not None
+            )
+
+            with cast_gradients_dtype(model, dtype=self.precision_grad_clip, enabled=cast_grads):
+                if has_dtensor_grads:
+                    # Use custom clipping that handles DTensor/Tensor mix
+                    total_norm = clip_grad_norm_fsdp(model.parameters(), self.grad_norm, device=model_device)
+                else:
+                    # Standard clipping for non-FSDP
+                    total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), self.grad_norm, foreach=True)
+
+            log_dict[f"optimizer/grad_norm (model_key {self.model_key})"] = total_norm.item()
+
+        if hasattr(self, "config"):
+            # only wandb log when config exists
+            if iteration % self.config.trainer.logging_iter == 0 and is_rank0() and wandb.run:
+                wandb.log(log_dict, step=iteration)

lipforcing/callbacks/param_count.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
+from lipforcing.callbacks.callback import Callback
+from lipforcing.utils.distributed import world_size
+import lipforcing.utils.logging_utils as logger
+import torch
+import wandb
+
+try:
+    from torch.distributed.tensor import DTensor
+except ImportError:
+    DTensor = None
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+def _get_local_numel(param: torch.Tensor) -> int:
+    """Get the local (sharded) number of elements for a parameter.
+
+    For DTensor (FSDP2), returns the local shard size.
+    For regular tensors, returns the full size.
+    """
+    if DTensor is not None and isinstance(param, DTensor):
+        return param._local_tensor.numel()
+    return param.numel()
+
+
+class ParamCountCallback(Callback):
+    def on_train_begin(self, model: FastGenModel, **kwargs) -> None:
+        # get modules
+        modules = {"model": model, **model.model_dict}
+
+        # iterate over modules
+        output = {}
+        for name, module in modules.items():
+            # Logical (full model) param counts
+            trainable_params = sum(p.numel() for p in module.parameters() if p.requires_grad)
+            total_params = sum(p.numel() for p in module.parameters())
+
+            # Local (sharded) param counts - what's actually in memory on this rank
+            local_trainable_params = sum(_get_local_numel(p) for p in module.parameters() if p.requires_grad)
+            local_total_params = sum(_get_local_numel(p) for p in module.parameters())
+
+            # check if parameter counts are different across ranks
+            if world_size() > 1:
+                trainable_params = self.gather_param_counts(trainable_params)
+                total_params = self.gather_param_counts(total_params)
+                local_trainable_params = self.gather_param_counts(local_trainable_params)
+                local_total_params = self.gather_param_counts(local_total_params)
+                if len(set(total_params)) == 1 and len(set(trainable_params)) == 1:
+                    trainable_params = trainable_params[0]
+                    total_params = total_params[0]
+                if len(set(local_total_params)) == 1 and len(set(local_trainable_params)) == 1:
+                    local_trainable_params = local_trainable_params[0]
+                    local_total_params = local_total_params[0]
+
+            # logging
+            module_name = module.__class__.__name__
+            output.update(
+                {
+                    f"{name}/trainable_params": trainable_params,
+                    f"{name}/total_params": total_params,
+                    f"{name}/local_trainable_params": local_trainable_params,
+                    f"{name}/local_total_params": local_total_params,
+                }
+            )
+            if isinstance(trainable_params, list):
+                logger.warning(f"Parameter counts differ across ranks for {module_name}.")
+                for rank, (p_train, p) in enumerate(zip(trainable_params, total_params)):
+                    logger.info(
+                        f"{name} ({module_name}) has {p_train * 1.e-6:.2f} M trainable and {p * 1.e-6:.2f} M total params on rank {rank}."
+                    )
+            else:
+                logger.info(
+                    f"{name} ({module_name}) has {trainable_params * 1.e-6:.2f} M trainable and {total_params * 1.e-6:.2f} M total params (logical)."
+                )
+
+            # Report local/sharded counts
+            if isinstance(local_trainable_params, list):
+                for rank, (p_train, p) in enumerate(zip(local_trainable_params, local_total_params)):
+                    logger.info(
+                        f"{name} ({module_name}) has {p_train * 1.e-6:.2f} M trainable and {p * 1.e-6:.2f} M total params LOCAL on rank {rank}."
+                    )
+            else:
+                is_sharded = local_total_params < total_params if not isinstance(total_params, list) else True
+                if is_sharded:
+                    logger.info(
+                        f"{name} ({module_name}) has {local_trainable_params * 1.e-6:.2f} M trainable and {local_total_params * 1.e-6:.2f} M total params LOCAL per rank (sharding ratio: {world_size()}x)."
+                    )
+                else:
+                    logger.info(f"{name} ({module_name}) is NOT sharded (local == logical params).")
+
+        if wandb.run:
+            wandb.run.summary.update(output)
+
+    def gather_param_counts(self, param_count):
+        """
+        Gather parameter counts across all ranks.
+
+        Args:
+            param_count: Parameter count to gather.
+
+        Returns:
+            List of parameter counts across all ranks.
+        """
+        param_count = torch.tensor(
+            [param_count], dtype=torch.long, device="cuda" if torch.cuda.is_available() else "cpu"
+        )
+        param_count_list = [torch.zeros_like(param_count) for _ in range(world_size())]
+        torch.distributed.all_gather(param_count_list, param_count)
+        return [p.item() for p in param_count_list]

lipforcing/callbacks/stdout_logger.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Simple stdout loss logger callback for testing without wandb."""
+
+from typing import Callable
+
+import torch
+from lipforcing.callbacks.callback import Callback
+from lipforcing.methods.model import FastGenModel
+import lipforcing.utils.logging_utils as logger
+
+
+class StdoutLoggerCallback(Callback):
+    """Prints loss values to stdout at every logging_iter."""
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        logging_iter = getattr(self.config.trainer, "logging_iter", 1) if self.config else 1
+        if iteration % logging_iter == 0:
+            parts = [f"iter {iteration:5d}"]
+            for k, v in sorted(loss_dict.items()):
+                if isinstance(v, torch.Tensor):
+                    parts.append(f"{k}={v.item():.6f}")
+                elif isinstance(v, (int, float)):
+                    parts.append(f"{k}={v:.6f}")
+            logger.info(" | ".join(parts))

lipforcing/callbacks/train_profiler.py

@@ -0,0 +1,138 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, Callable
+
+import torch
+import wandb
+
+from lipforcing.callbacks.callback import Callback
+from lipforcing.utils.distributed import is_rank0
+import lipforcing.utils.logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.methods import FastGenModel
+
+
+class TrainProfilerCallback(Callback):
+    """Callback for profiling training speed and detailed timing breakdowns.
+
+    Tracks:
+    - iter_time: seconds per iteration (wall clock time)
+    - data_load_time: time spent loading data
+    - avg_forward_time: average forward pass time across accumulation steps
+    - backward_time: time spent in backward pass
+    - optim_step_time: time spent in optimizer step
+    """
+
+    def __init__(self, every_n: int = 100, detailed: bool = True):
+        """Initialize the profiler callback.
+
+        Args:
+            every_n: Log metrics every N iterations
+            detailed: If True, log detailed timing breakdown. If False, only log iter_time.
+        """
+        # For iter_time tracking
+        self.last_log_time = None
+
+        # For detailed profiling
+        self.detailed = detailed
+        self.train_step_begin_time = None
+        self.accum_begin_times = None
+        self.backward_begin_times = None
+        self.optimizer_step_begin = None
+        self.step_end_time = None
+        self.every_n = every_n
+
+    def on_train_begin(self, model: FastGenModel, iteration: int = 0) -> None:
+        if hasattr(self, "config"):
+            # overwritten by logging_iter if self.config exists
+            self.every_n = self.config.trainer.logging_iter
+        logger.info(f"every_n to profile trainer: {self.every_n}")
+
+    def on_training_step_begin(
+        self,
+        model: FastGenModel,
+        iteration: int = 0,
+    ):
+        if self.detailed:
+            self.train_step_begin_time = time.perf_counter()
+            self.accum_begin_times = []
+            self.backward_begin_times = []
+
+    def on_training_accum_step_begin(
+        self, model: FastGenModel, data_batch: dict[str, torch.Tensor], iteration: int = 0, accum_iter: int = 0
+    ):
+        if self.detailed:
+            self.accum_begin_times.append(time.perf_counter())
+
+    def on_backward_begin(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+        accum_iter: int = 0,
+    ):
+        if self.detailed:
+            self.backward_begin_times.append(time.perf_counter())
+
+    def on_optimizer_step_begin(self, model: FastGenModel, iteration: int = 0):
+        if self.detailed:
+            self.optimizer_step_begin = time.perf_counter()
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        del data_batch, output_batch, loss_dict
+
+        if self.detailed:
+            self.step_end_time = time.perf_counter()
+
+        if hasattr(self, "config"):
+            # only wandb log when config exists
+            if iteration % self.every_n == 0 and is_rank0():
+                metrics = {}
+
+                # Calculate iter_time (wall clock time per iteration)
+                cur_time = time.time()
+                if self.last_log_time is not None:
+                    iter_time = (cur_time - self.last_log_time) / self.every_n
+                    logger.info(f"{iteration} : avg iteration time       {iter_time:.2f} seconds")
+                    metrics["profiler/avg_iteration_time"] = iter_time
+                self.last_log_time = cur_time
+
+                # Calculate detailed timing breakdown
+                if self.detailed and self.accum_begin_times and self.backward_begin_times:
+                    data_load_time = self.accum_begin_times[0] - self.train_step_begin_time
+                    forward_time = sum(
+                        [b - a for (b, a) in zip(self.backward_begin_times, self.accum_begin_times)]
+                    ) / len(self.accum_begin_times)
+                    backward_time = self.optimizer_step_begin - self.backward_begin_times[-1]
+                    optim_step_time = self.step_end_time - self.optimizer_step_begin
+
+                    logger.info(f"{iteration} : data loading time        {data_load_time:.2f}")
+                    logger.info(f"{iteration} : avg forward pass time    {forward_time:.2f}")
+                    logger.info(f"{iteration} : backward pass time       {backward_time:.2f}")
+                    logger.info(f"{iteration} : optimizer step time      {optim_step_time:.2f}")
+
+                    metrics.update(
+                        {
+                            "profiler/data_loading_time": data_load_time,
+                            "profiler/avg_forward_pass_time": forward_time,
+                            "profiler/backward_pass_time": backward_time,
+                            "profiler/optimizer_step_time": optim_step_time,
+                        }
+                    )
+
+                if wandb.run and metrics:
+                    wandb.log(metrics, step=iteration)

lipforcing/callbacks/wandb.py

@@ -0,0 +1,773 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+import os
+import subprocess
+import tempfile
+from dataclasses import dataclass, field
+import time
+from typing import Optional, Dict, Callable, TYPE_CHECKING
+import gc
+
+import numpy as np
+
+
+import torch
+import torchvision
+from torchvision.transforms import functional as tv_F
+
+import wandb
+import wandb.util
+
+from lipforcing.callbacks.callback import Callback
+from lipforcing.configs.config_utils import serialize_config
+from lipforcing.utils import basic_utils
+
+from lipforcing.utils.distributed import rank0_only, synchronize, world_size
+from lipforcing.utils import logging_utils as logger
+
+if TYPE_CHECKING:
+    from lipforcing.configs.config import BaseConfig
+    from lipforcing.methods import FastGenModel
+
+
+def tensor_to_wandb_video_with_audio(
+    video_tensor: torch.Tensor,
+    audio_path: str,
+    fps: int = 25,
+    vid_format: str = "mp4",
+    caption: str | None = None,
+) -> wandb.Video:
+    """Convert a [B, T, C, H, W] uint8 video tensor + audio file to wandb.Video with audio.
+
+    Takes the first sample in the batch. Writes video to a temp file, muxes audio
+    with ffmpeg, and returns a wandb.Video from the muxed output.
+
+    Args:
+        video_tensor: uint8 tensor of shape [B, T, C, H, W] (already in 0-255 range).
+        audio_path: Path to the audio .wav file.
+        fps: Video frame rate (default 25 for OmniAvatar).
+        vid_format: Video format (default "mp4").
+        caption: Optional caption for wandb.Video.
+
+    Returns:
+        wandb.Video with muxed audio, or silent video if muxing fails.
+    """
+    try:
+        # Take first sample: [T, C, H, W]
+        vid = video_tensor[0] if video_tensor.dim() == 5 else video_tensor
+        # Convert to [T, H, W, C] uint8 on CPU
+        vid = vid.permute(0, 2, 3, 1).cpu()
+
+        tmpdir = tempfile.mkdtemp()
+        silent_path = os.path.join(tmpdir, f"silent.{vid_format}")
+        muxed_path = os.path.join(tmpdir, f"muxed.{vid_format}")
+
+        # Write silent video — try torchvision.io first, fall back to raw ffmpeg pipe
+        T, H, W, C = vid.shape
+        try:
+            torchvision.io.write_video(silent_path, vid, fps=fps, video_codec="libx264")
+        except Exception:
+            # Fallback: pipe raw frames to ffmpeg
+            write_cmd = [
+                "ffmpeg", "-y",
+                "-f", "rawvideo", "-pix_fmt", "rgb24",
+                "-s", f"{W}x{H}", "-r", str(fps),
+                "-i", "pipe:0",
+                "-c:v", "libx264", "-pix_fmt", "yuv420p",
+                "-loglevel", "error",
+                silent_path,
+            ]
+            proc = subprocess.run(write_cmd, input=vid.numpy().tobytes(), capture_output=True, timeout=60)
+            if proc.returncode != 0:
+                raise RuntimeError(f"ffmpeg raw write failed: {proc.stderr.decode()}")
+
+        # Mux audio with ffmpeg
+        cmd = [
+            "ffmpeg", "-y",
+            "-i", silent_path,
+            "-i", audio_path,
+            "-c:v", "copy",
+            "-c:a", "aac",
+            "-shortest",
+            "-loglevel", "error",
+            muxed_path,
+        ]
+        result = subprocess.run(cmd, capture_output=True, timeout=30)
+
+        if result.returncode == 0 and os.path.exists(muxed_path):
+            return wandb.Video(muxed_path, fps=fps, format=vid_format, caption=caption)
+        else:
+            logger.warning(f"ffmpeg muxing failed (rc={result.returncode}): {result.stderr.decode()}")
+            return wandb.Video(video_tensor[:1].cpu().numpy(), fps=fps, format=vid_format, caption=caption)
+
+    except Exception as e:
+        logger.warning(f"Audio muxing failed, falling back to silent video: {e}")
+        return wandb.Video(video_tensor[:1].cpu().numpy(), fps=fps, format=vid_format, caption=caption)
+
+
+def to_wandb(
+    tensor: torch.Tensor,
+    rgb_range: float = 255.0,
+    normalized: bool = False,
+    max_plot_img: int = 16,
+    max_plot_vid: int = 2,
+    fps: int = 16,
+    channel_before_time: bool = True,
+    caption: str | None = None,
+    vid_format: str = "mp4",
+) -> wandb.Image | wandb.Video:
+    """
+    Convert a tensor to a wandb.Image or wandb.Video.
+
+    Args:
+        tensor (torch.Tensor): Input tensor of shape [B,C,H,W], [B,T,C,H,W], or [B,T,C,H,W,D].
+        rgb_range (float, optional): Output target RGB range (can almost definitely be kept as 255).
+            Defaults to 255.0.
+        normalized (bool, optional): Whether the tensor is normalized to [0,1]. Defaults to False which assumes [-1,1] range.
+        max_plot_img (int, optional): Max number of images to plot. Defaults to 16.
+        max_plot_vid (int, optional): Max number of videos to plot. Defaults to 2.
+        fps (int, optional): Frames per second. Defaults to 8.
+        channel_before_time (bool, optional): Whether the tensor is in the format [B,C,T,..]. Set False if the [B,T,C,..] format is used.
+        caption (str, optional): Caption for the image or video. Defaults to None.
+        vid_format (str, optional): Format of the video file. Defaults to "mp4".
+
+    Returns:
+        wandb.Image | wandb.Video: Format a tensor for logging to W&B.
+    """
+
+    if tensor.ndim == 5:
+        max_plot = max_plot_vid
+        if channel_before_time:
+            tensor = tensor.permute(0, 2, 1, 3, 4)
+    elif tensor.ndim == 4:
+        max_plot = max_plot_img
+    else:
+        raise ValueError(f"Tensor must be 4 or 5 dimensional, but got {tensor.ndim} dimensions")
+
+    # slice and adjust range
+    if normalized:
+        factor = rgb_range
+        offset = 0.0
+    else:
+        factor = rgb_range / 2.0
+        offset = rgb_range / 2.0
+    tensor = tensor[:max_plot].mul(factor).add(offset).clip_(0, rgb_range).to(torch.uint8)
+
+    # convert to wandb.Image or wandb.Video
+    assert tensor.shape[-3] == 3, "Make sure that the data is in ..., C, H, W format"
+    if tensor.ndim == 5:
+        return wandb.Video(tensor.cpu().numpy(), fps=fps, format=vid_format, caption=caption)
+    else:
+        image_grid = torchvision.utils.make_grid(tensor, nrow=4, pad_value=1)
+        image_grid = tv_F.to_pil_image(image_grid)
+        return wandb.Image(image_grid, caption=caption)
+
+
+def _to_wandb_with_audio(
+    tensor: torch.Tensor,
+    audio_path: str,
+    fps: int = 25,
+    rgb_range: float = 255.0,
+    normalized: bool = False,
+    vid_format: str = "mp4",
+    caption: str | None = None,
+    channel_before_time: bool = True,
+) -> wandb.Video:
+    """Convert a video tensor to wandb.Video with audio muxed in.
+
+    Handles the same normalization as to_wandb, then delegates to
+    tensor_to_wandb_video_with_audio for ffmpeg muxing.
+
+    Args:
+        tensor: [B, C, T, H, W] or [B, T, C, H, W] video tensor in [-1,1] or [0,1] range.
+        audio_path: Path to audio .wav file.
+        fps: Frame rate for the output video.
+        rgb_range: Target RGB range (255).
+        normalized: Whether tensor is in [0,1] (True) or [-1,1] (False).
+        vid_format: Video file format.
+        caption: Optional caption.
+        channel_before_time: Whether tensor is [B,C,T,H,W] (True) or [B,T,C,H,W] (False).
+
+    Returns:
+        wandb.Video with audio.
+    """
+    if channel_before_time:
+        tensor = tensor.permute(0, 2, 1, 3, 4)  # [B,C,T,H,W] -> [B,T,C,H,W]
+
+    # Normalize to uint8
+    if normalized:
+        factor = rgb_range
+        offset = 0.0
+    else:
+        factor = rgb_range / 2.0
+        offset = rgb_range / 2.0
+    tensor = tensor[:1].mul(factor).add(offset).clip_(0, rgb_range).to(torch.uint8)
+
+    return tensor_to_wandb_video_with_audio(
+        tensor, audio_path, fps=fps, vid_format=vid_format, caption=caption,
+    )
+
+
+def _load_audio_waveform(audio_path: str, target_sr: int = 16000, num_frames: int = 81, fps: float = 25.0) -> Optional[torch.Tensor]:
+    """Load audio waveform from .wav for SyncCScorer. Returns [L] float32 tensor or None."""
+    if not audio_path or not os.path.isfile(audio_path):
+        return None
+    try:
+        import scipy.io.wavfile as wavfile
+        from scipy import signal
+        sr, wav = wavfile.read(audio_path)
+        if wav.dtype == np.int16:
+            wav = wav.astype(np.float32) / 32768.0
+        elif wav.dtype != np.float32:
+            wav = wav.astype(np.float32)
+        wav = torch.from_numpy(wav)
+        if wav.ndim == 2:
+            wav = wav.mean(dim=1)
+        if sr != target_sr:
+            num_samples_new = int(len(wav) * target_sr / sr)
+            wav = torch.from_numpy(signal.resample(wav.numpy(), num_samples_new))
+        target_length = int(num_frames / fps * target_sr)
+        if wav.shape[0] < target_length:
+            wav = torch.nn.functional.pad(wav, (0, target_length - wav.shape[0]))
+        else:
+            wav = wav[:target_length]
+        return wav.to(torch.float32)
+    except Exception as e:
+        logger.warning(f"[SyncEval] Failed to load audio from {audio_path}: {e}")
+        return None
+
+
+@rank0_only
+def init_wandb(config: BaseConfig):
+    # wandb login
+    wandb_credential = config.log_config.wandb_credential
+    if os.path.isfile(wandb_credential):
+        os.environ["WANDB_API_KEY"] = open(wandb_credential, encoding="utf-8").read().strip("\n")
+        logger.info(f"Loading WANDB_API_KEY from {wandb_credential}")
+
+    wandb_config = config.log_config
+
+    # Resume with or generate a wandb id
+    logger.info(f"wandb_config.save_path: {wandb_config.save_path}")
+    os.makedirs(wandb_config.save_path, exist_ok=True)
+    wandb_id_path = f"{wandb_config.save_path}/wandb_id.txt"
+    resuming = getattr(config.trainer, "resume", True)
+    if os.path.isfile(wandb_id_path) and resuming:
+        wandb_id = open(wandb_id_path, encoding="utf-8").read().strip()
+        logger.info(f"Resuming with an existing wandb id: {wandb_id}")
+    else:
+        wandb_id = wandb.util.generate_id()
+        with open(wandb_id_path, "w", encoding="utf-8") as f:
+            f.write(f"{wandb_id}\n")
+        logger.info(f"Generating a wandb id: {wandb_id}")
+
+    # Get config as plain dict
+    config_resolved = serialize_config(config, return_type="dict")
+
+    # Initialize the wandb library.
+    wandb.init(
+        id=wandb_id,
+        project=wandb_config.project,
+        group=wandb_config.group,
+        name=wandb_config.name,
+        entity=getattr(wandb_config, "wandb_entity", None),
+        config=config_resolved,
+        dir=wandb_config.save_path,
+        resume="allow",
+        mode=wandb_config.wandb_mode,
+    )
+
+    # Save a copy of code to a wandb Artifact (this can be slow)
+    # Make code upload optional to avoid distributed training delays
+    upload_code = basic_utils.str2bool(os.getenv("WANDB_UPLOAD_CODE", "false"))
+    if upload_code:
+        logger.info("Uploading code to wandb (this may take a few minutes)...")
+        wandb.run.log_code(".")
+        logger.info("Code upload to wandb completed")
+    else:
+        logger.info("Wandb code upload disabled (set WANDB_UPLOAD_CODE=true to enable)")
+
+
+@dataclass
+class _LossDictRecord:
+    loss_dict: dict = field(default_factory=dict)
+    iter_count_dict: dict = field(default_factory=dict)
+
+    def add(self, loss_dict: Optional[Dict[str, torch.Tensor]]) -> None:
+        if loss_dict is not None:
+            for loss_name, loss_val in loss_dict.items():
+                scalar = loss_val.float().item() if torch.is_tensor(loss_val) else float(loss_val)
+                self.loss_dict[loss_name] = self.loss_dict.get(loss_name, 0.0) + scalar
+                self.iter_count_dict[loss_name] = self.iter_count_dict.get(loss_name, 0) + 1
+
+    def reset(self) -> None:
+        self.loss_dict = {}
+        self.iter_count_dict = {}
+
+    def gather_dict(self, dictionary: Dict[str, float | int]) -> Dict[str, float | int]:
+        n_ranks = world_size()
+        if n_ranks > 1:
+            dict_list = [None for _ in range(n_ranks)]
+            torch.distributed.all_gather_object(dict_list, dictionary)
+            # from list of dicts to dict of summed values
+            dictionary = {}
+            for d in dict_list:
+                for key, value in d.items():
+                    dictionary[key] = dictionary.get(key, 0.0) + value
+        return dictionary
+
+    def get_stat(self) -> Dict[str, float]:
+        # number of ranks that logged this loss
+        rank_dict = self.gather_dict({k: 1 for k in self.loss_dict.keys()})
+        # number of times this loss was computed
+        count_dict = self.gather_dict(self.iter_count_dict)
+        # sum of all losses
+        loss_dict = self.gather_dict(self.loss_dict)
+
+        avg_loss_dict = {}
+        for loss_name, loss_val in loss_dict.items():
+            count = count_dict.get(loss_name, 0)
+            ranks = rank_dict.get(loss_name, 1)
+            iter_count = count / ranks
+            avg_loss = (loss_val / count) * (ranks / world_size()) if count > 0 else 0.0
+            logger.info(f"avg_{loss_name}: {avg_loss:.4f}".ljust(30) + f"iter count: {iter_count}")
+            avg_loss_dict[loss_name] = avg_loss
+        self.reset()
+        return avg_loss_dict
+
+
+class WandbCallback(Callback):
+    """
+    The callback gets precision for data from model
+    """
+
+    def __init__(
+        self,
+        *args,
+        validation_logging_step: int = 1,
+        sample_logging_iter: Optional[int] = None,
+        vid_format: str = "mp4",
+        fps: int = 16,
+        syncnet_checkpoint_path: Optional[str] = None,
+        syncnet_vshift: int = 15,
+        syncnet_audio_sr: int = 16000,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+
+        self.validation_logging_step = validation_logging_step
+        self.sample_logging_iter = sample_logging_iter
+        self.val_sample_map = None
+        self._val_gen_videos: list[torch.Tensor] = []
+        self._val_gt_videos: list[torch.Tensor] = []
+        self._val_audio_paths: list[str | None] = []
+        self.vid_format = vid_format
+        self.fps = fps
+        self.syncnet_checkpoint_path = syncnet_checkpoint_path
+        self.syncnet_vshift = syncnet_vshift
+        self.syncnet_audio_sr = syncnet_audio_sr
+        self._syncnet_scorer = None
+        self._val_sync_c_scores: list[float] = []
+        self.loss_dict_record = _LossDictRecord()
+        self.val_loss_dict_record = _LossDictRecord()
+
+    def on_app_begin(self) -> None:
+        assert hasattr(self, "config"), "Missing config in WandbCallback."
+        init_wandb(self.config)
+        self.offload_module_in_decoding = self.config.trainer.offload_module_in_decoding
+        # disable offloading if using FSDP
+        if self.config.trainer.fsdp:
+            self.offload_module_in_decoding = False
+        if self.sample_logging_iter is None:
+            self.sample_logging_iter = self.config.trainer.logging_iter
+        synchronize()
+
+    def _get_syncnet_scorer(self):
+        if self._syncnet_scorer is not None:
+            return self._syncnet_scorer
+        if not self.syncnet_checkpoint_path:
+            return None
+        try:
+            from lipforcing.methods.reward.sync_c_scorer import SyncCScorer
+            self._syncnet_scorer = SyncCScorer(
+                checkpoint_path=self.syncnet_checkpoint_path,
+                input_fps=25.0,
+                audio_sample_rate=self.syncnet_audio_sr,
+                vshift=self.syncnet_vshift,
+                device="cuda" if torch.cuda.is_available() else "cpu",
+                dtype=torch.float32,
+            )
+            logger.info(f"[WandbCallback] Loaded SyncCScorer from {self.syncnet_checkpoint_path}")
+        except Exception as e:
+            logger.warning(f"[WandbCallback] Failed to load SyncCScorer: {e}")
+            self.syncnet_checkpoint_path = None
+        return self._syncnet_scorer
+
+    def on_dataloader_init_end(
+        self, model: FastGenModel, dataloader_train, dataloader_val, iteration: int = 0
+    ) -> None:
+        """Upload GT validation videos at the start so they're always available for comparison.
+
+        Not decorated with @rank0_only — all ranks must enter this method to stay
+        synchronized (synchronize() calls dist.barrier). Only rank 0 does the actual
+        VAE decode and wandb upload.
+        """
+        if dataloader_val is None:
+            return
+        # Skip GT upload if SKIP_GT_VAL_UPLOAD env var is set (avoids NCCL timeout)
+        if os.environ.get("SKIP_GT_VAL_UPLOAD", "0") == "1":
+            if wandb.run:
+                logger.info("SKIP_GT_VAL_UPLOAD=1 — skipping GT val video upload")
+            synchronize()
+            return
+        if iteration > 0:
+            if wandb.run:
+                logger.info("Resuming from checkpoint — skipping GT val video upload (already logged)")
+            synchronize()
+            return
+        if not hasattr(model.net, "vae"):
+            if wandb.run:
+                logger.info("No VAE loaded — skipping GT val video upload")
+            synchronize()
+            return
+
+        # Only rank 0 decodes and uploads; other ranks wait at the barrier below
+        if wandb.run:
+            logger.info("Uploading GT validation videos to wandb...")
+            device = model.device
+            try:
+                gt_videos = []
+                audio_paths = []
+                with torch.no_grad(), basic_utils.inference_mode(
+                    precision_amp=model.precision_amp_enc, device_type=device.type
+                ):
+                    for step, data in enumerate(dataloader_val):
+                        real = data["real"].to(device)  # [1, 16, 21, 64, 64]
+                        decoded = model.net.vae.decode(real[:1])  # [1, C, T, H, W]
+                        gt_videos.append(self._to_uint8_video(decoded))
+                        ap = None
+                        if "audio_path" in data:
+                            raw = data["audio_path"]
+                            if isinstance(raw, (list, tuple)) and len(raw) > 0 and raw[0]:
+                                ap = raw[0] if os.path.isfile(raw[0]) else None
+                        audio_paths.append(ap)
+                gt_list = []
+                for v, ap in zip(gt_videos, audio_paths):
+                    if ap:
+                        gt_list.append(tensor_to_wandb_video_with_audio(v, ap, fps=self.fps))
+                    else:
+                        gt_list.append(wandb.Video(v[0].numpy(), fps=self.fps, format="mp4"))
+                wandb.log({"val_gt/videos": gt_list}, step=0)
+                logger.info(f"Uploaded {len(gt_videos)} GT validation videos to wandb")
+            except Exception as e:
+                logger.warning(f"Failed to upload GT val videos: {e}")
+        synchronize()
+
+    @rank0_only
+    def on_optimizer_step_begin(self, model: FastGenModel, iteration: int = 0) -> None:
+        assert hasattr(self, "config"), "Missing config in WandbCallback."
+        if iteration % self.config.trainer.logging_iter == 0:
+            for name, scheduler in model.scheduler_dict.items():
+                wandb.log({f"optimizer/lr_{name}": scheduler.get_last_lr()[0]}, step=iteration)
+
+    def get_sample_map(
+        self, model: FastGenModel, data_batch: dict[str, torch.Tensor], output_batch: dict[str, torch.Tensor | Callable]
+    ) -> dict[str, wandb.Image | wandb.Video]:
+        # Collect generated and real data and create copies to avoid modifying the original dicts
+        sample_map = {}
+        gen_rand = output_batch["gen_rand"]
+        if isinstance(gen_rand, Callable):
+            synchronize()
+            gen_rand = gen_rand()
+            synchronize()
+
+        # Avoid modifying the original dicts
+        data_batch = data_batch.copy()
+        output_batch = output_batch.copy()
+
+        # Decide whether we want to visualize multistep teacher generation
+        if self.config.trainer.visualize_teacher:
+            assert "input_rand" in output_batch, "We need to know the noise to visualize teacher generation"
+            teacher_output = model.sample(
+                model.teacher,
+                output_batch["input_rand"][0:1],
+                data_batch["condition"][0:1],  # e.g. text condition encoded by the text encoder
+                data_batch["neg_condition"][0:1],  # e.g. negative text condition encoded by the text encoder
+            )
+            output_batch["gen_teacher"] = teacher_output
+
+        # Decode to pixel if it's in latent space
+        if hasattr(model.net, "init_preprocessors"):
+            torch.cuda.empty_cache()
+            device_nets = model.device
+
+            has_vae = hasattr(model.net, "vae")
+            if not has_vae:
+                model.net.init_vae()
+                model.net.vae.to(device=device_nets, dtype=model.precision)
+
+            if self.offload_module_in_decoding:
+                # offload the unneeded models to CPU (enable it if hitting OOM here)
+                logger.info(
+                    f"GPU Memory BEFORE moving nets to CPU: {torch.cuda.memory_allocated(device_nets) / 1024 ** 2:.2f} MB"
+                )
+                if hasattr(model, "fake_score"):
+                    model.fake_score = model.fake_score.to("cpu")
+                if hasattr(model, "teacher"):
+                    model.teacher = model.teacher.to("cpu")
+                logger.info(
+                    f"GPU Memory AFTER moving nets to CPU: {torch.cuda.memory_allocated(device_nets) / 1024 ** 2:.2f} MB"
+                )
+                synchronize()
+
+            with basic_utils.inference_mode(precision_amp=model.precision_amp_enc, device_type=device_nets.type):
+                if "real" in data_batch:
+                    # only generate one sample for video
+                    limit = 1 if len(data_batch["real"].shape) == 5 else len(data_batch["real"])
+                    data_batch["real"] = model.net.vae.decode(data_batch["real"][:limit])
+                if isinstance(gen_rand, dict):
+                    for k in gen_rand:
+                        limit = 1 if len(gen_rand[k].shape) == 5 else len(gen_rand[k])
+                        gen_rand[k] = model.net.vae.decode(gen_rand[k][:limit])
+                else:
+                    limit = 1 if len(gen_rand.shape) == 5 else len(gen_rand)
+                    gen_rand = model.net.vae.decode(gen_rand[:limit])
+
+                if "gen_teacher" in output_batch:
+                    output_batch["gen_teacher"] = model.net.vae.decode(output_batch["gen_teacher"][:limit])
+                if logger.LOG_LEVEL == "DEBUG" and "gen_rand_train" in output_batch:
+                    output_batch["gen_rand_train"] = model.net.vae.decode(output_batch["gen_rand_train"][:limit])
+
+            if not has_vae:
+                del model.net.vae
+
+            if self.offload_module_in_decoding:
+                # move back fake_score to gpu
+                if hasattr(model, "fake_score"):
+                    model.fake_score = model.fake_score.to(device_nets)
+                if hasattr(model, "teacher"):
+                    model.teacher = model.teacher.to(device_nets)
+                logger.info(
+                    f"GPU Memory AFTER moving nets back to GPU: {torch.cuda.memory_allocated(device_nets) / 1024 ** 2:.2f} MB"
+                )
+                synchronize()
+
+        if wandb.run:
+            if (
+                "condition_raw" in data_batch
+                and isinstance(data_batch["condition_raw"], (list, tuple))
+                and isinstance(data_batch["condition_raw"][0], str)
+            ):
+                caption = "\n".join(data_batch["condition_raw"][: len(gen_rand)])
+            else:
+                caption = None
+
+            # Check for audio path (from OmniAvatar dataloader) for audio-muxed video logging.
+            # The dataloader sets audio_path="" when audio.wav doesn't exist.
+            audio_path = None
+            if "audio_path" in data_batch:
+                ap = data_batch["audio_path"]
+                # default_collate turns strings into a list
+                if isinstance(ap, (list, tuple)) and len(ap) > 0:
+                    audio_path = ap[0] if ap[0] else None
+                elif isinstance(ap, str):
+                    audio_path = ap if ap else None
+                # Verify the file actually exists
+                if audio_path and not os.path.isfile(audio_path):
+                    logger.warning(f"audio_path does not exist, logging silent video: {audio_path}")
+                    audio_path = None
+
+            if isinstance(gen_rand, dict):
+                for k in gen_rand:
+                    sample_map[f"student/generation/{k}"] = to_wandb(
+                        gen_rand[k], caption=caption, vid_format=self.vid_format
+                    )
+            else:
+                if audio_path and gen_rand.ndim == 5:
+                    sample_map["student/generation"] = _to_wandb_with_audio(
+                        gen_rand, audio_path, fps=self.fps, vid_format=self.vid_format, caption=caption,
+                    )
+                else:
+                    sample_map["student/generation"] = to_wandb(gen_rand, caption=caption, vid_format=self.vid_format, fps=self.fps)
+            if "real" in data_batch:
+                if audio_path and data_batch["real"].ndim == 5:
+                    sample_map["data/real"] = _to_wandb_with_audio(
+                        data_batch["real"], audio_path, fps=self.fps, vid_format=self.vid_format, caption=caption,
+                    )
+                else:
+                    sample_map["data/real"] = to_wandb(data_batch["real"], caption=caption, vid_format=self.vid_format, fps=self.fps)
+            if "gen_teacher" in output_batch:
+                sample_map["teacher/generation"] = to_wandb(
+                    output_batch["gen_teacher"], caption=caption, vid_format=self.vid_format
+                )
+            if logger.LOG_LEVEL == "DEBUG" and "gen_rand_train" in output_batch:
+                sample_map["student/generation_train"] = to_wandb(
+                    output_batch["gen_rand_train"], caption=caption, vid_format=self.vid_format
+                )
+
+        return sample_map
+
+    def log_sample_map(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        suffix: str = "",
+        iteration: int = 0,
+        group: str = "train",
+    ) -> None:
+        sample_map = self.get_sample_map(model, data_batch, output_batch)
+        sample_map = {f"{group}_media/{k}{suffix}": v for k, v in sample_map.items()}
+        if wandb.run:
+            wandb.log(sample_map, step=iteration)
+        synchronize()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def log_stats(self, loss_dict_record: _LossDictRecord, iteration: int = 0, group: str = "train") -> None:
+        logger.info(f"logging {group} stats at iteration {iteration}" + "-" * 20)
+        # Collect distributed statistics
+        avg_loss_dict = loss_dict_record.get_stat()
+        stats = {f"{group}/{name}": val for name, val in avg_loss_dict.items()}
+        base_info = {"optimizer/iteration": iteration}
+
+        # log stats and base info
+        if wandb.run:
+            wandb.log(stats, step=iteration)
+            wandb.log(base_info, step=iteration)
+
+    def on_training_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        iteration: int = 0,
+    ) -> None:
+        self.loss_dict_record.add(loss_dict)
+        time_start = time.perf_counter()
+        logged = False
+        if iteration % self.config.trainer.logging_iter == 0 or iteration == 1:
+            self.log_stats(self.loss_dict_record, iteration=iteration, group="train")
+            logged = True
+        skip_early_sample = os.environ.get("SKIP_EARLY_SAMPLE_LOG", "0") == "1"
+        if iteration % self.sample_logging_iter == 0 or (iteration == 1 and not skip_early_sample):
+            self.log_sample_map(model, data_batch, output_batch, iteration=iteration, group="train")
+            logged = True
+        if logged:
+            time_taken = time.perf_counter() - time_start
+            logger.info(f"WandB logging complete after {time_taken:.2f} seconds")
+
+    @staticmethod
+    def _to_uint8_video(tensor: torch.Tensor, normalized: bool = False) -> torch.Tensor:
+        """Convert [B, C, T, H, W] float video to [B, T, C, H, W] uint8 on CPU."""
+        t = tensor.permute(0, 2, 1, 3, 4)  # [B, C, T, H, W] -> [B, T, C, H, W]
+        if normalized:
+            t = t.mul(255.0)
+        else:
+            t = t.mul(127.5).add(127.5)
+        return t.clamp(0, 255).to(torch.uint8).cpu()
+
+    def on_validation_step_end(
+        self,
+        model: FastGenModel,
+        data_batch: dict[str, torch.Tensor],
+        output_batch: dict[str, torch.Tensor | Callable],
+        loss_dict: dict[str, torch.Tensor],
+        step: int = 0,
+        iteration: int = 0,
+        idx: int = 0,
+    ) -> None:
+        self.val_loss_dict_record.add(loss_dict)
+
+        if step % self.validation_logging_step == 0:
+            has_vae = hasattr(model.net, "vae")
+            if not has_vae:
+                return
+
+            # AR-generate the video
+            gen_rand = output_batch.get("gen_rand")
+            if gen_rand is not None and isinstance(gen_rand, Callable):
+                synchronize()
+                gen_rand = gen_rand()
+                synchronize()
+
+            if gen_rand is None:
+                return
+
+            # VAE decode + video collection — rank 0 only to avoid FSDP deadlock.
+            # Other ranks wait at synchronize() below.
+            _rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
+            if _rank == 0:
+                device = model.device
+                with torch.no_grad(), basic_utils.inference_mode(
+                    precision_amp=model.precision_amp_enc, device_type=device.type
+                ):
+                    gen_decoded = model.net.vae.decode(gen_rand[:1])
+                    gt_decoded = model.net.vae.decode(data_batch["real"][:1].to(device))
+
+                self._val_gen_videos.append(self._to_uint8_video(gen_decoded))
+                self._val_gt_videos.append(self._to_uint8_video(gt_decoded))
+
+                # Extract audio path for muxing
+                audio_path = None
+                if "audio_path" in data_batch:
+                    ap = data_batch["audio_path"]
+                    if isinstance(ap, (list, tuple)) and len(ap) > 0 and ap[0]:
+                        audio_path = ap[0] if os.path.isfile(ap[0]) else None
+                self._val_audio_paths.append(audio_path)
+
+                # SyncNet-v2 evaluation on generated video
+                scorer = self._get_syncnet_scorer()
+                if scorer is not None and audio_path is not None:
+                    try:
+                        audio_wav = _load_audio_waveform(
+                            audio_path, target_sr=self.syncnet_audio_sr,
+                        )
+                        if audio_wav is not None:
+                            pixels = gen_decoded.clamp(-1.0, 1.0)
+                            u8 = ((pixels + 1.0) * 127.5).to(torch.uint8)
+                            face_frames = u8[0].permute(1, 0, 2, 3).contiguous()  # [T, 3, H, W]
+                            sync_c = scorer._score_single(face_frames, audio_wav)
+                            self._val_sync_c_scores.append(sync_c.item())
+                            logger.info(f"[SyncEval] val sample {step}: sync_c={sync_c.item():.3f}")
+                    except Exception as e:
+                        logger.warning(f"[SyncEval] Failed on val sample {step}: {e}")
+
+            synchronize()
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def on_validation_end(self, model: FastGenModel, iteration: int = 0, idx: int = 0) -> None:
+        self.log_stats(self.val_loss_dict_record, iteration=iteration, group=f"val{idx}")
+        if wandb.run and self._val_gen_videos:
+            gen_list = []
+            gt_list = []
+            for i, (gen_v, gt_v) in enumerate(zip(self._val_gen_videos, self._val_gt_videos)):
+                ap = self._val_audio_paths[i] if i < len(self._val_audio_paths) else None
+                caption = None
+                if i < len(self._val_sync_c_scores):
+                    caption = f"sync_c={self._val_sync_c_scores[i]:.3f}"
+                if ap:
+                    gen_list.append(tensor_to_wandb_video_with_audio(gen_v, ap, fps=self.fps, caption=caption))
+                    gt_list.append(tensor_to_wandb_video_with_audio(gt_v, ap, fps=self.fps))
+                else:
+                    gen_list.append(wandb.Video(gen_v[0].numpy(), fps=self.fps, format="mp4", caption=caption))
+                    gt_list.append(wandb.Video(gt_v[0].numpy(), fps=self.fps, format="mp4"))
+            wandb.log({
+                f"val{idx}/generated": gen_list,
+                f"val{idx}/reconstructed": gt_list,
+            }, step=iteration)
+            logger.info(f"Logged {len(self._val_gen_videos)} val videos at iteration {iteration}")
+            if self._val_sync_c_scores:
+                mean_sync_c = sum(self._val_sync_c_scores) / len(self._val_sync_c_scores)
+                wandb.log({f"val{idx}/sync_c_mean": mean_sync_c}, step=iteration)
+                for i, sc in enumerate(self._val_sync_c_scores):
+                    wandb.log({f"val{idx}/sync_c_sample_{i}": sc}, step=iteration)
+                logger.info(f"[SyncEval] val{idx} mean sync_c: {mean_sync_c:.3f} (n={len(self._val_sync_c_scores)})")
+        self._val_gen_videos = []
+        self._val_gt_videos = []
+        self._val_audio_paths = []
+        self._val_sync_c_scores = []

lipforcing/configs/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

lipforcing/configs/callbacks.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from lipforcing.utils import LazyCall as L
+
+from lipforcing.callbacks.ct_schedule import CTScheduleCallback
+from lipforcing.callbacks.grad_clip import GradClipCallback
+from lipforcing.callbacks.param_count import ParamCountCallback
+from lipforcing.callbacks.wandb import WandbCallback
+from lipforcing.callbacks.ema import EMACallback
+from lipforcing.callbacks.train_profiler import TrainProfilerCallback
+from lipforcing.callbacks.gpu_stats import GPUStatsCallback
+from lipforcing.callbacks.forced_weight_norm import ForcedWeightNormCallback
+from lipforcing.callbacks.gpu_mem_profiler import MemTrackerCallback
+
+
+CTSchedule_CALLBACK = dict(
+    ct_schedule=L(CTScheduleCallback)(q=2.0, ratio_limit=0.999, kimg_per_stage=12500),
+)
+
+EMA_CALLBACK = dict(
+    ema=L(EMACallback)(
+        type="constant", beta=0.9999, gamma=16.97, ema_halflife_kimg=500, ema_rampup_ratio=0.05, start_iter=0
+    ),
+)
+
+EMA_CONST_CALLBACKS = dict(
+    ema_9999=L(EMACallback)(type="constant", beta=0.9999, ema_name="ema_9999"),
+    ema_99995=L(EMACallback)(type="constant", beta=0.99995, ema_name="ema_99995"),
+    ema_9996=L(EMACallback)(type="constant", beta=0.9996, ema_name="ema_9996"),
+)
+
+EMA_POWER_CALLBACKS = dict(
+    ema_1=L(EMACallback)(type="power", gamma=96.99, ema_name="ema_1"),
+    ema_5=L(EMACallback)(type="power", gamma=16.97, ema_name="ema_5"),
+    ema_10=L(EMACallback)(type="power", gamma=6.94, ema_name="ema_10"),
+)
+
+ForcedWeightNorm_CALLBACK = dict(
+    forced_weight_norm=L(ForcedWeightNormCallback)(),
+)
+
+GradClip_CALLBACK = dict(
+    grad_clip=L(GradClipCallback)(grad_norm=10.0, model_key="net"),
+)
+
+GPUStats_CALLBACK = dict(
+    gpu_stats=L(GPUStatsCallback)(every_n=100),
+)
+
+ParamCount_CALLBACK = dict(
+    param_count=L(ParamCountCallback)(),
+)
+
+TrainProfiler_CALLBACK = dict(
+    train_profiler=L(TrainProfilerCallback)(every_n=100),
+)
+
+WANDB_CALLBACK = dict(
+    wandb=L(WandbCallback)(sample_logging_iter=None),
+)
+
+MemTracker_CALLBACK = dict(
+    mem_tracker=L(MemTrackerCallback)(),
+)

lipforcing/configs/config.py

@@ -0,0 +1,282 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Any, List, Optional, Dict
+
+import copy
+import attrs
+from omegaconf import DictConfig
+
+from lipforcing.utils import LazyCall as L
+from lipforcing.configs.callbacks import WANDB_CALLBACK
+from lipforcing.configs.opt import BaseOptimizerConfig, BaseSchedulerConfig
+from lipforcing.methods import FastGenModel
+
+
+@attrs.define(slots=False)
+class CuDNNConfig:
+    # If set to True, cudnn will use deterministic cudnn functions for better reproducibility.
+    deterministic: bool = False
+    # If set to True, cudnn will benchmark several algorithms and pick the fastest one.
+    benchmark: bool = True
+
+
+@attrs.define(slots=False)
+class LogConfig:
+    # Project name
+    project: str = "LipForcing"
+    # Experiment name
+    group: str = "default"
+    # Run/job name
+    name: str = "debug"
+    # W&B mode, can be "online" or "disabled".
+    wandb_mode: str = "online"
+    # W&B entity (team or username)
+    wandb_entity: Optional[str] = None
+    # Wandb credential path
+    wandb_credential: str = "./credentials/wandb_api.txt"
+
+    # save path
+    @property
+    def save_path(self) -> str:
+        return os.path.join(
+            os.environ.get("LIPFORCING_OUTPUT_ROOT", "outputs"), f"{self.project}/{self.group}/{self.name}"
+        )
+
+
+@attrs.define(slots=False)
+class EvalConfig:
+    # Number of samples to generate
+    num_samples: int = 50000
+    # Save a small batch of images
+    save_images: bool = False
+    # Minimum checkpoint to evaluate
+    min_ckpt: int = 0
+    # Maximum checkpoint to evaluate
+    max_ckpt: int = 100000000
+    # Directory to save samples
+    samples_dir: str = "samples"
+
+
+@attrs.define(slots=False)
+class BaseCheckpointerConfig:
+    save_dir: str = "checkpoints"
+    use_s3: bool = False
+    s3_container: str = "s3://checkpoints/lipforcing"
+    s3_credential: str = "./credentials/s3.json"
+
+    # path to pretrained model (from previous stages),
+    # it's used by loading fsdp/ddp trained ckpt to an fsdp/ddp pipeline
+    pretrained_ckpt_path: str = ""
+    # submodule names of model and keys of a pretrained checkpoint of the form {"model": {"submodule_key": ...}, ...}
+    pretrained_ckpt_key_map: Dict[str, str] = {"net": "net"}
+
+
+@attrs.define(slots=False)
+class SampleTConfig:
+    """Config for sampling t from a time distribution."""
+
+    # time distribution (currently supporting: uniform, lognormal, polynomial, logitnormal, shift, and log_t)
+    time_dist_type: str = "uniform"
+    # mu in lognormal, logitnormal, and log_t distributions
+    train_p_mean: float = -1.1
+    # sigma in lognormal, logitnormal, and log_t distributions
+    train_p_std: float = 2.0
+    # shift value in shifted sampling (t_shifted = t * shift / (t * (shift - 1) + 1))
+    shift: float = 5.0
+    # lowest value in truncated range
+    min_t: float = 0.002
+    # highest value in truncated range
+    max_t: float = 80.0
+    # If provided, it is in the form [t_max, ..., 0] where len(t_list) needs to equal student_sample_steps + 1
+    t_list: Optional[List[float]] = None
+    # degree of freedom in log-transformed student-t distribution
+    log_t_df: float = 0.01
+
+
+@attrs.define(slots=False)
+class SyncWindowCFGConfig:
+    """Sync-Window DMD (SW-DMD) — timestep-gated classifier-free guidance.
+
+    This is the paper's Sync-Window DMD (SW-DMD; Sec. 4.3, Eq. 6): the teacher
+    CFG scale is gated by the DMD re-noising timestep. When enabled, CFG uses the
+    configured guidance_scale only when t is inside the sync window [t_lo, t_hi]
+    (the shifted-timestep band corresponding to teacher ODE steps j in [20, 40]);
+    outside the window the effective scale is 1.0 (no-CFG, fidelity-preserving).
+    Both teacher forward passes always run for FSDP consistency.
+
+    When ``reverse=True``, the window is inverted: CFG is OFF inside [t_lo, t_hi]
+    and ON outside it (the reverse-window ablation).
+    """
+    enabled: bool = False
+    t_lo: float = 0.0
+    t_hi: float = 1.0
+    reverse: bool = False
+
+
+@attrs.define(slots=False)
+class BaseModelConfig:
+    # Use factory functions to ensure each instance gets its own copy
+    net: Optional[dict] = None  # set per-experiment (OmniAvatar configs assign the causal/bidirectional net)
+    teacher: Optional[dict] = None  # Usually not used, only used when teacher is different from net (i.e. Causvid)
+    fake_score_net: Optional[dict] = None  # When critic architecture differs from teacher (e.g. 1.3B fake_score with 14B teacher)
+
+    # guidance scale for classifier-free guidance in teacher diffusion model. None means no guidance.
+    guidance_scale: Optional[float] = None
+    # Sync-Window DMD (SW-DMD): timestep-gated CFG (applies when guidance_scale is not None)
+    sync_window_cfg: SyncWindowCFGConfig = attrs.field(factory=SyncWindowCFGConfig)
+
+    # enable skip layer guidance (currently only wan network has the skip_layers option in cfg)
+    skip_layers: List[int] | None = None
+
+    # optimizer and scheduler for the main net (i.e., one-step generator in DMD)
+    net_optimizer: dict = attrs.field(factory=lambda: copy.deepcopy(BaseOptimizerConfig))
+    net_scheduler: dict = attrs.field(factory=lambda: copy.deepcopy(BaseSchedulerConfig))
+
+    # sampling t from a given distribution
+    sample_t_cfg: SampleTConfig = attrs.field(factory=SampleTConfig)
+
+    # shape of the input to the model (defaults to CIFAR-10)
+    input_shape: List[int] = [3, 32, 32]
+    # device ("cuda" or "cpu")
+    device: str = "cuda"
+
+    # enable gradient scaler
+    grad_scaler_enabled: bool = False
+    grad_scaler_init_scale: float = 65536.0
+    grad_scaler_growth_interval: int = 2000
+
+    # path to the pretrained teacher model ckpt
+    pretrained_model_path: str = ""
+    # path to the pretrained student net ckpt (if different from the teacher)
+    pretrained_student_net_path: str = ""
+    # initialize student from the above checkpoints (can be turned off to only load weights to the teacher)
+    load_student_weights: bool = True
+
+    # enable preprocessors in the model
+    enable_preprocessors: bool = True
+
+    # EMA for the main net (requires EMACallback)
+    use_ema: Any = False
+
+    # multistep generation if larger than 1 (default: single-step generation)
+    student_sample_steps: int = 1
+    # sampling type in multistep generation ('sde', 'ode')
+    student_sample_type: str = "sde"
+
+    # Enable memory-efficient model loading with meta device:
+    # - Rank 0 loads pretrained weights normally
+    # - Other ranks use torch.device("meta") for ZERO memory allocation (just metadata)
+    # - FSDP materializes meta tensors and broadcasts weights from rank 0
+    # This dramatically speeds up initialization for large models (14B+):
+    # - Reduces RAM from N*model_size to 1*model_size
+    # - Eliminates disk I/O contention (N parallel reads -> 1 read)
+    # - Expected speedup: 30+ min -> <1 min for 14B models on 8 GPUs
+    fsdp_meta_init: bool = False
+
+    # whether to add the teacher model to the fsdp_dict
+    add_teacher_to_fsdp_dict: bool = True
+
+    # whether to find unused parameters in ddp
+    # - can be turned off for improved performance
+    # - however, it is required if the model has a discriminator or the net initializes unused modules (e.g., for logvar predictions)
+    ddp_find_unused_parameters: bool = True
+
+    # precision variables (choose from "float64", "float32", "bfloat16", or "float16")
+    # (precision of the time steps is handled in the noise scheduler, defaulting to float64 for numerical stability)
+
+    # precision for model/optimizer states and data - recommended to be float32 if precision_amp is not None
+    precision: str = "float32"
+    # AMP during training - if None or equal to precision, AMP is disabled during training.
+    precision_amp: str | None = None
+    # AMP during inference - if None or equal to precision, AMP is disabled during inference.
+    precision_amp_infer: str | None = None
+    # AMP during en-/decoding (e.g., for VAEs or text encoders) - if None or equal to precision, AMP is disabled during en-/decoding.
+    precision_amp_enc: str | None = None
+    # FSDP2 precision for parameter storage and gradient reduction.
+    # If None, defaults to `precision`. Useful for storing params/grads in float32 while computing in bfloat16.
+    precision_fsdp: str | None = None
+
+
+@attrs.define(slots=False)
+class BaseTrainerConfig:
+    cudnn: CuDNNConfig = attrs.field(factory=CuDNNConfig)
+    checkpointer: BaseCheckpointerConfig = attrs.field(factory=BaseCheckpointerConfig)
+
+    # Callbacks configs.
+    callbacks: dict = DictConfig(WANDB_CALLBACK)
+
+    # save checkpoint frequency
+    save_ckpt_iter: int = 5000
+    # test on validation set frequency
+    validation_iter: int = 1000
+    # skip validation before first training step (avoids compilation hang with FlexAttention)
+    skip_initial_validation: bool = False
+    # logging frequency
+    logging_iter: int = 1000
+    # maximum training iteration
+    max_iter: int = 1000000
+    # whether to visualize multistep teacher generation
+    visualize_teacher: bool = False
+
+    # Set the random seed.
+    seed: int = 0
+    # Validation seed
+    val_seed: int | None = None
+    # Resume
+    resume: bool = True
+
+    # DDP Parallelism
+    ddp: bool = False
+    # FSDP Parallelism
+    fsdp: bool = False
+    # Enable TensorFloat32 (convolution and matmul)
+    tf32_enabled: bool = True
+
+    # Number of gradient accumulation rounds
+    grad_accum_rounds: int = 1
+
+    # Global batch size (if not None, overrides grad_accum_rounds to match the specified batch size)
+    batch_size_global: int | None = None
+
+    # offload other modules to cpu during latent decoding
+    offload_module_in_decoding: bool = False
+
+    # apply cpu offloading in fsdp
+    fsdp_cpu_offload: bool = False
+    # Fallback minimum number of parameters for FSDP wrapping
+    # (10M wraps large models into fairly small shards)
+    # The FastGenNetwork should provide a fully_shard method that can be used to shard the network.
+    # If we need to shard a different module, we fall back to an auto-sharding policy based on this value.
+    fsdp_min_num_params: int = 10_000_000
+    # Sharding group size for FSDP. If None, fully shard across all ranks.
+    # If set, creates a 2D mesh with (replicate, shard) dimensions.
+    fsdp_sharding_group_size: Optional[int] = None
+
+    # global variables
+    global_vars: Optional[dict] = None
+    global_vars_val: List[dict | None] = [None]
+
+    # augment config
+    augment_pipe: Optional[DictConfig] = None
+
+
+@attrs.define(slots=False)
+class BaseConfig:
+    # Log config.
+    log_config: LogConfig = attrs.field(factory=LogConfig)
+
+    # Trainer configs.
+    trainer: BaseTrainerConfig = attrs.field(factory=BaseTrainerConfig)
+
+    # Model configs.
+    model: BaseModelConfig = attrs.field(factory=BaseModelConfig)
+    model_class: DictConfig = L(FastGenModel)(config=None)
+
+    # Data configs.
+    dataloader_train: dict = attrs.field(factory=lambda: DictConfig({"batch_size": 1}))  # placeholder; experiments assign the real loader
+    dataloader_val: Any = None
+
+    # Eval configs.
+    eval: EvalConfig = attrs.field(factory=EvalConfig)

lipforcing/configs/config_utils.py

@@ -0,0 +1,266 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from typing import Any, Optional, List, Dict
+import inspect
+from copy import deepcopy
+
+import attrs
+import yaml
+from omegaconf import DictConfig, OmegaConf, ListConfig
+from hydra import compose, initialize
+from hydra.core.config_store import ConfigStore
+
+import importlib
+from dataclasses import fields as dataclass_fields
+import attr
+from dataclasses import is_dataclass
+import lipforcing.utils.logging_utils as logger
+
+
+def import_config_from_python_file(config_file: str) -> Any:
+    """
+    Import a config from a python file.
+
+    Args:
+        config_file (str): The path to the python file.
+
+    Returns:
+        Any: The config object.
+    """
+
+    if not config_file.endswith(".py"):
+        raise ValueError("Config file must be a Python file with a .py extension. " f"Received: {config_file}")
+
+    if not os.path.isfile(config_file):
+        raise FileNotFoundError(f"Config file ({config_file}) not found.")
+
+    # Convert to importable module format.
+    config_module = config_file.replace("/", ".").replace(".py", "")
+
+    # Import the module
+    try:
+        config = importlib.import_module(config_module)
+    except ImportError as e:
+        logger.error(f"Failed to import config from python file: {e}")
+        raise e
+
+    return config.create_config()
+
+
+def config_from_dict(ref_instance: Any, kwargs: Any) -> Any:
+    """
+    Construct an instance of the same type as ref_instance using the provided dictionary or data or unstructured data
+
+    Args:
+        ref_instance: The reference instance to determine the type and fields when needed
+        kwargs: A dictionary of keyword arguments to use for constructing the new instance or primitive data or unstructured data
+
+    Returns:
+        Any: A new instance of the same type as ref_instance constructed using the provided kwargs or the primitive data or unstructured data
+
+    Raises:
+        AssertionError: If the fields do not match or if extra keys are found.
+        Exception: If there is an error constructing the new instance.
+    """
+    is_type = is_attrs_or_dataclass(ref_instance)
+    if not is_type:
+        return kwargs
+    else:
+        ref_fields = set(get_fields(ref_instance))
+        assert isinstance(kwargs, dict) or isinstance(kwargs, DictConfig), "kwargs must be a dictionary or a DictConfig"
+        keys = set(kwargs.keys())
+
+        # ref_fields must equal to or include all keys
+        extra_keys = keys - ref_fields
+        assert (
+            ref_fields == keys or keys.issubset(ref_fields)
+        ), f"Fields mismatch: {ref_fields} != {keys}. Extra keys found: {extra_keys} \n \t when constructing {type(ref_instance)} with {keys}"
+
+        resolved_kwargs: Dict[str, Any] = {}
+        for f in keys:
+            resolved_kwargs[f] = config_from_dict(getattr(ref_instance, f), kwargs[f])
+        try:
+            new_instance = type(ref_instance)(**resolved_kwargs)
+        except Exception as e:
+            logger.error(f"Error when constructing {type(ref_instance)} with {resolved_kwargs}")
+            logger.error(e)
+            raise e
+        return new_instance
+
+
+def override_config_with_opts(config: Any, opts: Optional[List[str]] = None) -> Any:
+    """
+    Override the config with the opts.
+
+    Args:
+        config (Any): The config object.
+        opts (Dict[str, Any]): Dict for the overrides.
+
+    Returns:
+        Any: The config object.
+    """
+
+    # Convert Config object to a DictConfig object for Hydra
+
+    if not isinstance(config, DictConfig):
+        config_dict = attrs.asdict(config)
+        config_dict = DictConfig(content=config_dict, flags={"allow_objects": True})
+    else:
+        config_dict = config
+
+    # Use Hydra to handle overrides
+    cs = ConfigStore.instance()
+    cs.store(name="config", node=config_dict)
+
+    if opts is None:
+        opts = []
+
+    if len(opts) > 0 and opts[0] != "-":
+        raise ValueError(f"opts must start with '-' to separate from other arguments. Got: {opts}")
+    opts = opts[1:]
+    with initialize(version_base=None):
+        try:
+            cfg = compose(config_name="config", overrides=opts)
+        except Exception as e:
+            raise ValueError(f"Failed to compose config with opts: {e}")
+
+        OmegaConf.resolve(cfg)
+
+    config = config_from_dict(config, cfg)
+
+    return config
+
+
+def is_attrs_or_dataclass(obj) -> bool:
+    """
+    Check if the object is an instance of an attrs class or a dataclass.
+
+    Args:
+        obj: The object to check.
+
+    Returns:
+        bool: True if the object is an instance of an attrs class or a dataclass, False otherwise.
+    """
+    return is_dataclass(obj) or attr.has(type(obj))
+
+
+def get_fields(obj):
+    """
+    Get the fields of an attrs class or a dataclass.
+
+    Args:
+        obj: The object to get fields from. Must be an instance of an attrs class or a dataclass.
+
+    Returns:
+        list: A list of field names.
+
+    Raises:
+        ValueError: If the object is neither an attrs class nor a dataclass.
+    """
+    if is_dataclass(obj):
+        return [field.name for field in dataclass_fields(obj)]
+    elif attr.has(type(obj)):
+        return [field.name for field in attr.fields(type(obj))]
+    else:
+        raise ValueError("The object is neither an attrs class nor a dataclass.")
+
+
+def serialize_config(
+    config: Any,
+    return_type: str = "dict",
+    path: str | bytes | None = None,
+    filename: str = "config.yaml",
+    include_defaults: bool = False,
+) -> Dict[str, Any] | str:
+    """
+    Serialize a config (BaseConfig or DictConfig) to various formats.
+
+    Args:
+        config: The config to serialize (BaseConfig attrs object or DictConfig).
+        return_type: Output format - "dict" (plain dict), "yaml" (YAML string), or "file" (save to file).
+        path: Directory path to save the file. Required if return_type is "file".
+        filename: Name of the file to save. Only used if return_type is "file".
+        include_defaults: If True, add default parameter values from _target_ classes.
+
+    Returns:
+        Dict[str, Any] if return_type is "dict"
+        str (YAML) if return_type is "yaml" or "file"
+
+    Raises:
+        ValueError: If return_type is "file" but path is not provided.
+    """
+    if return_type == "file" and path is None:
+        raise ValueError("path must be provided when return_type is 'file'")
+
+    # Deep copy to avoid modifying original
+    config = deepcopy(config)
+
+    # Normalize to DictConfig with object support
+    if not isinstance(config, DictConfig):
+        config_dict = attrs.asdict(config)
+        config_omegaconf = DictConfig(content=config_dict, flags={"allow_objects": True})
+    else:
+        config_omegaconf = config
+
+    def is_serializable(item) -> bool:
+        try:
+            OmegaConf.to_yaml(item)
+            return True
+        except Exception:
+            return False
+
+    def get_default_params(cls_or_func):
+        if callable(cls_or_func):
+            signature = inspect.signature(cls_or_func)
+        else:
+            signature = inspect.signature(cls_or_func.__init__)
+        params = signature.parameters
+        return {name: param.default for name, param in params.items() if param.default is not inspect.Parameter.empty}
+
+    def process_config(conf):
+        if isinstance(conf, DictConfig):
+            for key, value in conf.items():
+                if isinstance(value, (DictConfig, ListConfig)):
+                    # Optionally add default params from _target_ classes
+                    if include_defaults:
+                        try:
+                            if "_target_" in value:
+                                default_params = get_default_params(value["_target_"])
+                                for default_key, default_v in default_params.items():
+                                    if default_key not in value:
+                                        value[default_key] = default_v
+                        except Exception as e:
+                            logger.error(f"Failed to add default argument values: {e}")
+                    process_config(value)
+                else:
+                    if not is_serializable(value) and value is not None:
+                        conf[key] = str(value)
+        elif isinstance(conf, ListConfig):
+            for i, item in enumerate(conf):
+                if isinstance(item, (DictConfig, ListConfig)):
+                    process_config(item)
+                else:
+                    if not is_serializable(item) and item is not None:
+                        conf[i] = str(item)
+        else:
+            raise NotImplementedError("Input config must be a DictConfig or ListConfig.")
+        return conf
+
+    config_omegaconf = process_config(config_omegaconf)
+    result_dict: Dict[str, Any] = OmegaConf.to_container(config_omegaconf, resolve=True)  # type: ignore
+
+    if return_type == "dict":
+        return result_dict
+
+    # For yaml and file, convert to YAML string
+    yaml_str = yaml.dump(result_dict, default_flow_style=False, sort_keys=True)
+
+    if return_type == "file":
+        os.makedirs(path, exist_ok=True)  # type: ignore
+        with open(f"{path}/{filename}", "w") as f:
+            f.write(yaml_str)
+        logger.info(f"Config is saved at {path}/{filename}")
+
+    return yaml_str

lipforcing/configs/discriminator.py

@@ -0,0 +1,23 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from omegaconf import DictConfig
+
+from lipforcing.utils import LazyCall as L
+from lipforcing.networks.discriminators import Discriminator_VideoDiT
+
+# 1.3B patchify: spatial-2, temporal-1; inner_dim=1536; layer=30
+Discriminator_Wan_1_3B_Config: DictConfig = L(Discriminator_VideoDiT)(
+    feature_indices=None,
+    num_blocks=30,
+    disc_type="dit_simple_conv3d",
+    inner_dim=1536 // 4,
+)
+
+# 14B patchify: spatial-2, temporal-1; inner_dim=5120; layer=40
+Discriminator_Wan_14B_Config: DictConfig = L(Discriminator_VideoDiT)(
+    feature_indices=None,
+    num_blocks=40,
+    disc_type="dit_simple_conv3d",
+    inner_dim=5120 // 4,
+)

lipforcing/configs/experiments/OmniAvatar/config_df.py

@@ -0,0 +1,247 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""DF (shift=5) for the 14B causal student — LoRA + selective unfreeze + t769 schedule.
+
+Flattened, self-contained Diffusion-Forcing config for the 14B causal
+student (shift=5, LoRA + selective unfreeze, t769 schedule).
+
+Effect: the DF student is only ever trained at the noise levels used at
+inference (input on step 1 = t=0.999, input on step 2 = t=0.769),
+avoiding a train/test schedule mismatch.
+
+Combines:
+
+1) 14B LoRA + selective unfreeze student:
+   - model_size="14B" student
+   - pretrained 14B student checkpoint (set via OMNIAVATAR_STUDENT_CKPT_14B)
+   - merge_lora=False (PEFT injects LoRA on transformer blocks)
+   - unfreeze_modules on the audio path + patch embedding
+   - lora_rank=128, lora_alpha=64
+   - FSDP + bf16 fwd / fp32 master+optim
+
+2) t769 schedule narrowing:
+   - sample_t_cfg.t_list = [0.999, 0.769, 0.0]
+   - student_sample_steps = 2
+"""
+
+import os
+
+
+from lipforcing.utils import LazyCall as L
+
+# Methods-level config (defines the attrs Config/ModelConfig classes and the
+# default create_config()). We keep importing this — it is NOT an experiment
+# chain file.
+import lipforcing.configs.methods.config_omniavatar_df as config_df_default
+
+from lipforcing.networks.OmniAvatar.network_causal import CausalOmniAvatarWan
+from lipforcing.datasets.omniavatar_dataloader import OmniAvatarDataLoader, create_omniavatar_dataloader
+
+
+# ---- Paths (override via env vars) ----
+OMNIAVATAR_ROOT = os.getenv("OMNIAVATAR_ROOT", "./OmniAvatar")
+DATA_ROOT = os.getenv("OMNIAVATAR_DATA_ROOT", "./data/v2v_training_data")
+# Pretrained 1.3B student checkpoint (used by the base network specs below).
+# The 14B release loads STUDENT_CKPT_14B instead; this is kept so the 1.3B
+# student can be enabled later without restructuring the config.
+STUDENT_CKPT_1_3B = os.getenv(
+    "OMNIAVATAR_STUDENT_CKPT_1_3B",
+    "/path/to/omniavatar_1.3b.pt",
+)
+DATA_LIST = os.getenv("OMNIAVATAR_DATA_LIST", f"{DATA_ROOT}/train_list.txt")
+VAL_LIST = os.getenv("OMNIAVATAR_VAL_LIST", f"{DATA_ROOT}/val_list.txt")
+MASK_PATH = os.getenv(
+    "MASK_PATH",
+    "/path/to/mask.png",
+)
+VAE_PATH = os.getenv(
+    "OMNIAVATAR_VAE_PATH",
+    os.path.join(OMNIAVATAR_ROOT, "pretrained_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth"),
+)
+
+# Pretrained 14B student checkpoint. Override via OMNIAVATAR_STUDENT_CKPT_14B
+# to use a different checkpoint.
+STUDENT_CKPT_14B = os.getenv(
+    "OMNIAVATAR_STUDENT_CKPT_14B",
+    "/path/to/omniavatar_14b_init.pt",
+)
+
+# Directory holding the base Wan2.1-T2V-14B diffusion safetensors shards.
+WAN_BASE_DIR = os.getenv(
+    "OMNIAVATAR_WAN_BASE", f"{OMNIAVATAR_ROOT}/pretrained_models/Wan2.1-T2V-14B"
+)
+WAN_14B_BASE = ",".join(
+    f"{WAN_BASE_DIR}/diffusion_pytorch_model-{i:05d}-of-00006.safetensors"
+    for i in range(1, 7)
+)
+
+# Submodules to keep fully trainable alongside LoRA on the transformer blocks.
+# Paths are dotted, relative to the CausalOmniAvatarWan instance (so they
+# include the "_core." prefix where the actual modules live).
+DEFAULT_UNFREEZE_MODULES = [
+    "_core.audio_proj",
+    "_core.audio_cond_projs",
+    "_core.patch_embedding",
+]
+
+
+# ---- Student network config (1.3B base — overridden to 14B below) ----
+CausalOmniAvatar_V2V_1_3B_Config: dict = L(CausalOmniAvatarWan)(
+    model_size="1.3B",
+    in_dim=65,
+    mode="v2v",
+    use_audio=True,
+    audio_hidden_size=32,
+    chunk_size=3,
+    total_num_frames=21,
+    base_model_paths=f"{OMNIAVATAR_ROOT}/pretrained_models/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors",
+    omniavatar_ckpt_path=STUDENT_CKPT_1_3B,
+    net_pred_type="flow",
+    schedule_type="rf",
+    use_dynamic_rope=False,
+    stochastic_attn_configs=[
+        {"local_attn_size": 7,  "sink_size": 1, "weight": 0.2},   # sink=1, window=6
+        {"local_attn_size": 10, "sink_size": 1, "weight": 0.2},   # sink=1, window=9
+        {"local_attn_size": 13, "sink_size": 1, "weight": 0.2},   # sink=1, window=12
+        {"local_attn_size": 9,  "sink_size": 3, "weight": 0.2},   # sink=3, window=6
+        {"local_attn_size": 12, "sink_size": 3, "weight": 0.2},   # sink=3, window=9
+    ],
+)
+
+
+def create_config():
+    # ================================================================== #
+    # Base: methods-level OmniAvatar Diffusion Forcing config             #
+    # ================================================================== #
+    config = config_df_default.create_config()
+
+    # ================================================================== #
+    # DF (shift=5) overrides: LR, precision                               #
+    # ================================================================== #
+    # Learning rate
+    config.model.net_optimizer.lr = 1e-5
+
+    # Precision
+    config.model.precision = "bfloat16"
+    config.model.precision_fsdp = "float32"
+
+    # Input shape: 512x512 @ 81 frames -> latent [16, 21, 64, 64]
+    config.model.input_shape = [16, 21, 64, 64]
+
+    # Student network
+    config.model.net = CausalOmniAvatar_V2V_1_3B_Config
+    config.model.net.total_num_frames = config.model.input_shape[1]
+
+    # Timestep schedule — shift=5.0 matches OmniAvatar's default scheduler
+    config.model.sample_t_cfg.time_dist_type = "shifted"
+    config.model.sample_t_cfg.shift = 5.0
+    config.model.sample_t_cfg.min_t = 0.001
+    config.model.sample_t_cfg.max_t = 0.999
+    config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0]
+
+    # Diffusion forcing settings
+    config.model.student_sample_steps = 4
+
+    # Dataloader
+    config.dataloader_train = L(OmniAvatarDataLoader)(
+        data_list_path=DATA_LIST,
+        latentsync_mask_path=MASK_PATH,
+        batch_size=2,
+        num_workers=4,
+        neg_text_emb_path=os.getenv("NEG_TEXT_EMB_PATH", None),
+        use_ref_sequence=True,
+        load_ode_path=False,
+    )
+
+    # Validation dataloader — 10 fixed samples, finite iterator, batch_size=1
+    config.dataloader_val = L(create_omniavatar_dataloader)(
+        data_list_path=VAL_LIST,
+        latentsync_mask_path=MASK_PATH,
+        batch_size=1,
+        num_workers=2,
+        neg_text_emb_path=os.getenv("NEG_TEXT_EMB_PATH", None),
+        use_ref_sequence=True,
+        load_ode_path=False,
+    )
+
+    # VAE for visual logging (decodes latents to video for wandb)
+    config.model.vae_path = VAE_PATH
+
+    # ---- Optional: on-the-fly preprocessing (opt-in via env) ----
+    # When OMNIAVATAR_ON_THE_FLY=1, samples lacking precomputed .pt are encoded
+    # from raw (sub_clip.mp4 + audio.wav + prompt.txt) by frozen encoders in the
+    # MAIN training process; encoded tensors are optionally cached for reuse. The
+    # precompute fast path is unchanged, and with the env unset this block is a
+    # no-op so the precompute-based data path is used unchanged.
+    if os.getenv("OMNIAVATAR_ON_THE_FLY", "0") == "1":
+        config.dataloader_train.on_the_fly = True
+        config.dataloader_train.cache_encoded = os.getenv("OMNIAVATAR_CACHE_ENCODED", "1") == "1"
+        config.dataloader_train.cache_dir = os.getenv("OMNIAVATAR_CACHE_DIR", None)
+        config.dataloader_train.vae_path = VAE_PATH
+        config.dataloader_train.wav2vec_path = os.getenv("OMNIAVATAR_WAV2VEC_PATH", None)
+        config.dataloader_train.text_encoder_path = os.getenv("OMNIAVATAR_TEXT_ENCODER_PATH", None)
+
+    # Training (5K steps — matches the released DF run and paper App. D.1)
+    config.trainer.max_iter = 5000
+    config.trainer.logging_iter = 1
+    config.trainer.save_ckpt_iter = 500
+    config.trainer.validation_iter = 500
+    config.trainer.skip_initial_validation = True
+    config.trainer.callbacks.wandb.sample_logging_iter = 500
+
+    config.log_config.group = "stage1_df"
+    config.log_config.name = "df_14b_lora_t769"
+
+    # ================================================================== #
+    # Switch student to 14B + FSDP                                        #
+    # ================================================================== #
+    # ---- Switch student to 14B (omit this block for the 1.3B base student) ----
+    config.model.net.model_size = "14B"
+    config.model.net.base_model_paths = WAN_14B_BASE
+    config.model.net.omniavatar_ckpt_path = STUDENT_CKPT_14B
+    # The 14B teacher uses merge_lora=True to fuse the adapter into the base
+    # before training. Mirror it for the 14B student so DF starts from the
+    # fused state, not a base+LoRA stacked state.
+    config.model.net.merge_lora = True
+
+    # ---- DDP -> FSDP ----
+    config.trainer.ddp = False
+    config.trainer.fsdp = True
+    config.trainer.fsdp_min_num_params = int(1e8)
+    config.trainer.fsdp_cpu_offload = False
+    config.trainer.fsdp_sharding_group_size = None  # default = world_size
+    # Mixed-precision FSDP: bf16 fwd/bwd, fp32 master+optim.
+    config.model.precision = "bfloat16"
+    config.model.precision_fsdp = "float32"
+    # Meta-init disabled (RoPE Python-attr issue at 14B DF).
+    config.model.fsdp_meta_init = False
+
+    # ---- Effective batch 16 = 2/GPU * 4 GPUs * grad_accum 2 (matches train_stage1_df.sh) ----
+    config.trainer.grad_accum_rounds = 2
+
+    # ================================================================== #
+    # Full fine-tune -> LoRA + selective unfreeze                         #
+    # ================================================================== #
+    # ---- Switch from full FT to LoRA + selective unfreeze ----
+    config.model.net.merge_lora = False
+    config.model.net.unfreeze_modules = DEFAULT_UNFREEZE_MODULES
+
+    # LoRA hyperparameters. Match the V2V adapter we're loading from
+    # (rank=128, alpha=64).
+    config.model.net.lora_rank = 128
+    config.model.net.lora_alpha = 64
+
+    # ================================================================== #
+    # t769 2-step schedule                                                #
+    # ================================================================== #
+    # 2-step schedule overrides. With student_sample_steps=2, the student
+    # is trained on the two intervals (0.999 -> 0.769) and (0.769 -> 0.0)
+    # — the exact intervals SF t769 inference uses.
+    config.model.sample_t_cfg.t_list = [0.999, 0.769, 0.0]
+    config.model.student_sample_steps = 2
+
+    return config
+
+
+config = create_config()

lipforcing/configs/experiments/OmniAvatar/config_sf.py

@@ -0,0 +1,395 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""SF (Re-DMD beta=2 + TAEW) for the 14B causal student + 14B fake_score,
+LoRA + selective unfreeze, t769 schedule, matched critic LR 2e-6.
+
+Self-contained Self-Forcing experiment config for the 14B causal student +
+14B critic (Re-DMD, beta=2, TAEW decoder, t769 2-step schedule).
+
+Combines:
+- 14B student (causal) + 14B fake_score (bidirectional), both LoRA +
+  selective unfreeze, initialized from a pretrained 14B audio adapter checkpoint.
+- 14B teacher = OmniAvatar-LS (the lip-sync-finetuned OmniAvatar teacher,
+  paper App. B.2; bidirectional, frozen, merge_lora=True).
+- Sliding-window attention: sink=1 + window=7, dynamic RoPE.
+- Re-DMD reward path (SyncNet-v2 sync-C), beta=2, TAEW decoder.
+- t769 2-step schedule: t_list=[0.999, 0.769, 0.0], student_sample_steps=2.
+- Effective batch 16 = 1/GPU * 4 GPUs * grad_accum=4; FSDP bf16/fp32.
+- Critic (fake_score) LR matched to the student at 2e-6 (paper D.1).
+"""
+
+import os
+
+
+from lipforcing.utils import LazyCall as L
+
+# Methods-level configs (define the attrs Config/ModelConfig classes, network
+# classes, RewardConfig, and the Re-DMD model class). These are NOT experiment
+# chain files and are kept.
+import lipforcing.configs.methods.config_omniavatar_sf as config_sf_default
+from lipforcing.configs.methods.config_omniavatar_sf import RewardConfig
+
+from lipforcing.networks.OmniAvatar.network import OmniAvatarWan
+from lipforcing.networks.OmniAvatar.network_causal import CausalOmniAvatarWan
+from lipforcing.datasets.omniavatar_dataloader import OmniAvatarDataLoader, create_omniavatar_dataloader
+from lipforcing.methods.omniavatar_self_forcing_re_dmd import OmniAvatarSelfForcingReDMD
+
+
+# ---- Paths (override via CLI or env) ----
+OMNIAVATAR_ROOT = os.getenv("OMNIAVATAR_ROOT", "./OmniAvatar")
+DATA_ROOT = os.getenv("OMNIAVATAR_DATA_ROOT", "./data/v2v_training_data")
+TEACHER_CKPT = os.getenv(
+    "OMNIAVATAR_TEACHER_CKPT",
+    "/path/to/omniavatar_teacher.pt",
+)
+# Pretrained 1.3B student checkpoint (used by the base network specs below).
+# The 14B release loads STUDENT_CKPT_14B instead; this is kept so the 1.3B
+# student can be enabled later without restructuring the config.
+STUDENT_CKPT_1_3B = os.getenv(
+    "OMNIAVATAR_STUDENT_CKPT_1_3B",
+    "/path/to/omniavatar_1.3b.pt",
+)
+
+# Initial 14B checkpoint (same as the teacher): provides initial LoRA values
+# plus audio/patch-embedding weights for both the student and the fake_score
+# before training.
+STUDENT_CKPT_14B = os.getenv(
+    "OMNIAVATAR_STUDENT_CKPT_14B",
+    "/path/to/omniavatar_14b_init.pt",
+)
+
+# Directory holding the base Wan2.1-T2V-14B diffusion safetensors shards.
+WAN_BASE_DIR = os.getenv(
+    "OMNIAVATAR_WAN_BASE", f"{OMNIAVATAR_ROOT}/pretrained_models/Wan2.1-T2V-14B"
+)
+WAN_14B_BASE = ",".join(
+    f"{WAN_BASE_DIR}/diffusion_pytorch_model-{i:05d}-of-00006.safetensors"
+    for i in range(1, 7)
+)
+
+# Submodule paths (relative to each network) to keep fully trainable
+# alongside the LoRA A/B matrices on the transformer blocks. Different
+# prefix between the causal and bidirectional classes:
+#   - CausalOmniAvatarWan (student): WanModel lives at self._core
+#   - OmniAvatarWan (fake_score):    WanModel lives at self.model
+STUDENT_UNFREEZE = [
+    "_core.audio_proj",
+    "_core.audio_cond_projs",
+    "_core.patch_embedding",
+]
+FAKE_SCORE_UNFREEZE = [
+    "model.audio_proj",
+    "model.audio_cond_projs",
+    "model.patch_embedding",
+]
+
+# Reward checkpoints: SyncNet scorer + TAEW decoder for the reward-path pixel decode.
+SYNCNET_CKPT = os.getenv("SYNCNET_CKPT", "/path/to/syncnet_v2.model")
+TAEW_CKPT = os.getenv("TAEW_CKPT", "/path/to/taew2_1.pth")
+
+# Critic (fake_score) LR: matched to the student LR (paper D.1).
+CRITIC_LR = 2e-6
+RUN_NAME = "sf_14b_lora_t769"
+
+
+# ---- Network configs ----
+OmniAvatar_V2V_14B_Teacher: dict = L(OmniAvatarWan)(
+    model_size="14B",
+    in_dim=65,
+    mode="v2v",
+    use_audio=True,
+    audio_hidden_size=32,
+    base_model_paths=WAN_14B_BASE,
+    omniavatar_ckpt_path=TEACHER_CKPT,
+    merge_lora=True,
+    net_pred_type="flow",
+    schedule_type="rf",
+)
+
+OmniAvatar_V2V_1_3B_FakeScore: dict = L(OmniAvatarWan)(
+    model_size="1.3B",
+    in_dim=65,
+    mode="v2v",
+    use_audio=True,
+    audio_hidden_size=32,
+    base_model_paths=f"{OMNIAVATAR_ROOT}/pretrained_models/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors",
+    omniavatar_ckpt_path=STUDENT_CKPT_1_3B,
+    merge_lora=False,  # Fake score is trainable, keep LoRA separate
+    net_pred_type="flow",
+    schedule_type="rf",
+)
+
+CausalOmniAvatar_V2V_1_3B_Student: dict = L(CausalOmniAvatarWan)(
+    model_size="1.3B",
+    in_dim=65,
+    mode="v2v",
+    use_audio=True,
+    audio_hidden_size=32,
+    chunk_size=3,
+    total_num_frames=21,
+    base_model_paths=f"{OMNIAVATAR_ROOT}/pretrained_models/Wan2.1-T2V-1.3B/diffusion_pytorch_model.safetensors",
+    omniavatar_ckpt_path=STUDENT_CKPT_1_3B,
+    net_pred_type="flow",
+    schedule_type="rf",
+    # Sliding window attention for AR rollout.
+    # Defaults: full causal (no window constraint during SF).
+    local_attn_size=-1,
+    sink_size=0,
+    use_dynamic_rope=True,
+)
+
+
+def create_config():
+    # ================================================================== #
+    # Base: methods-level OmniAvatar Self-Forcing config                  #
+    # ================================================================== #
+    config = config_sf_default.create_config()
+
+    # ================================================================== #
+    # Self-Forcing overrides: LR, optimizer, FSDP                         #
+    # ================================================================== #
+    # Learning rates and optimizer (Adam beta1=0.0, as used by Self-Forcing)
+    config.model.net_optimizer.lr = 2e-6
+    config.model.net_optimizer.betas = (0.0, 0.999)
+    config.model.fake_score_optimizer.lr = 2e-6
+    config.model.fake_score_optimizer.betas = (0.0, 0.999)
+
+    # Multi-GPU: FSDP required (DDP OOMs on student update at ~79GB/GPU)
+    config.trainer.fsdp = True
+
+    # Precision
+    config.model.precision = "bfloat16"
+    config.model.precision_fsdp = "float32"
+
+    # Input shape: 512x512 @ 81 frames -> latent [16, 21, 64, 64]
+    config.model.input_shape = [16, 21, 64, 64]
+    config.model.fake_score_pred_type = "x0"
+    config.model.guidance_scale = 4.5
+    # Sync-Window DMD (SW-DMD; paper Sec. 4.3 / Eq. 6): teacher CFG is gated to the
+    # sync window. t_lo/t_hi = the shifted-timestep band [0.556, 0.882] = teacher ODE
+    # steps j in [20, 40] (the sync-favoring band from the trajectory analysis).
+    config.model.sync_window_cfg.enabled = False  # base default; enabled in the sliding-window block below
+    config.model.sync_window_cfg.t_lo = 0.556
+    config.model.sync_window_cfg.t_hi = 0.882
+
+    # Networks (base specs): 14B teacher + 1.3B student + 1.3B fake_score.
+    # net and fake_score are switched to 14B in the 14B block below.
+    config.model.net = CausalOmniAvatar_V2V_1_3B_Student
+    config.model.net.total_num_frames = config.model.input_shape[1]
+    config.model.teacher = OmniAvatar_V2V_14B_Teacher  # OmniAvatar-LS (paper App. B.2)
+    config.model.fake_score_net = OmniAvatar_V2V_1_3B_FakeScore
+
+    # GAN disabled by default.
+    config.model.gan_loss_weight_gen = 0
+    config.model.student_update_freq = 5  # 1:5 ratio
+
+    # Student weights: do NOT copy 14B teacher weights onto 1.3B student.
+    config.model.load_student_weights = False
+    # Load the DF-initialized student from the Stage-1 checkpoint. This DF init
+    # must match the sliding-window (sink=1, window=7) architecture set below.
+    config.trainer.checkpointer.pretrained_ckpt_path = os.getenv(
+        "OMNIAVATAR_DF_CKPT",
+        "/path/to/df_init_checkpoint.pth",
+    )
+    config.trainer.checkpointer.pretrained_ckpt_key_map = {"net": "net"}
+
+    # Timestep schedule — shift=5.0 matches OmniAvatar's inference scheduler
+    config.model.sample_t_cfg.time_dist_type = "shifted"
+    config.model.sample_t_cfg.shift = 5.0
+    config.model.sample_t_cfg.min_t = 0.001
+    config.model.sample_t_cfg.max_t = 0.999
+    config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0]
+
+    # Self-Forcing specific
+    config.model.enable_gradient_in_rollout = True
+    config.model.start_gradient_frame = 0
+    config.model.same_step_across_blocks = True
+    config.model.context_noise = 0.0
+
+    # Dataloader (OmniAvatarDataLoader provides infinite iteration)
+    config.dataloader_train = L(OmniAvatarDataLoader)(
+        data_list_path=os.getenv("OMNIAVATAR_DATA_LIST", f"{DATA_ROOT}/train_list.txt"),
+        latentsync_mask_path=os.getenv(
+            "MASK_PATH",
+            "/path/to/mask.png",
+        ),
+        batch_size=8,
+        num_workers=2,
+        neg_text_emb_path=os.getenv("NEG_TEXT_EMB_PATH", None),
+        use_ref_sequence=True,
+    )
+
+    # Validation dataloader — 10 fixed samples, finite iterator, batch_size=1
+    VAL_LIST = os.getenv("OMNIAVATAR_VAL_LIST", f"{DATA_ROOT}/val_list.txt")
+    VAE_PATH = os.getenv(
+        "OMNIAVATAR_VAE_PATH",
+        os.path.join(OMNIAVATAR_ROOT, "pretrained_models/Wan2.1-T2V-1.3B/Wan2.1_VAE.pth"),
+    )
+    config.dataloader_val = L(create_omniavatar_dataloader)(
+        data_list_path=VAL_LIST,
+        latentsync_mask_path=os.getenv(
+            "MASK_PATH",
+            "/path/to/mask.png",
+        ),
+        batch_size=1,
+        num_workers=2,
+        neg_text_emb_path=os.getenv("NEG_TEXT_EMB_PATH", None),
+        use_ref_sequence=True,
+        load_ode_path=False,
+    )
+    config.model.vae_path = VAE_PATH
+
+    # ---- Optional: on-the-fly preprocessing (opt-in via env) ----
+    # When OMNIAVATAR_ON_THE_FLY=1, samples lacking precomputed .pt are encoded
+    # from raw (sub_clip.mp4 + audio.wav + prompt.txt) by frozen encoders in the
+    # MAIN training process; encoded tensors are optionally cached for reuse. The
+    # precompute fast path is used by default, and with the env unset this block
+    # is a no-op.
+    if os.getenv("OMNIAVATAR_ON_THE_FLY", "0") == "1":
+        config.dataloader_train.on_the_fly = True
+        config.dataloader_train.cache_encoded = os.getenv("OMNIAVATAR_CACHE_ENCODED", "1") == "1"
+        config.dataloader_train.cache_dir = os.getenv("OMNIAVATAR_CACHE_DIR", None)
+        config.dataloader_train.vae_path = VAE_PATH
+        config.dataloader_train.wav2vec_path = os.getenv("OMNIAVATAR_WAV2VEC_PATH", None)
+        config.dataloader_train.text_encoder_path = os.getenv("OMNIAVATAR_TEXT_ENCODER_PATH", None)
+
+    # Training — 1.3B base: bs=8, grad_accum=2 -> eff batch 64 (14B override below = eff 16)
+    config.trainer.grad_accum_rounds = 2
+    config.trainer.max_iter = 600  # released 14B checkpoint = step 600 (paper App. D.1)
+    config.trainer.logging_iter = 1
+    config.trainer.save_ckpt_iter = 100
+    config.trainer.validation_iter = 100
+    config.trainer.skip_initial_validation = True
+
+    # Wandb sample logging (video generation) every 100 steps
+    config.trainer.callbacks.wandb.sample_logging_iter = 100
+    config.trainer.callbacks.wandb.fps = 25  # OmniAvatar is 25 fps
+    config.trainer.callbacks.wandb.syncnet_checkpoint_path = SYNCNET_CKPT
+
+    config.log_config.group = "stage2_sf"
+    config.log_config.wandb_entity = ""
+
+    # ================================================================== #
+    # Sliding-window attention (sink=1, window=7)                         #
+    # ================================================================== #
+    # Sliding window: 1 sink + 6 rolling = 7 total visible frames
+    config.model.net.local_attn_size = 7
+    config.model.net.sink_size = 1
+    config.model.net.use_dynamic_rope = True
+
+    # 2-step distillation: t_list[0] and t_list[2] from the 4-step schedule
+    config.model.sample_t_cfg.t_list = [0.999, 0.833, 0.0]
+    config.model.student_sample_steps = 2
+
+    # Enable Sync-Window DMD (SW-DMD) for this experiment
+    config.model.sync_window_cfg.enabled = True
+
+    # ================================================================== #
+    # Re-DMD reward weighting (sync-C, beta=2)                            #
+    # ================================================================== #
+    # Switch model class to the Re-DMD variant.
+    config.model_class._target_ = OmniAvatarSelfForcingReDMD
+
+    # Reward sub-config (RewardConfig attrs class survives OmegaConf serialization)
+    config.model.reward = RewardConfig(
+        enabled=True,
+        checkpoint_path=SYNCNET_CKPT,
+        input_fps=25.0,
+        audio_sample_rate=16000,
+        vshift=15,
+    )
+
+    # Top-level reward knobs (read by _apply_reward_weighting)
+    config.model.reward_beta = 2
+    config.model.center_reward = False
+    config.model.clamp_reward = None
+
+    # VAE path (required for reward decode in _decode_gen_to_pixels)
+    assert getattr(config.model, "vae_path", "") != "", (
+        "config.model.vae_path must be set for Re-DMD — the reward path VAE-decodes "
+        "the generator output to pixels. Either set OMNIAVATAR_VAE_PATH env or "
+        "ensure the default OmniAvatar install path is present."
+    )
+
+    # Data: raw waveform loading (required for audio_waveform in batch)
+    config.dataloader_train.load_raw_audio = True
+    config.dataloader_train.raw_audio_sample_rate = 16000
+    config.dataloader_train.raw_audio_num_frames = 81
+    config.dataloader_train.raw_audio_fps = 25.0
+
+    # ================================================================== #
+    # TAEW decoder for reward-path pixel decode                           #
+    # ================================================================== #
+    config.model.reward.decoder_kind = "taew"
+    config.model.reward.taew_checkpoint_path = TAEW_CKPT
+
+    # ================================================================== #
+    # 14B student + critic: LoRA, t769 schedule                           #
+    # ================================================================== #
+    # ---- Student: 1.3B causal -> 14B causal LoRA (omit to keep the 1.3B student) ----
+    config.model.net.model_size = "14B"
+    config.model.net.base_model_paths = WAN_14B_BASE
+    config.model.net.omniavatar_ckpt_path = STUDENT_CKPT_14B
+    config.model.net.merge_lora = False
+    config.model.net.unfreeze_modules = STUDENT_UNFREEZE
+    config.model.net.lora_rank = 128
+    config.model.net.lora_alpha = 64
+
+    # ---- Fake_score: 1.3B bidirectional -> 14B bidirectional LoRA ----
+    config.model.fake_score_net.model_size = "14B"
+    config.model.fake_score_net.base_model_paths = WAN_14B_BASE
+    config.model.fake_score_net.omniavatar_ckpt_path = STUDENT_CKPT_14B
+    config.model.fake_score_net.merge_lora = False
+    config.model.fake_score_net.unfreeze_modules = FAKE_SCORE_UNFREEZE
+    config.model.fake_score_net.lora_rank = 128
+    config.model.fake_score_net.lora_alpha = 64
+
+    # ---- Teacher: stays 14B + merge_lora=True (frozen, full state) ----
+    # No changes.
+
+    # ---- Schedule: t769 ----
+    config.model.sample_t_cfg.t_list = [0.999, 0.769, 0.0]
+    config.model.student_sample_steps = 2
+
+    # ---- Effective batch 16 = 1/GPU * 4 GPUs * grad_accum=4 ----
+    config.dataloader_train.batch_size = 1
+    config.trainer.grad_accum_rounds = 4
+
+    # ---- FSDP knobs (mirror 14B DF LoRA setup) ----
+    config.trainer.ddp = False
+    config.trainer.fsdp = True
+    config.trainer.fsdp_min_num_params = int(1e8)
+    config.trainer.fsdp_cpu_offload = False
+    config.trainer.fsdp_sharding_group_size = None
+    config.model.precision = "bfloat16"
+    config.model.precision_fsdp = "float32"
+    # Meta-init enabled (3-network 14B run); RoPE buffers are re-materialized via
+    # reset_parameters() on both network classes.
+    config.model.fsdp_meta_init = True
+
+    # ================================================================== #
+    # Critic LR = student LR (matched, 2e-6)                             #
+    # ================================================================== #
+    # Keep the student optimizer unchanged; reduce only the critic.
+    config.model.fake_score_optimizer.lr = CRITIC_LR
+
+    config.log_config.name = RUN_NAME
+
+    assert abs(config.model.net_optimizer.lr - 2e-6) < 1e-12, (
+        f"Expected student LR to stay at 2e-6, got {config.model.net_optimizer.lr}"
+    )
+    assert abs(config.model.fake_score_optimizer.lr - CRITIC_LR) < 1e-12, (
+        f"Expected critic LR {CRITIC_LR}, got {config.model.fake_score_optimizer.lr}"
+    )
+    assert config.model.sample_t_cfg.t_list == [0.999, 0.769, 0.0], (
+        f"Expected t769 schedule, got {config.model.sample_t_cfg.t_list}"
+    )
+    assert config.model.student_sample_steps == 2, (
+        f"Expected 2 student sample steps, got {config.model.student_sample_steps}"
+    )
+
+    return config
+
+
+config = create_config()

lipforcing/configs/methods/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

lipforcing/configs/methods/config_dmd2.py

@@ -0,0 +1,79 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import copy
+import attrs
+from omegaconf import DictConfig
+
+from lipforcing.utils import LazyCall as L
+from lipforcing.configs.config import (
+    BaseModelConfig,
+    BaseConfig,
+)
+from lipforcing.configs.opt import BaseOptimizerConfig, BaseSchedulerConfig
+from lipforcing.methods import DMD2Model
+from lipforcing.configs.callbacks import (
+    WANDB_CALLBACK,
+    GradClip_CALLBACK,
+    GPUStats_CALLBACK,
+    TrainProfiler_CALLBACK,
+    ParamCount_CALLBACK,
+)
+
+
+@attrs.define(slots=False)
+class ModelConfig(BaseModelConfig):
+    # optimizer and scheduler for the fake score net
+    fake_score_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseOptimizerConfig))
+    fake_score_scheduler: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseSchedulerConfig))
+
+    discriminator: DictConfig = None  # opt-in GAN (set e.g. Discriminator_Wan_14B_Config to enable)
+    # optimizer and scheduler for the discriminator
+    discriminator_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseOptimizerConfig))
+    discriminator_scheduler: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseSchedulerConfig))
+
+    # student update frequency
+    student_update_freq: int = 5
+
+    # weight for the GAN loss in the student update phase
+    gan_loss_weight_gen: float = 0.001
+
+    # use the same t and noise to perturb the real data and fake data
+    gan_use_same_t_noise: bool = False
+
+    # perform dsm loss in a space specified by fake_score_pred_type (None means use original net_pred_type)
+    fake_score_pred_type: str | None = None
+
+    # R1 regularization weight (0 means no R1 reg, recommended value when using R1 reg: 100-1000)
+    gan_r1_reg_weight: float = 0.0
+    # R1 regularization noise scale (it only takes effect when gan_r1_reg_weight > 0)
+    gan_r1_reg_alpha: float = 0.1
+
+
+@attrs.define(slots=False)
+class Config(BaseConfig):
+    model: ModelConfig = attrs.field(factory=ModelConfig)
+    model_class: DictConfig = L(DMD2Model)(
+        config=None,
+    )
+
+
+def create_config():
+    config = Config()
+    config.trainer.callbacks = DictConfig(
+        {
+            **GradClip_CALLBACK,
+            **GPUStats_CALLBACK,
+            **TrainProfiler_CALLBACK,
+            **ParamCount_CALLBACK,
+            **WANDB_CALLBACK,
+        }
+    )
+
+    config.dataloader_train.batch_size = 256
+    config.model.discriminator_scheduler.warm_up_steps = [0]
+    config.model.fake_score_scheduler.warm_up_steps = [0]
+    config.model.net_scheduler.warm_up_steps = [0]
+    config.model.sample_t_cfg.time_dist_type = "polynomial"
+
+    return config

lipforcing/configs/methods/config_omniavatar_df.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Method config for OmniAvatar Diffusion Forcing (Stage 1 alternative to ODE KD)."""
+
+import attrs
+from omegaconf import DictConfig
+
+from lipforcing.utils import LazyCall as L
+from lipforcing.configs.config import BaseConfig, BaseModelConfig
+from lipforcing.methods.omniavatar_diffusion_forcing import OmniAvatarDiffusionForcingModel
+from lipforcing.callbacks.wandb import WandbCallback
+from lipforcing.configs.callbacks import (
+    GradClip_CALLBACK,
+    ParamCount_CALLBACK,
+    TrainProfiler_CALLBACK,
+    GPUStats_CALLBACK,
+)
+
+
+@attrs.define(slots=False)
+class ModelConfig(BaseModelConfig):
+    context_noise: float = 0.0
+    vae_path: str = ""  # Path to WanVAE for visual logging (optional)
+
+
+@attrs.define(slots=False)
+class Config(BaseConfig):
+    model: ModelConfig = attrs.field(factory=ModelConfig)
+    model_class: DictConfig = L(OmniAvatarDiffusionForcingModel)(
+        config=None,
+    )
+
+
+def create_config():
+    config = Config()
+    # OmniAvatar uses 25fps video
+    OMNIAVATAR_WANDB = dict(wandb=L(WandbCallback)(sample_logging_iter=None, fps=25))
+    config.trainer.callbacks = DictConfig(
+        {
+            **GradClip_CALLBACK,
+            **GPUStats_CALLBACK,
+            **TrainProfiler_CALLBACK,
+            **ParamCount_CALLBACK,
+            **OMNIAVATAR_WANDB,
+        }
+    )
+
+    config.dataloader_train.batch_size = 1
+    config.model.student_sample_steps = 4
+    config.model.net_scheduler.warm_up_steps = [0]
+
+    return config

lipforcing/configs/methods/config_omniavatar_sf.py

@@ -0,0 +1,102 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Method config for OmniAvatar Self-Forcing distillation."""
+
+import attrs
+from omegaconf import DictConfig
+
+from lipforcing.utils import LazyCall as L
+from typing import Optional
+
+from lipforcing.configs.methods.config_self_forcing import (
+    Config as SFConfig,
+    ModelConfig as SFModelConfig,
+)
+from lipforcing.methods.omniavatar_self_forcing import OmniAvatarSelfForcingModel
+from lipforcing.configs.callbacks import (
+    WANDB_CALLBACK,
+    GradClip_CALLBACK,
+    ParamCount_CALLBACK,
+    TrainProfiler_CALLBACK,
+    GPUStats_CALLBACK,
+    EMA_CALLBACK,
+)
+
+
+@attrs.define(slots=False)
+class RewardConfig:
+    """Config for the Re-DMD reward scorer (SyncNet-v2 sync-C)."""
+    enabled: bool = True
+    checkpoint_path: str = ""
+    input_fps: float = 25.0
+    audio_sample_rate: int = 16000
+    vshift: int = 15
+    # Opt-in TAEW decoder. Default "vae" preserves WanVideoVAE.decode behavior.
+    # When "taew", the Re-DMD model loads a TAEHVDecoderWrapper from
+    # taew_checkpoint_path and uses it in place of self.net.vae for the
+    # reward-path pixel decode.
+    decoder_kind: str = "vae"
+    taew_checkpoint_path: str = ""
+
+
+@attrs.define(slots=False)
+class OmniAvatarModelConfig(SFModelConfig):
+    # Separate fake_score config (allows 1.3B fake_score with 14B teacher)
+    fake_score: Optional[DictConfig] = None
+
+    # Wan VAE path for decoding validation videos in the wandb callback.
+    # Logging-only; declared here so it survives config serialization.
+    vae_path: str = ""
+
+    # Re-DMD reward config. None = reward disabled (vanilla SF).
+    reward: Optional[RewardConfig] = None
+    reward_beta: float = 0.25
+    center_reward: bool = False
+    clamp_reward: Optional[list] = None
+
+    # How to combine per-sample reward weights with per-sample VSD losses.
+    #   - "per_sample" (default): mean_i(exp(beta*r_i) * L_i). Per-sample
+    #       coupling without normalization. Preserves reward magnitude as
+    #       a loss-scale signal.
+    #   - "self_normalized": sum_i(exp(beta*r_i) * L_i) / sum_j(exp(beta*r_j)).
+    #       Self-normalized importance sampling (paper's Z(c) partition).
+    #       Shift-invariant; loss bounded in [min(L), max(L)]. Better for
+    #       additive combination with GAN/other aux losses.
+    #   - "legacy_batch_mean": mean_i(exp(beta*r_i)) * mean_i(L_i). Batch-mean
+    #       form that decouples reward from loss at batch > 1; not recommended.
+    reward_weighting_mode: str = "per_sample"
+
+    # Diagnostic video dump at every generator step (rank 0 only).
+    save_reward_debug_video: bool = False
+    reward_debug_dir: str = "logs/redmd_debug"
+
+
+@attrs.define(slots=False)
+class Config(SFConfig):
+    model: OmniAvatarModelConfig = attrs.field(factory=OmniAvatarModelConfig)
+    model_class: DictConfig = L(OmniAvatarSelfForcingModel)(
+        config=None,
+    )
+
+
+def create_config():
+    config = Config()
+    config.trainer.callbacks = DictConfig(
+        {
+            **GradClip_CALLBACK,
+            **GPUStats_CALLBACK,
+            **TrainProfiler_CALLBACK,
+            **ParamCount_CALLBACK,
+            **EMA_CALLBACK,
+            **WANDB_CALLBACK,
+        }
+    )
+
+    config.dataloader_train.batch_size = 1
+    config.model.student_sample_steps = 4
+    config.model.discriminator_scheduler.warm_up_steps = [0]
+    config.model.fake_score_scheduler.warm_up_steps = [0]
+    config.model.net_scheduler.warm_up_steps = [0]
+
+    return config

lipforcing/configs/methods/config_self_forcing.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import attrs
+from omegaconf import DictConfig
+
+from lipforcing.utils import LazyCall as L
+from lipforcing.configs.methods.config_dmd2 import (
+    Config as DMD2Config,
+    ModelConfig as DMD2ModelConfig,
+)
+from lipforcing.methods import SelfForcingModel
+from lipforcing.configs.callbacks import (
+    WANDB_CALLBACK,
+    GradClip_CALLBACK,
+    ParamCount_CALLBACK,
+    TrainProfiler_CALLBACK,
+    GPUStats_CALLBACK,
+    EMA_CALLBACK,
+)
+
+
+@attrs.define(slots=False)
+class ModelConfig(DMD2ModelConfig):
+    enable_gradient_in_rollout: bool = True
+    start_gradient_frame: int = 0
+    same_step_across_blocks: bool = True
+    last_step_only: bool = False
+    context_noise: float = 0.0
+
+
+@attrs.define(slots=False)
+class Config(DMD2Config):
+    model: ModelConfig = attrs.field(factory=ModelConfig)
+    model_class: DictConfig = L(SelfForcingModel)(
+        config=None,
+    )
+
+
+def create_config():
+    config = Config()
+    config.trainer.callbacks = DictConfig(
+        {
+            **GradClip_CALLBACK,
+            **GPUStats_CALLBACK,
+            **TrainProfiler_CALLBACK,
+            **ParamCount_CALLBACK,
+            **EMA_CALLBACK,
+            **WANDB_CALLBACK,
+        }
+    )
+
+    config.dataloader_train.batch_size = 256
+    config.model.student_sample_steps = 4
+    config.model.discriminator_scheduler.warm_up_steps = [0]
+    config.model.fake_score_scheduler.warm_up_steps = [0]
+    config.model.net_scheduler.warm_up_steps = [0]
+
+    return config