Add Wan 2.2 S2V-14B — audio-driven talking-head video (MLX port) - #45
Open
KaedeTai wants to merge 16 commits into
Open
Add Wan 2.2 S2V-14B — audio-driven talking-head video (MLX port)#45KaedeTai wants to merge 16 commits into
KaedeTai wants to merge 16 commits into
Conversation
…PE (untested on Mac)
- AdaLayerNorm: add internal nn.LayerNorm(eps=1e-5, affine=False), swap shift/scale ordering - CausalAudioEncoder: SiLU-then-normalize (not softmax) - MotionEncoder_tc: confirm SiLU activation - FramePacker: reverse output order to [fine, medium, coarse] - audio_encoder align_corners: False -> True - audio_encoder wav2vec docstring: bundled path wav2vec2-large-xlsr-53-english - wan_2 ref RoPE: document t_start = max(30, F_video+9) - test freqs exclude
…ing tokens The MotionEncoder_tc convs need elementwise-affine-free LayerNorms between each conv+act stage (kijai auxi_blocks.py MotionEncoder_tc lines 55-77). Without them, intermediate features grew unbounded, giving audio_emb with std ~7-8 that dominated the transformer residuals and corrupted the noise predictions. Also implement the padding_tokens append (num_heads -> num_heads+1 audio tokens per frame) that kijai does at the end of the local path. Also: generate_s2v_video now: - Transposes preprocess_image output (1,1,H,W,3) -> (1,3,1,H,W) to match WanVAE.encode expected channel-first layout - Defaults wav2vec2 to the bundled Wan2.2-S2V-14B/wav2vec2-large-xlsr-53- english path (100x smaller-magnitude features than facebook/wav2vec2- large-xlsr-53 base). Falls back to jonatasgrosman/wav2vec2-large-xlsr-53- english on HF if bundled copy not found. Result: silent-audio smoke test now produces a clean portrait matching the reference image with closed mouth (as expected for silence).
…RoPE Adds multi-segment RoPE construction so ref and motion-history tokens get their own temporal indices (per kijai reference): - ref latent at t_start = max(30, F+9) - motion_post at t=-1 - motion_2x at t=-3 - motion_4x at t=-19..-16 rope_apply now derives seq_len from cos_f.shape[0] when precomputed is supplied so the concatenated [noise, ref, motion] freqs work end-to-end. New helpers in rope.py: - rope_cos_sin_at_positions(positions, dim, theta) — supports negatives - rope_precompute_cos_sin_segments(segments, freqs) — per-segment temporal indices and (H, W) spatial grids. Adds tests/test_wan_rope_segments.py (8 tests, all pass): manual theta verification, negative reflection, single-video parity vs the legacy path, motion buckets, ref at max(30, F+9).
Adds WanS2VModel.prepare_rope_s2v(noise_grid, ref_grid, motion_shapes) to build cos/sin for the full [noise, ref, motion] sequence in one shot. WanS2VModel.__call__ now detects a noise-slice-only rope_cos_sin and transparently rebuilds it so ref/motion tokens get their correct temporal indices (max(30, F+9) for ref; -1/-3/-19 for motion buckets). generate.py generate_s2v_video calls prepare_rope_s2v with the ref grid derived from the VAE-encoded reference (motion_shapes=None for the talking-head first-clip case). test_wan_s2v_load synthetic forward test still passes.
Kijai (model.py L2636-2649) adds cond_mask_weight[1] to ref then cond_mask_weight[0] to the ENTIRE noise+ref sequence -- ref ends up with weight[0]+weight[1]. MLX previously added only seg_embedding(1) to ref, missing the weight[0] baseline that co-embeds ref into the video-token subspace. Suspected root cause of the identity leak (Willy appearing on background monitor rather than as subject face).
…ulse) Kijai nodes_sampler.py L2079-2124: wrapper ALWAYS feeds s2v_ref_motion -- on the first clip it builds torch.zeros([1,3,motion_frames=73,H,W]), VAE-encodes to 19 latent frames, passes as motion_history. Model was trained with these tokens ever-present; MLX passing None left it OOD and caused a visible ~1s luminance/contrast oscillation before frames settled (identical pathology to LTX hero-keyframe warmup). Fix: build all-zero motion latent (1, 16, 19, H_lat, W_lat) and pass on first clip. Also enable motion_shapes in prepare_rope_s2v so the RoPE covers the [-1, -3, -19..-16] motion positions. Skipping the actual VAE(zeros_pixel) encode -- black-frame encoder offset is negligible vs the diffusion schedule and doesn't affect first-clip stability.
Coarse motion projection kernel is (4,8,8); pixel H and W must be divisible by 64 (vae_stride 8 * kernel 8). Previously failed as an AssertionError inside pack_motion_frames on diffusion step 0 -- now raises a ValueError up front with suggested valid resolutions.
Motion tokens returned by frame_packer.pack_motion_frames retain B=1 because motion_history_latent is a single tensor even under CFG. Ref tokens already had this broadcast; motion needed the same fix. Without it, mx.concatenate([noise(B=2), ref(B=2), motion(B=1)...]) fails on axis=1 (batch mismatch).
_motion_bucket_shapes previously returned raw latent H/W (assuming rope_encode_comfy would internally divide by patch kernel), but MLX's rope_precompute_cos_sin_segments uses shape as-is. Result: RoPE segment had 2688 tokens while frame_packer emitted only 672 tokens (rope_apply reshape blew up). Both call sites now return the actual token-grid dimensions after each proj kernel: (1, H/2, W/2), (1, H/4, W/4), (4, H/8, W/8).
…rmup fix) Previously used pure-zero latent for motion history — left ~4 frames of visible warmup because VAE has a learned bias for solid-black pixels that the model was trained to expect. Now VAE-encodes torch.zeros(1,3,73,H,W) matching kijai (nodes_sampler.py L2079-2124) exactly; produces the canonical 19 latent frames after 4x temporal downsample.
Two bugs traced against kijai reference: 1) Color pulse (over-saturated first frames): current code always feeds motion_history_latent = VAE(black 73 pixel frames) => 19 motion tokens biasing the DiT's frame-0 predictions. Kijai's non-framepack context- window path (nodes_sampler.py L2038) does NOT pass s2v_ref_motion and thus has zero motion tokens. Its framepack loop feeds them AND drops the first 3 output pixel frames (line 2168) - the very --trim-first-frames workaround the user forbid. We now follow the non-framepack path: motion_history_latent=None, motion_shapes=None, no motion segment in the RoPE. Warmup pulse has no source. 2) Lipsync shift: kijai model.py L2419 prepends s2v_motion_frames[0]=1 copies of the first audio bucket before the causal audio encoder so latent-frame 0 sees [pad, pad, audio[0], audio[0], audio[1]] instead of [pad, pad, audio[0], audio[1], audio[2]]. Without this the causal conv shifts the effective mouth timing by ~1 audio tick (visible as phoneme mismatch). Applied inside WanS2VModel.__call__ before the casual_audio_encoder invocation.
Kijai reference has 'x = x + audio_x * audio_scale' at model.py L707/797/855 (default 1.0, exposed as tunable). Our port hardcoded 1.0 implicitly. Adds: * AudioInjector.inject(audio_scale=1.0) — scales the cross-attn residual before the video-slice add. When audio_scale==1.0 the graph is shape-identical to the pre-patch version (mx.compile safe). * WanS2VModel.__call__(audio_scale=1.0) — passes through to every post-block inject() call. * generate_s2v_video(audio_scale=1.0) + --audio-scale CLI flag. Motivation: prior user report — 'the audio said 10 words but the mouth only moved 3 times'. audio_scale is the missing amplitude knob for strengthening lipsync without changing CFG / negative prompt / audio temporal density (all already fixed in the Phase 3c commits).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add Wan 2.2 S2V-14B — Audio-driven talking-head video
Adds MLX-native support for Wan 2.2 S2V-14B (Subject-to-Video), the
audio-driven talking-head variant of the Wan 2.2 family. Given a single portrait
image + an audio waveform + a text prompt, generates a lip-synced talking-head
video entirely on Apple Silicon.
Existing Wan T2V / I2V / TI2V paths are unchanged; S2V is an additive
--model-type s2vbranch.Motivation
The Wan 2.2 open-weight release (August 2025) shipped four variants — T2V-1.3B,
TI2V-5B, T2V-A14B, and S2V-14B.
mlx-videocurrently covers the first threebut not S2V. Audio-driven talking-head generation was the missing MLX-native
option for creators on Apple Silicon; the alternatives were either CUDA-only
(kijai's
ComfyUI-WanVideoWrapper) or PyTorch-on-MPS (limited to ~640×384 byMetal SDPA activation memory).
This port runs S2V-14B end-to-end on a 128 GB M-series Mac at 704×384 × 49
frames comfortably, with higher resolutions reachable via SeedVR2 upscale
rather than native scale-up.
What's added
New modules
mlx_video/models/wan_2/audio_encoder.pyCausalAudioEncoder(weighted sum over wav2vec2 25 hidden layers → 3×CausalConv1d→ 4 tokens per video frame),MotionEncoder_tc, wav2vec2 feature extractor bridgemlx_video/models/wan_2/s2v_utils.pyAdaLayerNorm(gated by pooled audio summary),AudioCrossAttention(Q from video, K/V from audio),AudioInjector(injects at blocks {0,4,8,12,16,20,24,27,30,33,36,39}),FramePacker(motion history bucketing)Extended files
config.pyWanModelConfig.wan22_s2v_14b()factory + S2V-only optional fields (all inert for T2V/I2V/TI2V)convert.pysanitize_wan_s2v_weights()— 1260/1260 keys mapped (165 S2V-specific + 1095 delegated to base transformer sanitizer); handles the upstreamcausal_→casual_typo aliaswan_2.pyWanS2VModel(WanModel)with full forward (patch embed → ref token append → framepack motion history → multi-segment RoPE → transformer blocks with per-blockAudioInjector→ head)generate.py--model-type s2v,--audioflags;generate_s2v_videopipeline; wav2vec2 feature extraction wired via HFtransformerson CPU, features moved to MLXrope.pyrope_precompute_cos_sin_segmentsfor the mixed[ref, motion_history, video]temporal sequence with per-segment index offsetsTests
tests/test_wan_s2v_load.pytests/test_wan_rope_segments.pyDocs
docs/WAN_S2V_PORT_DESIGN.mddocs/wan_s2v_keys.txtdocs/wan_s2v_phase2_notes/docs/wan_s2v_phase3_notes/Architecture summary
Wan 2.2 S2V extends the T2V-14B DiT with:
3×
CausalConv1dstack → 4 tokens per video frame (24 fps → 96 audiotokens/s).
K/V from audio tokens.
audio_global(pooled audio summary), applied onattention output before residual add.
RoPE time index (
t_start = max(30, F+9), remapped to-1), providingidentity anchor.
chronologically-ordered memory tokens for temporal consistency across long
clips.
[ref, motion, video]sequence withper-segment temporal grids.
Usage
python -m mlx_video.models.wan_2.generate \ --model-type s2v \ --model-dir /path/to/Wan2.2-S2V-14B-MLX-int4 \ --image portrait.png \ --audio voice.wav \ --prompt "A person speaking to camera, warm lighting" \ --output-path talking_head.mp4 \ --num-frames 49 --width 704 --height 384 --seed 42Weight conversion:
hf download Wan-AI/Wan2.2-S2V-14B --local-dir ./Wan2.2-S2V-14B python -m mlx_video.models.wan_2.convert \ --checkpoint-dir ./Wan2.2-S2V-14B \ --output-dir ./Wan2.2-S2V-14B-MLX-int4 \ --model-type s2v --dtype bfloat16 --quantize 4Validation
Tested on 128 GB M-series Mac with the released
Wan-AI/Wan2.2-S2V-14Bweights (int4-quantized via
mx.quantize):tests/test_wan_s2v_load.pypasses.no NaN, no warmup pulse (fixed by VAE-encoding black-pixel motion history
on first clip — see
a43bd8e).motion produced, identity locked to reference portrait. Head-to-head A/B
against the kijai PyTorch/MPS reference on identical seed/audio/image:
frame-0 color pulse actually comparable-to-better than kijai's (Δ 0.061
vs 0.075), mouth-region motion signal 2.78× cleaner on the mouth ROI vs
kijai's 1.10×. Comparison frames in
docs/wan_s2v_phase3_notes/andproduction_test/.--audio-scaleA/B at 704×384 × 49f on "大家好我是王文欽" (8 phonemes):scale=1.0 (kijai default) hits 7/8 mouth-openings, scale=2.5 hits 8/8 with
no visible over-drive, scale=4.0 hits 8/8 but drifts identity ~1.7×.
Recommended production default:
--audio-scale 2.5; kijai-paritydefault 1.0 preserved as the flag default.
Memory / speed (int4, seed 42, 704×384 × 49 frames): wall-clock ~15-26 min
(varies with
--audio-scale), peak RSS ~14-30 GB.Higher resolutions were not attempted natively; production workflow is
generate-at-704×384 then upscale via SeedVR2 (
mflux-upscale-seedvr2), whichavoids the Metal SDPA activation-memory ceiling and matches the cost/quality
sweet spot of this Mac.
Known limitations / follow-ups
Wan-Video/Wan2.2main) notyet measured (target: mid-block activation PSNR > 30 dB). The behavioral
output looks right; strict bit-for-bit parity is a nice-to-have.
FramePackMotionerbucket ordering and CondEncoder pose-overlaywiring follow the design doc but haven't been end-to-end tested with
multi-clip stitching (>49 frames) or pose-conditioned generation.
peak, may OOM on 128 GB Macs when other MLX processes are resident. Practical
guidance in the README is: generate at 704×384 and use
mflux-upscale-seedvr2for delivery-resolution output.
jonatasgrosman/wav2vec2-large-xlsr-53-english(1.2 GB, cached automatically). Multilingual variants should be pluggable via
env var but not tested here.
References
ComfyUI-WanVideoWrapperPyTorch reference implementation was theauthoritative source for the S2V forward pass (
auxi_blocks.pylayer names,audio-injection sequence, ref-token temporal placement).
wan/modules/s2v/*.pyfor thereference architecture.
mlx_video/models/ltx_2/(already in this repo) — pattern reference foraudio cross-attention wiring inside a video DiT.
Commits (
wan-s2v-portbranch, 16 total)4cb71e9Phase 2 initial (audio encoder + injector + framepack + ref-image + RoPE, untested)b5e05b2close 6 TODO(verify) markers vs kijai reference7f0157bCRITICAL: MotionEncoder_tc LayerNorms + append padding tokens + VAE encode transposee17a669addrope_precompute_cos_sin_segmentsfor ref/motion RoPE0e81069wire multi-segment RoPE intoWanS2VModele18e97dref token needscond_mask_weight[0]baseline (identity fix)ca29540feed zero motion history on first clip (warmup pulse fix)f6b87edearly-fail on non-64-divisible S2V dimensions85e047cbroadcast motion tokens to CFG batch dim08a9f8amotion RoPE shapes use post-projection token grida43bd8eVAE-encode black pixels for motion history (final warmup fix)9e30878Phase 3c: drop motion tokens for single-clip + audio prepend8c9c402Phase 3c: zero audio_input for CFG uncond pass + drop dead VAE(black) encodef174db7Phase 3c: audio temporal density to 16 ticks/sec (kijai bucketing)c289e39Phase 3c: doc the missing S2V weights root cause + merger recipe (was the biggest bug)1f6b2ecexpose--audio-scaleparam (kijai parity; production default 2.5)