wlyu/raystream_cfpp

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 (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:

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:

# configs/cfpp_causal_cd_framewise_ti2v704.yaml
generator_ckpt: /path/to/ar_diffusion/model.pt
ucpe_ckpt:      /path/to/bidirectional_teacher/last.ckpt

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

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

@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 (arXiv 2602.02214) and UCPE. Contact: Weijie Lyu (weijielyu1@gmail.com)