| --- |
| license: mit |
| library_name: wan2.2 |
| pipeline_tag: image-to-video |
| tags: |
| - video-generation |
| - diffusion |
| - camera-control |
| - causal-video |
| - wan2.2 |
| - ucpe |
| base_model: Wan-AI/Wan2.2-TI2V-5B |
| --- |
| |
| # RayStream CF++ — bidirectional teacher + AR-diffusion student |
|
|
| Checkpoints from the **Causal Forcing++ (CF++)** frame-wise rebuild: camera-controlled, |
| image-conditioned video generation on **Wan2.2-TI2V-5B at 704×1280**, with UCPE camera |
| conditioning (PRoPE) throughout. |
|
|
| Code: [github.com/weijielyu/RayStream_CF](https://github.com/weijielyu/RayStream_CF) (`main`) |
|
|
| This repo holds the two checkpoints at the **front** of the CF++ pipeline: |
|
|
| ``` |
| UCPE bidirectional teacher (this repo: bidirectional_teacher/) |
| └─> AR-diffusion (this repo: ar_diffusion/) teacher-forced warmup, step 6200 |
| └─> Causal-CD consistency distillation, effective step 3000 |
| └─> DMD final few-step causal student |
| ``` |
|
|
| ## Contents |
|
|
| ``` |
| . |
| ├── bidirectional_teacher/ |
| │ └── last.ckpt # 10.1 GiB — PyTorch Lightning ckpt, run ti2v_relray_absmap_comp8_704_fixedloss |
| └── ar_diffusion/ |
| └── model.pt # 19.2 GiB — CF++ AR-diffusion, step 6200 (generator only) |
| ``` |
|
|
| | File | Format | Top-level keys | |
| |---|---|---| |
| | `bidirectional_teacher/last.ckpt` | Lightning (DeepSpeed stage-1 consolidated) | `state_dict` with `pipe.dit.*` prefix | |
| | `ar_diffusion/model.pt` | CF++ trainer save | `generator` | |
|
|
| The AR checkpoint is **generator-only** — do not pass `--use_ema` when loading it |
| (`ema_start_step` was set past the run length, so no EMA weights exist). |
|
|
| ## Model configuration |
|
|
| Both share the same geometry and camera setup: |
|
|
| - **Base**: Wan2.2-TI2V-5B — DiT dim 3072, 24 heads (head_dim 128), 30 layers, `in_dim`/`z_dim` 48 |
| - **Resolution**: 704×1280, 81 pixel frames @ 16 fps |
| - **VAE**: `WanVideoVAE38` (48-ch, temporal 4×) → 21 latent frames = 81 pixel frames |
| - **Latent shape**: `[1, 21, 48, 44, 80]`; patch factor 32 → `frame_seq_length` 880, `seq_len` 18480 |
| - **Camera**: UCPE Unified Camera Model (`x_fov`, `xi`), method `relray_absmap` |
| (Relative Ray Encoding + absolute orientation via up-vector and latitude map), |
| `attn_compress=8`. Injected as a parallel `cam_self_attn` (PRoPE) branch next to DiT |
| self-attention — ~300 tensors / 142 M params (2.8% of the model). PRoPE is KV-cache |
| compatible, which is what lets the camera branch survive autoregressive distillation. |
| - **I2V**: frame 0 is the clean condition image — pinned at timestep 0, kept clean, excluded |
| from the loss, never noised or augmented. |
|
|
| The AR student grafts the teacher's camera branch **frozen** (`freeze_camera_branch: true`). |
|
|
| ## ⚠️ This is the *corrected* teacher |
|
|
| An earlier bidirectional teacher (published in `wlyu/ucpe_checkpoints` as |
| `wan22_bidirectional_ucpe/`) baked in an I2V training bug: `training_loss` noised **and** |
| supervised frame 0, even though TI2V-5B pins frame 0 clean at t=0 and inference re-pins it every |
| step. Frame 0 was therefore fed noised while labeled t=0 — so it never learned the clean |
| first-frame anchor, and first-frame content ghosted onto later frames through bidirectional |
| attention. |
|
|
| The fix re-pins frame 0 after `add_noise` and excludes it from both `noise_pred` and the training |
| target. The model was **retrained from scratch** with the fix; that retrain is the checkpoint here |
| (run `ti2v_relray_absmap_comp8_704_fixedloss`, 10 epochs / 10000 steps, batch 32 on 4×8 B200). |
|
|
| **Use this one, not `wlyu/ucpe_checkpoints/wan22_bidirectional_ucpe/`**, which is superseded. |
| |
| ## Usage |
| |
| ### Load the teacher's camera branch into a CF++ model |
| |
| The teacher is consumed by CF++ as the frozen camera branch. `load_ucpe_camera_weights` accepts |
| this Lightning checkpoint directly — it unwraps `state_dict`, keeps keys containing |
| `cam_self_attn`, strips the `pipe.dit.` prefix, and loads with `strict=False`: |
| |
| ```python |
| from utils.camera_control import patch_causal_dit, load_ucpe_camera_weights |
| |
| patch_causal_dit(model, method="relray_absmap", attn_compress=8) |
| load_ucpe_camera_weights(model, "bidirectional_teacher/last.ckpt") |
| # -> Loaded 300 camera branch params, 825 missing (expected: base model keys), 0 unexpected |
| ``` |
| |
| In CF++ configs this is just `ucpe_ckpt: /path/to/last.ckpt`. |
| |
| ### Continue training from the AR checkpoint |
| |
| Warm-start Causal-CD (or resume AR) by pointing `generator_ckpt` at the AR weights. Note this is |
| a **weights-only** warm start — the optimizer state and step counter are not restored, so the step |
| counter restarts at 0 and the effective step is `6200 + new_step`: |
| |
| ```yaml |
| # configs/cfpp_causal_cd_framewise_ti2v704.yaml |
| generator_ckpt: /path/to/ar_diffusion/model.pt |
| ucpe_ckpt: /path/to/bidirectional_teacher/last.ckpt |
| ``` |
| |
| ```bash |
| export PYTHONPATH=$PWD TOKENIZERS_PARALLELISM=false HF_HUB_OFFLINE=1 |
| torchrun --nnodes=1 --nproc_per_node=8 --master_port=29562 train.py \ |
| --config_path configs/cfpp_causal_cd_framewise_ti2v704.yaml \ |
| --disable-wandb --no_visualize --logdir /path/to/out |
| ``` |
| |
| You also need the Wan2.2-TI2V-5B base weights (`hf download Wan-AI/Wan2.2-TI2V-5B`) symlinked as |
| `wan_models/Wan2.2-TI2V-5B` in the code tree. |
| |
| ### Download |
| |
| ```bash |
| hf download wlyu/raystream_cfpp --local-dir ./raystream_cfpp |
| |
| # or just one: |
| hf download wlyu/raystream_cfpp --include 'ar_diffusion/*' --local-dir ./raystream_cfpp |
| ``` |
| |
| ## Training |
| |
| Both trained on 8×B200 (teacher: 4 nodes), data = **PanShot** panoramic video with ground-truth |
| camera poses at 704p. |
|
|
| **AR-diffusion** used upstream Causal-Forcing's optimizer recipe verbatim — lr 2e-6, beta1 0.0, |
| no weight decay, `num_frame_per_block: 1` (frame-wise), teacher forcing on, batch 8 — with only |
| the 704p/5B geometry, camera control, and I2V conditioning added. It ran to step 6200 (~6 s/step), |
| past the ~5700 plateau of an earlier run. |
|
|
| Expect **low-motion, near-static output** from the AR checkpoint. That is inherent to teacher |
| forcing on panoramic data and is normal at this stage; motion dynamics are restored later, at the |
| DMD stage. Upstream guidance is AR ≥2K steps (5–10K better), Causal-CD ≥3K, and DMD ≤1K at bs64 — |
| beyond ~1K, DMD over-fits to the teacher's static-biased distribution and motion degrades. |
|
|
| ## Evaluation |
|
|
| PanShot test split, held out at **clip** level (0/340 exact overlap with train) but **not |
| scene-disjoint** — roughly 33 of 51 test source videos also appear in training as different |
| segments. Metrics span video quality (FVD, FID, CLIP-T) and camera accuracy (FOV error, |
| distortion k1/k2, pitch/roll, rotation/translation error, CAMMC), computed with the UCPE metric |
| harness. |
|
|
| ## Citation |
|
|
| ```bibtex |
| @misc{raystream_cfpp_2026, |
| author = {Lyu, Weijie}, |
| title = {RayStream CF++: Causal Forcing with UCPE Camera Control}, |
| year = {2026}, |
| url = {https://huggingface.co/wlyu/raystream_cfpp} |
| } |
| ``` |
|
|
| Built on [thu-ml/Causal-Forcing](https://github.com/thu-ml/Causal-Forcing) (arXiv 2602.02214) and |
| UCPE. Contact: Weijie Lyu (`weijielyu1@gmail.com`) |
|
|
