cs-mshah's picture
Document EMA vs non-EMA checkpoints
fd8c7c9 verified
|
Raw
History Blame Contribute Delete
4.85 kB
metadata
license: other
license_name: research-preview-see-card
library_name: fastvideo
tags:
  - video
  - video-decomposition
  - video-harmonization
  - layer-decomposition
  - wan2.2
  - diffusion
  - fastvideo
  - layergen
extra_gated_heading: Request access to the LayerGen weights
extra_gated_description: >-
  Access is reviewed manually. Please tell us who you are and how you intend to
  use these research checkpoints.
extra_gated_fields:
  Name: text
  Affiliation: text
  Intended use: text

LayerGen — Combined decompose + harmonize (fg-pred type-4) · combined_abl_fgpred_type4

A single model that does both directions of video layer generation:

  • decompose: composite + mask → foreground + background
  • harmonize: foreground + background → composite

It is the fg-pred type-4 ablation of the combined model: harmonization also predicts the pristine foreground, and the corrupted-foreground observation is a dedicated 5th stream type (id 4). Built on Wan2.2-14B as a multi-stream diffusion transformer with a per-task embedding selected at inference.

  • Base model: Wan2.2-14B diffusers WanPipeline (40-layer / 5120-dim DiT, AutoencoderKLWan VAE, UMT5EncoderModel)
  • Trained at: 480×832, 41 frames (11 latent frames), flow-shift 3.0 · 4-node B200, gbs32 · grounded editcap2 captions

Checkpoints in this repo

Path Which weights Step
checkpoint-8000/ EMA (exponential moving average) 8000
nonema/checkpoint-7600/ raw / non-EMA (the "pure" training weights) 7600

Both are model-only DCP exports (roles.* = the trained DiT only, ~53 GB fp32 each; the frozen Wan VAE / text encoder come from init_from). Each dir is a self-contained checkpoint-<step>/ with dcp/ inside, plus a metadata.json (embedded training config). EMA weights are usually the better choice for sampling; the non-EMA weights are provided for completeness / ablation. (Note the different steps: the EMA export is at 8000, the raw export at 7600.)

Requirements

  • FastVideo with the layer-decomp stack (fastvideo/train/entrypoint/infer_layer_decomp.py, fastvideo.train.models.layer_decomp.LayerDecompWanModel). Reference: branch prompt-fix-infer, commit b7ccac76. (Currently internal research code — you need access to it to run inference.)
  • A local Wan2.2-14B diffusers model for the VAE + architecture; point init_from in the config at it.
  • One 80 GB GPU (A100-80G / H100 / H200 / B200). FlashAttention recommended (TORCH_SDPA also works).

Inference

Because this is a combined checkpoint, set FASTVIDEO_COMBINED_INFER=1 (applies the trained per-task embedding under a single-task config) and pick the config for the direction you want. Being a type-4 model, it also needs FASTVIDEO_FG_CORRUPT_TYPE_ID=4.

# Point the chosen config:  init_from -> your local Wan2.2-14B diffusers dir,
#                           data_path -> your encoded latent pool (see encode_evals.py).
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN        # or TORCH_SDPA
export FASTVIDEO_COMBINED_INFER=1                    # REQUIRED: apply the trained per-task embedding
export FASTVIDEO_FG_CORRUPT_TYPE_ID=4                # REQUIRED: 5-type stream layout (fg-corrupt = its own type id 4)
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export TOKENIZERS_PARALLELISM=false WANDB_MODE=disabled

# EMA weights:      --checkpoint ./checkpoint-8000
# non-EMA weights:  --checkpoint ./nonema/checkpoint-7600
# DECOMPOSE:  --config configs/infer_decompose.yaml     HARMONIZE:  --config configs/infer_harmonize.yaml
torchrun --standalone --nnodes 1 --nproc_per_node 1 \
  -m fastvideo.train.entrypoint.infer_layer_decomp \
  --config configs/infer_decompose.yaml \
  --checkpoint ./checkpoint-8000 \
  --out ./out --steps 30 --model-only \
  --training.distributed.num_gpus 1 \
  --training.distributed.hsdp_replicate_dim 1 \
  --training.distributed.hsdp_shard_dim 1 \
  --training.distributed.sp_size 1

⚠️ Both env vars matter. FASTVIDEO_FG_CORRUPT_TYPE_ID=4 is required or the DCP load fails with a shape mismatch ([5,5120] vs [4,5120]); FASTVIDEO_COMBINED_INFER=1 is required or the per-task embedding is not applied and quality degrades.

Outputs are written under ./out/samples/<clip_id>/.

Longer clips: trained at 41 frames. For 61f/81f set FASTVIDEO_ROPE_T_INTERP=(41-1)/(T_infer-1) and bump num_latent_t (16 for 61f, 21 for 81f).

Provenance

Run combined_abl_fgpred_type4 (wandb abl_fgpred_type4_v2). EMA export at step 8000; raw (non-EMA) export at step 7600. Part of the LayerGen layer-decomposition project.

License

Research preview released for reproducibility. Built on Wan2.2 (see the Wan-AI license for the base model). Contact the authors for reuse terms.