README.md

---
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`)