LayerGen — RGBA decoder (decoder-only VAE head) · rgba_decoder_packs_b200
A decoder-only RGBA VAE head (~73 M params) that turns the LayerGen decompose foreground latent
— E(fg-over-gray) — into a true RGBA cutout (premultiplied RGB + a soft alpha matte), instead of
the gray-background fg_pred. It is a widened copy of the Wan VAE decoder (shared trunk, 4-channel
head) finetuned with Wan-Alpha reconstruction losses.
⚠️ This is not a standalone generator. It is a routing add-on for the decompose task: at inference it runs a second decode of the same foreground latent produced by a decompose checkpoint, emitting an extra
fg_clean.mp4(RGBA cutout composited on a neutral checkerboard). Use it together with a decompose model, e.g.cs-mshah/decompose_grounded_ema_gbs32_4n_b200.
Checkpoint in this repo
| Path | Which weights | Step |
|---|---|---|
checkpoint-45000/ |
raw / non-EMA (final step) | 45000 |
Model-only DCP export (106 tensors, ~280 MB fp32; the roles.student.transformer decoder only — the
frozen Wan VAE encoder comes from init_from). metadata.json holds the embedded training config;
configs/rgba_decoder_packs.yaml is the config that builds the decoder for inference.
Requirements
- FastVideo with the layer-decomp stack (
fastvideo/train/entrypoint/infer_layer_decomp.py,fastvideo.train.models.rgba_decoder.RgbaDecoderModel). Reference: branchprompt-fix-infer, commitb7ccac76. (Currently internal research code — you need access to it to run inference.) - A local Wan2.2-14B diffusers model for the Wan VAE (
init_fromin the config /--rgba-vae). - A decompose LayerGen checkpoint (this add-on decodes its foreground latent).
- One 80 GB GPU (the decoder itself is tiny; the decompose 14B DiT is the memory driver).
Inference
Run normal decompose inference and add the two --rgba-decoder-* flags (RGBA routing activates only
for task=decompose):
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA
export FASTVIDEO_COMBINED_INFER=0
export FASTVIDEO_FG_CORRUPT_TYPE_ID=2
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
torchrun --standalone --nnodes 1 --nproc_per_node 1 \
-m fastvideo.train.entrypoint.infer_layer_decomp \
--config <decompose infer config>.yaml \
--checkpoint <a DECOMPOSE checkpoint dir> \
--rgba-decoder-ckpt ./checkpoint-45000 \
--rgba-decoder-config configs/rgba_decoder_packs.yaml \
--rgba-vae /path/to/Wan2.2-14B-diffusers \
--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
For each clip this writes an extra fg_clean.mp4 (the RGBA cutout on a checkerboard) alongside the
standard decompose panels. The RGBA fg uses the same foreground latent as fg_pred, so the gray
and RGBA views are exactly matched.
Provenance
Run rgba_decoder_packs_b200 / rgbadec2, checkpoint at step 45000 (final; raw / non-EMA). Trained on
the soft-alpha packs (premult + soft alpha; missions / peeloff / kubric / bedlam) at 480×832, 17-frame
clips, bf16, with Wan-Alpha unit-weight losses (rgb / alpha / perceptual / edge). 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.