SceneWorks
/

scail2-mlx

@@ -20,18 +20,28 @@ base_model: zai-org/SCAIL-2
 > Capabilities (from upstream): character animation from a reference image + driving video, **cross-identity character replacement**, zero-shot animal-driving, end-to-end *and* pose-rendered driving, and (experimental) multi-reference. Image output is `num_frames == 1`.
 ## What changed vs. upstream
-- The main DiT (`model/1/fsdp2_rank_0000_checkpoint.pt`, an FSDP2/SAT checkpoint) was key-remapped to the `SCAIL2Model` parameter naming using the upstream `convert.py` contract (fused `query_key_value`→`q`/`k`/`v`, `key_value`→`k`/`v`, `clip_feature_key_value_list`→`k_img`/`v_img`) and cast **fp32 → bf16**, then saved as a single safetensors. The remap is bit-faithful (validated: 987 source keys → 1307 model keys, an exact key+shape match against `SCAIL2Model.from_config(config-14b.json)`).
-- The text encoder (UMT5-XXL), VAE (Wan2.1), and image encoder (open-CLIP XLM-RoBERTa ViT-H/14) are the **stock upstream files**, included here so the snapshot is self-contained.
-## Contents
-| file | source | notes |
-|---|---|---|
-| `dit.safetensors` | converted | SCAIL-2 14B DiT, **bf16**, 1307 tensors (~31 GB) |
-| `config.json` | upstream `configs/config-14b.json` | `model_type: i2v`, `dim 5120`, `ffn 13824`, `40` layers/heads, `in_dim 20`, `mask_dim 28`, `out_dim 16` |
-| `Wan2.1_VAE.pth` | upstream, stock | z16 VAE, stride (4,8,8) |
-| `umt5-xxl/` | upstream, stock | UMT5-XXL encoder (bf16) + tokenizer |
-| `models_clip_open-clip-xlm-roberta-large-vit-huge-14-onlyvisual.pth` | upstream, stock | reference-image visual encoder (1280-dim) |
-| `bias-aware-dpo-lora.pt` | upstream, stock | optional Bias-Aware DPO refinement LoRA |
 ## Architecture (summary)
 Wan2.1-14B **I2V** dense DiT. Conditioning is a **token-axis packed** stream — reference + video + pose patch-embedded (three Conv3d stems) with additive 28-channel color-coded mask embeddings, concatenated into one self-attention sequence — plus a **per-source RoPE** with integer T/H/W shifts (the `replace_flag` flips the reference H-shift, toggling animation vs. replacement). The reference image is encoded by the CLIP visual tower and injected via Wan-I2V image cross-attention. Sampling is plain CFG (guide 5.0), flow-matching UniPC/DPM++.

 > Capabilities (from upstream): character animation from a reference image + driving video, **cross-identity character replacement**, zero-shot animal-driving, end-to-end *and* pose-rendered driving, and (experimental) multi-reference. Image output is `num_frames == 1`.
 ## What changed vs. upstream
+Every component is repackaged to the safetensors layout the SceneWorks Rust/MLX loaders consume — no PyTorch at runtime:
+- **DiT** (`model/1/fsdp2_rank_0000_checkpoint.pt`, an FSDP2/SAT checkpoint) was key-remapped to the `SCAIL2Model` parameter naming using the upstream `convert.py` contract (fused `query_key_value`→`q`/`k`/`v`, `key_value`→`k`/`v`, `clip_feature_key_value_list`→`k_img`/`v_img`) and cast **fp32 → bf16** → `dit.safetensors`. Bit-faithful (987 source keys → 1307 model keys; exact key+shape match against `SCAIL2Model.from_config(config-14b.json)`).
+- **VAE** (`Wan2.1_VAE.pth`, the stock Wan2.1 z16 VAE) → `vae.safetensors` (**f32**, channels-last conv transpose, keys unchanged — the `sanitize_wan_vae_weights` contract shared with Bernini/wan). Loaded by `mlx_gen_wan::WanVae`.
+- **Text encoder** (`umt5-xxl/models_t5_umt5-xxl-enc-bf16.pth`, stock UMT5-XXL) → `t5_encoder.safetensors` (**bf16**, sole rename `.ffn.gate.0.`→`.ffn.gate_proj.`). Loaded by `mlx_gen_wan::Umt5Encoder` with `tokenizer.json`.
+- **Image encoder** (`models_clip_...onlyvisual.pth`, open-CLIP XLM-RoBERTa ViT-H/14) → `clip.safetensors` (**f32**, de-prefixed `visual.*` keys). Loaded by `mlx_gen_scail2::ScailClip` (32-layer visual tower, `use_31_block` penultimate features).
+The converted VAE/UMT5 are byte-size-identical (modulo safetensors header) to Bernini/wan's already-validated Wan2.1 VAE + umt5-xxl safetensors — confirming SCAIL-2 ships the stock components.
+## Contents (turnkey MLX snapshot)
+| file | source | loader | notes |
+|---|---|---|---|
+| `dit.safetensors` | converted | `Scail2Dit` | SCAIL-2 14B DiT, **bf16**, 1307 tensors (~31 GB) |
+| `vae.safetensors` | converted | `WanVae` | Wan2.1 z16 VAE, **f32**, stride (4,8,8) (~0.5 GB) |
+| `t5_encoder.safetensors` | converted | `Umt5Encoder` | UMT5-XXL encoder, **bf16** (~11 GB) |
+| `clip.safetensors` | converted | `ScailClip` | open-CLIP ViT-H/14 visual tower, **f32**, 1280-dim (~2.5 GB) |
+| `tokenizer.json` | upstream, stock | `load_tokenizer` | UMT5-XXL HF tokenizer (root copy) |
+| `config.json` | upstream `configs/config-14b.json` | `Scail2Config` | `model_type: i2v`, `dim 5120`, `ffn 13824`, `40` layers/heads, `in_dim 20`, `mask_dim 28`, `out_dim 16` |
+| `bias-aware-dpo-lora.pt` | upstream, stock | (sc-5451) | optional Bias-Aware DPO refinement LoRA |
+Quantization (Q4/Q8) is applied at **load time** by the SceneWorks worker (`mlx_gen` `.quantize()`); this snapshot ships the dense bf16/f32 weights.
+The raw upstream pickles (`Wan2.1_VAE.pth`, `umt5-xxl/models_t5_umt5-xxl-enc-bf16.pth`, `models_clip_...onlyvisual.pth`) remain in this repo for provenance; the Rust loaders use only the converted safetensors above, so a lean SceneWorks pull can skip them.
 ## Architecture (summary)
 Wan2.1-14B **I2V** dense DiT. Conditioning is a **token-axis packed** stream — reference + video + pose patch-embedded (three Conv3d stems) with additive 28-channel color-coded mask embeddings, concatenated into one self-attention sequence — plus a **per-source RoPE** with integer T/H/W shifts (the `replace_flag` flips the reference H-shift, toggling animation vs. replacement). The reference image is encoded by the CLIP visual tower and injected via Wan-I2V image cross-attention. Sampling is plain CFG (guide 5.0), flow-matching UniPC/DPM++.