diff --git a/video/cosmos3/README.md b/video/cosmos3/README.md new file mode 100644 index 000000000..761b012fa --- /dev/null +++ b/video/cosmos3/README.md @@ -0,0 +1,152 @@ +Cosmos 3 Nano +============= + +[NVIDIA Cosmos 3](https://github.com/NVIDIA/cosmos) Nano (16B) text-to-video +and image-to-video generation on Apple Silicon via MLX. Cosmos 3 is a +world foundation model for **physical AI** — robotics, autonomous driving, +and industrial simulation. The model weights are downloaded from the +[Hugging Face Hub](https://huggingface.co/nvidia/Cosmos3-Nano). + +| Config | RAM (bf16) | RAM (8-bit) | +|--------|-----------|------------| +| Cosmos3-Nano 16B | ~48 GB | ~24 GB | + +> **Model scope:** Cosmos 3 was trained on robotics manipulation, autonomous +> driving, and industrial/factory environments. It produces strong physical +> motion for on-distribution inputs (dashcam driving, robot arms, factory +> floors) but does not generalize well to arbitrary creative prompts. This +> matches [NVIDIA's model card](https://huggingface.co/nvidia/Cosmos3-Nano). + +Installation +------------ + +Install the dependencies: + +```shell +pip install -r requirements.txt +``` + +Download the model weights (~32 GB at bf16): + +```shell +hf download nvidia/Cosmos3-Nano --local-dir weights/Cosmos3-Nano +``` + +Saving videos as MP4 requires [ffmpeg](https://ffmpeg.org/) on your PATH. +If ffmpeg is not installed, output will be saved as GIF instead. + +Usage +----- + +> **Note:** The examples below use `--quantize` for 8-bit mode (~24 GB). +> Without `--quantize`, the model runs at bf16 and requires ~48 GB. + +### Text-to-Video + +Generate a video with an on-distribution prompt: + +```shell +python txt2video.py 'A car driving through a suburban intersection on a sunny day' \ + --quantize --output out.mp4 +``` + +Higher resolution: + +```shell +python txt2video.py 'A delivery truck backing into a warehouse loading dock' \ + --size 832x480 --frames 16 --steps 30 --guidance 6.0 --seed 42 \ + --quantize --output out_480p.mp4 +``` + +### Image-to-Video + +Generate a video conditioned on an input image (provide your own JPEG/PNG): + +```shell +python img2video.py 'A robot arm reaches toward a red block on a table' \ + --image your_image.jpg --quantize --output out_i2v.mp4 +``` + +### Audio + +Joint video+audio generation (stereo 48 kHz, muxed into MP4): + +```shell +python txt2video.py 'A robot arm pushes a metal box across a table' \ + --enable-audio --quantize --output out_audio.mp4 +``` + +### Quantization + +Pass `--quantize` (or `-q`) to quantize the transformer weights to 8-bit, +reducing model weight memory from ~32 GB to ~16 GB (total runtime memory +is higher due to activations and VAE): + +```shell +python txt2video.py 'A forklift moving pallets in a warehouse' \ + --quantize --output out_q8.mp4 +``` + +### Disabling the cache + +For additional memory savings at the expense of speed, use `--no-cache`: + +```shell +python txt2video.py 'A robot arm sorting objects on a conveyor belt' \ + --quantize --no-cache --output out_nocache.mp4 +``` + +### Options + +- **Negative prompts:** `--n-prompt 'blurry, low quality'` (default: model's + built-in negative prompt) +- **Guidance scale:** `--guidance 6.0` (default) +- **Denoising steps:** `--steps 30` (default) +- **Random seed:** `--seed 42` + +For all options, use `python txt2video.py --help`. + +Performance +----------- + +Measured on M4 Max (128 GB), 8-bit quantized, 30 denoising steps: + +| Resolution | Frames | Generation time | Peak memory | +|------------|--------|----------------|-------------| +| 256x256 | 16 | ~38s | ~17 GB | +| 480p (832x480) | 16 | ~252s | ~24 GB | +| 720p (1280x720) | 16 | ~591s | ~48 GB | + +The pipeline caches text token KV pairs across denoising steps (text +embeddings are constant), which significantly reduces per-step compute +at lower resolutions. + +Hardware requirements: + +- **8-bit quantized, 256p:** ~17 GB peak (measured on M4 Max); 24 GB+ recommended +- **8-bit quantized, 480p:** ~24 GB peak; 48 GB+ recommended +- **bf16 full precision:** ~48 GB peak; 48 GB+ Mac (M4 Max or higher) + +Architecture +------------ + +Cosmos 3 uses a Mixture-of-Transformers (MoT) design with two pathways: + +- **Understanding (reasoner):** causal self-attention (Qwen3-VL text backbone) +- **Generation (diffuser):** full bidirectional attention for video/audio synthesis + +Video VAE: Wan2.2 AutoencoderKL (16x spatial, 4x temporal downsampling). +Audio: Cosmos3 AVAEAudioTokenizer (Oobleck decoder, stereo 48 kHz). +Scheduler: UniPC multi-step predictor-corrector. + +License +------- + +Model weights are under [NVIDIA OpenMDW 1.1](https://openmdw.ai/license/1-1/) +(commercial and non-commercial use permitted). + +References +---------- + +1. [NVIDIA Cosmos 3](https://github.com/NVIDIA/cosmos) +2. [Cosmos3-Nano model card](https://huggingface.co/nvidia/Cosmos3-Nano) diff --git a/video/cosmos3/cosmos3/__init__.py b/video/cosmos3/cosmos3/__init__.py new file mode 100644 index 000000000..c97269c97 --- /dev/null +++ b/video/cosmos3/cosmos3/__init__.py @@ -0,0 +1 @@ +"""NVIDIA Cosmos 3 on Apple Silicon via MLX.""" diff --git a/video/cosmos3/cosmos3/attention.py b/video/cosmos3/cosmos3/attention.py new file mode 100644 index 000000000..0feeb61d8 --- /dev/null +++ b/video/cosmos3/cosmos3/attention.py @@ -0,0 +1,240 @@ +"""Cosmos 3 dual-pathway Mixture-of-Transformers attention. + +The MoT attention has two pathways: + - Understanding (reasoner): causal self-attention with standard Q/K/V + - Generation (diffuser): full attention over [und + gen] tokens with separate Q/K/V + +Both share the same RoPE and layer structure. +""" + +from typing import Optional, Tuple + +import mlx.core as mx +import mlx.nn as nn + +from .rope import Cosmos3RotaryEmbedding, apply_rotary_pos_emb + + +class Cosmos3Attention(nn.Module): + """Dual-pathway packed attention for Cosmos 3 MoT.""" + + def __init__( + self, + hidden_size: int = 4096, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + head_dim: int = 128, + mrope_section: list[int] | None = None, + rope_theta: float = 5_000_000.0, + rms_norm_eps: float = 1e-6, + ): + super().__init__() + self.hidden_size = hidden_size + self.num_heads = num_attention_heads + self.num_kv_heads = num_key_value_heads + self.head_dim = head_dim + self.scale = head_dim ** -0.5 + + # Understanding pathway Q/K/V/O projections + self.to_q = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False) + self.to_k = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.to_v = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.to_out = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False) + + # QK norms (per-head RMSNorm) + self.norm_q = nn.RMSNorm(head_dim, eps=rms_norm_eps) + self.norm_k = nn.RMSNorm(head_dim, eps=rms_norm_eps) + + # Generation pathway projections + self.add_q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=False) + self.add_k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.add_v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=False) + self.to_add_out = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=False) + self.norm_added_q = nn.RMSNorm(head_dim, eps=rms_norm_eps) + self.norm_added_k = nn.RMSNorm(head_dim, eps=rms_norm_eps) + + # RoPE + self.rope = Cosmos3RotaryEmbedding( + head_dim=head_dim, + mrope_section=mrope_section or [24, 20, 20], + rope_theta=rope_theta, + ) + + def _project_and_reshape( + self, + x: mx.array, + proj: nn.Linear, + num_heads: int, + ) -> mx.array: + """Project and reshape to [batch, seq_len, num_heads, head_dim].""" + batch, seq_len, _ = x.shape + out = proj(x) + return out.reshape(batch, seq_len, num_heads, self.head_dim) + + def __call__( + self, + hidden_states: mx.array, + position_ids: mx.array, + understanding_mask: Optional[mx.array] = None, + generation_tokens: Optional[mx.array] = None, + cache: Optional[Tuple[mx.array, mx.array]] = None, + ) -> Tuple[mx.array, Optional[mx.array], Optional[Tuple[mx.array, mx.array]]]: + """Forward pass. + + Args: + hidden_states: [batch, und_len, hidden_size] understanding tokens + position_ids: [3, batch, total_len] position IDs per axis + understanding_mask: optional attention mask + generation_tokens: [batch, gen_len, hidden_size] or None + cache: optional (keys, values) KV cache tuple + + Returns: + (und_output, gen_output, updated_cache) + gen_output is None if generation_tokens is None + """ + batch, und_len, _ = hidden_states.shape + + # Understanding pathway Q/K/V + q = self._project_and_reshape(hidden_states, self.to_q, self.num_heads) + k = self._project_and_reshape(hidden_states, self.to_k, self.num_kv_heads) + v = self._project_and_reshape(hidden_states, self.to_v, self.num_kv_heads) + + # QK normalization + q = self.norm_q(q) + k = self.norm_k(k) + + # Apply RoPE + und_position_ids = position_ids[:, :, :und_len] + cos, sin = self.rope(und_position_ids, seq_len=und_len) + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + # KV cache + if cache is not None: + k_cache, v_cache = cache + k = mx.concatenate([k_cache, k], axis=1) + v = mx.concatenate([v_cache, v], axis=1) + + new_cache = (k, v) + + # Save un-expanded keys/values for generation pathway + k_unexpanded, v_unexpanded = k, v + + # GQA: repeat KV heads to match query heads for understanding attention + k_attn, v_attn = k, v + if self.num_kv_heads != self.num_heads: + repeat_factor = self.num_heads // self.num_kv_heads + k_attn = mx.repeat(k, repeat_factor, axis=2) + v_attn = mx.repeat(v, repeat_factor, axis=2) + + # Compute attention: [batch, seq_len, num_heads, head_dim] + # Transpose to [batch, num_heads, seq_len, head_dim] for SDPA + q_t = mx.transpose(q, (0, 2, 1, 3)) + k_t = mx.transpose(k_attn, (0, 2, 1, 3)) + v_t = mx.transpose(v_attn, (0, 2, 1, 3)) + + # Causal attention for understanding pathway + q_len = q_t.shape[2] + k_len = k_t.shape[2] + + if q_len == 1: + # Single-token generation: no mask needed + attn_out = mx.fast.scaled_dot_product_attention( + q_t, k_t, v_t, scale=self.scale + ) + else: + # Prefill or multi-token: apply causal mask + # Create full causal mask over key length, take last q_len rows + full_mask = nn.MultiHeadAttention.create_additive_causal_mask( + k_len, dtype=q_t.dtype + ) + # When cache is present, we only have q_len query positions + # corresponding to the last q_len rows of the full causal mask + mask = full_mask[-q_len:] + attn_out = mx.fast.scaled_dot_product_attention( + q_t, k_t, v_t, scale=self.scale, mask=mask + ) + + # Transpose back and project: [batch, seq_len, num_heads * head_dim] + attn_out = mx.transpose(attn_out, (0, 2, 1, 3)) + attn_out = attn_out.reshape(batch, -1, self.num_heads * self.head_dim) + und_output = self.to_out(attn_out) + + # Generation pathway + gen_output = None + if generation_tokens is not None: + gen_output = self._generation_forward( + generation_tokens, hidden_states, + k_unexpanded, v_unexpanded, + position_ids, und_len, + ) + + return und_output, gen_output, new_cache, (k_unexpanded, v_unexpanded) + + def generation_only_forward( + self, + gen_tokens: mx.array, + und_kv: Tuple[mx.array, mx.array], + position_ids: mx.array, + und_len: int, + ) -> mx.array: + """Generation pathway only, using cached understanding K/V. + + Skips the entire understanding pathway (Q/K/V projection, attention, output). + Uses pre-computed understanding keys/values for cross-attention. + """ + return self._generation_forward( + gen_tokens, None, und_kv[0], und_kv[1], + position_ids, und_len, + ) + + def _generation_forward( + self, + gen_tokens: mx.array, + und_tokens: mx.array, + und_keys: mx.array, + und_values: mx.array, + position_ids: mx.array, + und_len: int, + ) -> mx.array: + """Generation pathway: full attention over [und + gen] tokens. + + Bidirectional attention for diffusion. + """ + batch, gen_len, _ = gen_tokens.shape + + # Generation Q/K/V + q_gen = self._project_and_reshape(gen_tokens, self.add_q_proj, self.num_heads) + k_gen = self._project_and_reshape(gen_tokens, self.add_k_proj, self.num_kv_heads) + v_gen = self._project_and_reshape(gen_tokens, self.add_v_proj, self.num_kv_heads) + + # QK normalization + q_gen = self.norm_added_q(q_gen) + k_gen = self.norm_added_k(k_gen) + + # Apply RoPE to generation tokens + gen_position_ids = position_ids[:, :, und_len : und_len + gen_len] + cos, sin = self.rope(gen_position_ids, seq_len=gen_len) + q_gen, k_gen = apply_rotary_pos_emb(q_gen, k_gen, cos, sin) + + # Concatenate [und + gen] keys/values for full attention + k_full = mx.concatenate([und_keys, k_gen], axis=1) + v_full = mx.concatenate([und_values, v_gen], axis=1) + + # GQA expansion + if self.num_kv_heads != self.num_heads: + repeat_factor = self.num_heads // self.num_kv_heads + k_full = mx.repeat(k_full, repeat_factor, axis=2) + v_full = mx.repeat(v_full, repeat_factor, axis=2) + + # Full (non-causal) attention + q_t = mx.transpose(q_gen, (0, 2, 1, 3)) + k_t = mx.transpose(k_full, (0, 2, 1, 3)) + v_t = mx.transpose(v_full, (0, 2, 1, 3)) + + attn_out = mx.fast.scaled_dot_product_attention( + q_t, k_t, v_t, scale=self.scale + ) + + attn_out = mx.transpose(attn_out, (0, 2, 1, 3)) + attn_out = attn_out.reshape(batch, gen_len, self.num_heads * self.head_dim) + return self.to_add_out(attn_out) diff --git a/video/cosmos3/cosmos3/convert.py b/video/cosmos3/cosmos3/convert.py new file mode 100644 index 000000000..2a89ca6a4 --- /dev/null +++ b/video/cosmos3/cosmos3/convert.py @@ -0,0 +1,83 @@ +"""Convert Cosmos 3 weights from HuggingFace safetensors to MLX format. + +Handles: +- Weight name mapping (to_out.0 -> to_out, etc.) +- Reasoner-only mode: strips generation/diffusion weights +- Loading sharded safetensors from HuggingFace Hub +""" + +import re + +import mlx.core as mx + +# Patterns for generation/diffusion-only weights (skipped in reasoner mode) +GENERATION_PATTERNS = [ + r".*_moe_gen.*", # MoE generation layers + r".*add_q_proj.*", # Generation Q projections + r".*add_k_proj.*", # Generation K projections + r".*add_v_proj.*", # Generation V projections + r".*to_add_out.*", # Generation output projections + r".*norm_added_.*", # Generation QK norms + r".*norm_moe_gen.*", # Generation final norm + r"proj_in\..*", # Diffusion input projection + r"proj_out\..*", # Diffusion output projection + r"audio_proj_in\..*", # Audio input projection + r"audio_proj_out\..*", # Audio output projection + r"action_proj_in\..*", # Action input projection + r"action_proj_out\..*", # Action output projection + r"time_embedder\..*", # Timestep embedder + r"action_modality_embed", # Action modality embedding + r"audio_modality_embed", # Audio modality embedding +] + +_GENERATION_RE = [re.compile(p) for p in GENERATION_PATTERNS] + +# Weight name remapping: HuggingFace -> MLX +# Most names stay the same. The main change is to_out.0.weight -> to_out.weight +WEIGHT_NAME_MAP = { + # The HF diffusers model wraps output projection in nn.ModuleList + # so it's to_out.0.weight instead of to_out.weight +} + + +def map_weight_name(name: str) -> str: + """Map a HuggingFace weight name to the MLX model name. + + The main transformation: to_out.0.weight -> to_out.weight + (and to_add_out.0.weight -> to_add_out.weight) + """ + # Strip the .0 index from output projections + name = re.sub(r"\.to_out\.0\.", ".to_out.", name) + name = re.sub(r"\.to_add_out\.0\.", ".to_add_out.", name) + return name + + +def _is_generation_weight(name: str) -> bool: + """Check if a weight belongs to the generation/diffusion pathway.""" + return any(p.search(name) for p in _GENERATION_RE) + + +def convert_weights( + weights: dict[str, mx.array], + reasoner_only: bool = True, +) -> dict[str, mx.array]: + """Convert a dict of weights from HF naming to MLX naming. + + Args: + weights: dict mapping HF weight names to arrays + reasoner_only: if True, strip generation/diffusion weights + + Returns: + dict mapping MLX weight names to arrays + """ + converted = {} + for name, tensor in weights.items(): + # Skip generation weights in reasoner mode + if reasoner_only and _is_generation_weight(name): + continue + + # Map the name + mlx_name = map_weight_name(name) + converted[mlx_name] = tensor + + return converted diff --git a/video/cosmos3/cosmos3/decode_audio.py b/video/cosmos3/cosmos3/decode_audio.py new file mode 100644 index 000000000..140d4edeb --- /dev/null +++ b/video/cosmos3/cosmos3/decode_audio.py @@ -0,0 +1,217 @@ +"""Standalone audio decode: load HuggingFace sound tokenizer weights and decode audio latents. + +Implements the Oobleck decoder with weight-normalized convolutions and SnakeBeta activations. +Loads weights directly from safetensors — no nn.Module tree needed. +""" + +import json +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + + +def _weight_norm_conv1d(x: mx.array, weight_v: mx.array, weight_g: mx.array, + bias: mx.array = None, stride: int = 1, + padding: int = 0, dilation: int = 1) -> mx.array: + """Apply weight-normalized Conv1d. + + Weight normalization: w = g * (v / ||v||) + PyTorch weight_v: [out_ch, in_ch, kernel] -> MLX: [out_ch, kernel, in_ch] + + Args: + x: [B, C, T] channels-first + weight_v: [out_ch, kernel, in_ch] MLX layout (already transposed) + weight_g: [out_ch, 1, 1] magnitude + bias: [out_ch] optional + """ + # Compute normalized weight: g * v / ||v|| + # Norm over (in_ch, kernel) dims — axes 1,2 in MLX layout + v_norm = mx.sqrt(mx.sum(weight_v * weight_v, axis=(1, 2), keepdims=True) + 1e-12) + weight = weight_g * weight_v / v_norm + + # MLX conv1d expects [B, T, C] channels-last + x_cl = mx.transpose(x, (0, 2, 1)) # [B, T, C] + out = mx.conv1d(x_cl, weight, stride=stride, padding=padding, dilation=dilation) + if bias is not None: + out = out + bias + return mx.transpose(out, (0, 2, 1)) # back to [B, C, T] + + +def _weight_norm_conv_transpose1d(x: mx.array, weight_v: mx.array, weight_g: mx.array, + bias: mx.array = None, stride: int = 1, + padding: int = 0) -> mx.array: + """Apply weight-normalized ConvTranspose1d. + + PyTorch ConvTranspose1d weight: [in_ch, out_ch, kernel] + MLX ConvTranspose1d weight: [out_ch, kernel, in_ch] + + Args: + x: [B, C, T] channels-first + weight_v: [out_ch, kernel, in_ch] MLX layout (already transposed) + weight_g: [in_ch, 1, 1] magnitude (norm over out_ch*kernel per input channel) + bias: [out_ch] optional + """ + # Weight norm: g * v / ||v|| — for ConvTranspose1d, PyTorch weight_g is [in_ch, 1, 1] + # norming over (out_ch, kernel). In MLX layout [out_ch, kernel, in_ch], that's axes (0, 1). + v_norm = mx.sqrt(mx.sum(weight_v * weight_v, axis=(0, 1), keepdims=True) + 1e-12) + g = weight_g.reshape(1, 1, -1) # [in_ch, 1, 1] -> [1, 1, in_ch] + weight = g * weight_v / v_norm + + # MLX conv_transpose1d: input [B, T, C_in], weight [C_out, K, C_in] + x_cl = mx.transpose(x, (0, 2, 1)) # [B, T, C] + out = mx.conv_transpose1d(x_cl, weight, stride=stride, padding=padding) + if bias is not None: + out = out + bias + return mx.transpose(out, (0, 2, 1)) # [B, C, T] + + +def _snake_beta(x: mx.array, alpha: mx.array, beta: mx.array, + log_scale: bool = True) -> mx.array: + """SnakeBeta activation: x + (1/b) * sin²(a*x). + + Args: + x: [B, C, T] channels-first + alpha: [1, C, 1] + beta: [1, C, 1] + log_scale: if True, alpha and beta are in log space + """ + if log_scale: + a = mx.exp(alpha) + b = mx.exp(beta) + else: + a = alpha + b = beta + return x + (1.0 / (b + 1e-9)) * mx.power(mx.sin(a * x), 2) + + +def decode_audio( + audio_latents: mx.array, + sound_tokenizer_dir: str, +) -> mx.array: + """Decode audio latents to waveform using HF sound tokenizer weights. + + Args: + audio_latents: [B, latent_dim, T_latent] or [latent_dim, T_latent] + Audio latents from the diffusion model (channels-first) + sound_tokenizer_dir: path to sound_tokenizer/ directory + + Returns: + [B, 2, T_audio] stereo waveform in [-1, 1] + """ + if audio_latents.ndim == 2: + audio_latents = mx.expand_dims(audio_latents, 0) + + snd_path = Path(sound_tokenizer_dir) + + with open(snd_path / "config.json") as f: + config = json.load(f) + + raw_weights = mx.load(str(snd_path / "diffusion_pytorch_model.safetensors")) + + # Extract and transpose decoder weights + weights = {} + for k, v in raw_weights.items(): + if not k.startswith("decoder."): + continue + name = k[len("decoder."):] + + # Transpose Conv1d weight_v: PyTorch [O, I, K] -> MLX [O, K, I] + if name.endswith(".weight_v") and v.ndim == 3: + # For ConvTranspose1d (conv_t1), PyTorch shape is [I, O, K] + # For Conv1d, PyTorch shape is [O, I, K] + if "conv_t" in name: + # ConvTranspose1d: [I, O, K] -> MLX [O, K, I] + v = mx.transpose(v, (1, 2, 0)) + else: + # Conv1d: [O, I, K] -> MLX [O, K, I] + v = mx.transpose(v, (0, 2, 1)) + + weights[name] = v.astype(mx.float32) + + log_scale = config.get("snake_logscale", True) + strides = config.get("dec_strides", [2, 4, 5, 6, 8]) + + x = audio_latents.astype(mx.float32) + + # conv1: Conv1d(64, 5120, 7, padding=3) with weight norm + x = _weight_norm_conv1d(x, weights["conv1.weight_v"], weights["conv1.weight_g"], + weights.get("conv1.bias"), padding=3) + mx.eval(x) + + # Decoder blocks (5 blocks, strides reversed = [8, 6, 5, 4, 2]) + for block_idx in range(5): + prefix = f"block.{block_idx}" + stride = strides[-(block_idx + 1)] # reverse order + + # Snake activation before upsample + x = _snake_beta(x, + weights[f"{prefix}.snake1.alpha"], + weights[f"{prefix}.snake1.beta"], + log_scale) + + # ConvTranspose1d upsample + pad = stride // 2 + x = _weight_norm_conv_transpose1d( + x, + weights[f"{prefix}.conv_t1.weight_v"], + weights[f"{prefix}.conv_t1.weight_g"], + weights.get(f"{prefix}.conv_t1.bias"), + stride=stride, padding=pad, + ) + mx.eval(x) + + # 3 residual units + for unit_idx in range(1, 4): + unit_prefix = f"{prefix}.res_unit{unit_idx}" + residual = x + + # Snake1 + Conv1 (dilated) + x = _snake_beta(x, + weights[f"{unit_prefix}.snake1.alpha"], + weights[f"{unit_prefix}.snake1.beta"], + log_scale) + + # Dilation pattern: 1, 3, 9 + dilation = 3 ** (unit_idx - 1) + pad_d = dilation * 3 # kernel=7, (7-1)//2 * dilation + x = _weight_norm_conv1d( + x, + weights[f"{unit_prefix}.conv1.weight_v"], + weights[f"{unit_prefix}.conv1.weight_g"], + weights.get(f"{unit_prefix}.conv1.bias"), + padding=pad_d, dilation=dilation, + ) + + # Snake2 + Conv2 (1x1) + x = _snake_beta(x, + weights[f"{unit_prefix}.snake2.alpha"], + weights[f"{unit_prefix}.snake2.beta"], + log_scale) + x = _weight_norm_conv1d( + x, + weights[f"{unit_prefix}.conv2.weight_v"], + weights[f"{unit_prefix}.conv2.weight_g"], + weights.get(f"{unit_prefix}.conv2.bias"), + ) + + # Trim to match residual + if x.shape[-1] > residual.shape[-1]: + x = x[..., :residual.shape[-1]] + elif x.shape[-1] < residual.shape[-1]: + residual = residual[..., :x.shape[-1]] + + x = x + residual + + mx.eval(x) + + # Final: snake1 + conv2 (output conv) + x = _snake_beta(x, weights["snake1.alpha"], weights["snake1.beta"], log_scale) + # conv2 is Conv1d(320, 2, 7, padding=3) — no bias key in the weights + x = _weight_norm_conv1d(x, weights["conv2.weight_v"], weights["conv2.weight_g"], + weights.get("conv2.bias"), padding=3) + mx.eval(x) + + x = mx.clip(x, -1.0, 1.0) + return x diff --git a/video/cosmos3/cosmos3/decode_vae.py b/video/cosmos3/cosmos3/decode_vae.py new file mode 100644 index 000000000..52a01afcf --- /dev/null +++ b/video/cosmos3/cosmos3/decode_vae.py @@ -0,0 +1,444 @@ +"""Standalone VAE decode: load HuggingFace VAE weights and decode latents to pixels. + +Implements the full WanResidualUpBlock decoder architecture including: +- Mid-block self-attention (WanAttentionBlock) +- Residual upsampling blocks with DupUp3D shortcut +- Learned Conv2d after nearest-exact interpolation (WanResample) +- Temporal upsampling via CausalConv3d +""" + +import json +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + + +def _transpose_conv3d_weight(w: mx.array) -> mx.array: + """Transpose Conv3D weight from PyTorch [O,I,D,H,W] to MLX-friendly [O,D,H,W,I].""" + return mx.transpose(w, (0, 2, 3, 4, 1)) + + +def _transpose_conv2d_weight(w: mx.array) -> mx.array: + """Transpose Conv2D weight from PyTorch [O,I,H,W] to MLX [O,H,W,I].""" + return mx.transpose(w, (0, 2, 3, 1)) + + +def _conv3d_forward(x: mx.array, weight: mx.array, bias: mx.array = None, + stride=(1,1,1), padding=(1,1,1), causal=True) -> mx.array: + """Run Conv3D via per-frame 2D decomposition. + + Args: + x: [B, T, H, W, C] channels-last + weight: [O, kD, kH, kW, I] MLX layout + bias: [O] optional + stride: (sD, sH, sW) + padding: (pD, pH, pW) + causal: if True, use causal temporal padding + """ + b, t, h, w, c = x.shape + o_ch, kd, kh, kw, i_ch = weight.shape + sd, sh, sw = stride + + # Causal temporal padding + if causal: + pad_t = 2 * padding[0] + pad_zeros = mx.zeros((b, pad_t, h, w, c), dtype=x.dtype) + x = mx.concatenate([pad_zeros, x], axis=1) + else: + # Symmetric padding + pad_t = padding[0] + if pad_t > 0: + pad_zeros = mx.zeros((b, pad_t, h, w, c), dtype=x.dtype) + x = mx.concatenate([pad_zeros, x, pad_zeros], axis=1) + + t_padded = x.shape[1] + t_out = (t_padded - kd) // sd + 1 + + outputs = [] + for ti in range(t_out): + t_start = ti * sd + accum = None + for d in range(kd): + frame = x[:, t_start + d] # [B, H, W, C] + w_2d = weight[:, d, :, :, :] # [O, kH, kW, I] + conv_out = mx.conv2d(frame, w_2d, stride=(sh, sw), + padding=(padding[1], padding[2])) + accum = conv_out if accum is None else accum + conv_out + if bias is not None: + accum = accum + bias + outputs.append(accum) + + return mx.stack(outputs, axis=1) + + +def _rms_norm(x: mx.array, gamma: mx.array, eps: float = 1e-6) -> mx.array: + """RMS norm with gamma. Gamma may have trailing singleton dims.""" + g = gamma.reshape(-1) + rms = mx.sqrt(mx.mean(x * x, axis=-1, keepdims=True) + eps) + return x / rms * g + + +def _resnet_block(x: mx.array, weights: dict, prefix: str) -> mx.array: + """Run one residual block.""" + residual = x + + x = _rms_norm(x, weights[f"{prefix}.norm1.gamma"]) + x = nn.silu(x) + x = _conv3d_forward(x, + weights[f"{prefix}.conv1.weight"], + weights.get(f"{prefix}.conv1.bias")) + mx.eval(x) + + x = _rms_norm(x, weights[f"{prefix}.norm2.gamma"]) + x = nn.silu(x) + x = _conv3d_forward(x, + weights[f"{prefix}.conv2.weight"], + weights.get(f"{prefix}.conv2.bias")) + mx.eval(x) + + # Skip connection (conv_shortcut for channel changes) + skip_key = f"{prefix}.conv_shortcut.weight" + if skip_key in weights: + residual = _conv3d_forward(residual, + weights[skip_key], + weights.get(f"{prefix}.conv_shortcut.bias"), + padding=(0, 0, 0)) + + return x + residual + + +def _attention_block(x: mx.array, weights: dict, prefix: str) -> mx.array: + """Run self-attention block (WanAttentionBlock). + + Operates per-frame: reshapes [B, T, H, W, C] to [B*T, H, W, C], + applies norm → qkv → attention → proj, then reshapes back. + + Args: + x: [B, T, H, W, C] channels-last + weights: dict with keys like "{prefix}.norm.gamma", "{prefix}.to_qkv.weight", etc. + """ + identity = x + b, t, h, w, c = x.shape + + # Reshape to per-frame: [B*T, H, W, C] + x = x.reshape(b * t, h, w, c) + + # RMS norm + x = _rms_norm(x, weights[f"{prefix}.norm.gamma"]) + + # QKV projection via 1x1 conv: [B*T, H, W, C] -> [B*T, H, W, 3*C] + # HF uses Conv2d(C, 3*C, 1) which is a 1x1 conv + qkv_w = weights[f"{prefix}.to_qkv.weight"] # [3C, 1, 1, C] in MLX layout + qkv_b = weights.get(f"{prefix}.to_qkv.bias") + qkv = mx.conv2d(x, qkv_w, padding=(0, 0)) + if qkv_b is not None: + qkv = qkv + qkv_b + + # Reshape for attention: [B*T, H*W, 3*C] -> split to q, k, v each [B*T, H*W, C] + qkv = qkv.reshape(b * t, h * w, 3 * c) + q, k, v = mx.split(qkv, 3, axis=-1) + + # Single-head scaled dot-product attention + # [B*T, 1, H*W, C] for SDPA + scale = c ** -0.5 + q = mx.expand_dims(q, 1) + k = mx.expand_dims(k, 1) + v = mx.expand_dims(v, 1) + attn = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + attn = attn.squeeze(1) # [B*T, H*W, C] + + # Reshape back to spatial: [B*T, H, W, C] + attn = attn.reshape(b * t, h, w, c) + + # Output projection via 1x1 conv + proj_w = weights[f"{prefix}.proj.weight"] + proj_b = weights.get(f"{prefix}.proj.bias") + out = mx.conv2d(attn, proj_w, padding=(0, 0)) + if proj_b is not None: + out = out + proj_b + + # Reshape back to [B, T, H, W, C] + out = out.reshape(b, t, h, w, c) + + return out + identity + + +def _nearest_upsample_2x(x_2d: mx.array) -> mx.array: + """Nearest-neighbor 2x upsampling for a 2D tensor [B, H, W, C].""" + # Repeat along H and W + b, h, w, c = x_2d.shape + x_2d = mx.repeat(x_2d, 2, axis=1) # [B, 2H, W, C] + x_2d = mx.repeat(x_2d, 2, axis=2) # [B, 2H, 2W, C] + return x_2d + + +def _wan_resample_upsample2d(x: mx.array, weights: dict, prefix: str) -> mx.array: + """WanResample upsample2d: nearest-exact interpolation + learned Conv2d per frame. + + Args: + x: [B, T, H, W, C] channels-last + weights: dict with keys like "{prefix}.resample.1.weight", "{prefix}.resample.1.bias" + """ + b, t, h, w, c = x.shape + + # Process per-frame + frames = [] + for ti in range(t): + frame = x[:, ti] # [B, H, W, C] + + # 2x nearest upsample + frame = _nearest_upsample_2x(frame) # [B, 2H, 2W, C] + + # Learned Conv2d(C, out_dim, 3, padding=1) + conv_w = weights[f"{prefix}.resample.1.weight"] # [O, kH, kW, I] + conv_b = weights.get(f"{prefix}.resample.1.bias") + frame = mx.conv2d(frame, conv_w, padding=(1, 1)) + if conv_b is not None: + frame = frame + conv_b + + frames.append(frame) + + return mx.stack(frames, axis=1) + + +def _wan_resample_upsample3d(x: mx.array, weights: dict, prefix: str) -> mx.array: + """WanResample upsample3d: temporal conv doubling + spatial upsample + conv. + + Temporal: CausalConv3d(C, 2C, (3,1,1)) → reshape to interleave → doubles T + Spatial: nearest 2x + Conv2d + + Args: + x: [B, T, H, W, C] channels-last + weights: dict with keys for time_conv and resample + """ + b, t, h, w, c = x.shape + + # Temporal upsampling via CausalConv3d(C, 2C, (3,1,1), padding=(1,0,0)) + time_conv_w = weights[f"{prefix}.time_conv.weight"] # [2C, kD, 1, 1, C] + time_conv_b = weights.get(f"{prefix}.time_conv.bias") + x_t = _conv3d_forward(x, time_conv_w, time_conv_b, + stride=(1, 1, 1), padding=(1, 0, 0), causal=True) + mx.eval(x_t) + # x_t: [B, T, H, W, 2C] + # Reshape to interleave frames: split channels in half, interleave along time + x_t = x_t.reshape(b, t, h, w, 2, c) + # Interleave: [B, T, H, W, 2, C] -> [B, 2T, H, W, C] + x_t = mx.transpose(x_t, (0, 1, 4, 2, 3, 5)) # [B, T, 2, H, W, C] + x_t = x_t.reshape(b, t * 2, h, w, c) + + # Spatial upsampling: nearest 2x + learned Conv2d per frame + x = _wan_resample_upsample2d(x_t, weights, prefix) + return x + + +def _dup_up_3d_residual(x: mx.array, in_c: int, out_c: int, + factor_t: int, factor_s: int) -> mx.array: + """DupUp3D residual shortcut: channel-repeat + reshape for skip connection. + + HF DupUp3D: repeat_interleave on channels, then reshape to interleave spatially. + + Args: + x: [B, T, H, W, C] channels-last (PyTorch: [B, C, T, H, W]) + in_c: input channels + out_c: output channels + factor_t: temporal upsample factor (1 or 2) + factor_s: spatial upsample factor (2) + """ + b, t, h, w, c = x.shape + factor = factor_t * factor_s * factor_s + repeats = out_c * factor // in_c + + # repeat_interleave on channel dim + x = mx.repeat(x, repeats, axis=-1) # [B, T, H, W, C*repeats] + + # Reshape to expose upsample factors + # PyTorch: [B, out_c, factor_t, factor_s, factor_s, T, H, W] + # MLX channels-last: [B, T, H, W, out_c, factor_t, factor_s, factor_s] + x = x.reshape(b, t, h, w, out_c, factor_t, factor_s, factor_s) + + # Permute to interleave spatial/temporal factors + # Target: [B, T*factor_t, H*factor_s, W*factor_s, out_c] + # Intermediate: [B, T, factor_t, H, factor_s, W, factor_s, out_c] + x = mx.transpose(x, (0, 1, 5, 2, 6, 3, 7, 4)) + x = x.reshape(b, t * factor_t, h * factor_s, w * factor_s, out_c) + + return x + + +def decode_latents( + latents: mx.array, + vae_dir: str, + latents_mean: list[float] = None, + latents_std: list[float] = None, +) -> mx.array: + """Decode latents to video frames using HuggingFace VAE weights directly. + + Implements the full WanResidualUpBlock decoder architecture. + + Args: + latents: [B, T, H, W, z_dim] denoised latents (channels-last) + vae_dir: path to vae/ directory with config.json and safetensors + latents_mean: per-channel mean for denormalization + latents_std: per-channel std for denormalization + + Returns: + [B, T_out, H_out, W_out, 3] decoded video frames in [0, 1] + """ + vae_path = Path(vae_dir) + + # Load config + with open(vae_path / "config.json") as f: + config = json.load(f) + + # Load weights + raw_weights = mx.load(str(vae_path / "diffusion_pytorch_model.safetensors")) + + # Extract decoder weights and post_quant_conv, transpose Conv3D and Conv2D + weights = {} + pqc_weight = None + pqc_bias = None + for k, v in raw_weights.items(): + if k == "post_quant_conv.weight": + # [O, I, 1, 1, 1] -> [O, 1, 1, 1, I] for Conv3D, but since kernel=1 + # it's effectively a linear transform on channels + pqc_weight = _transpose_conv3d_weight(v).astype(mx.bfloat16) + continue + if k == "post_quant_conv.bias": + pqc_bias = v.astype(mx.bfloat16) + continue + if not k.startswith("decoder."): + continue + name = k[len("decoder."):] + + # Transpose Conv3D weights: [O,I,D,H,W] -> [O,D,H,W,I] + if "conv" in name and name.endswith(".weight") and v.ndim == 5: + v = _transpose_conv3d_weight(v) + # Transpose Conv2D weights: [O,I,H,W] -> [O,H,W,I] + elif name.endswith(".weight") and v.ndim == 4: + v = _transpose_conv2d_weight(v) + + weights[name] = v.astype(mx.bfloat16) + + # Denormalize latents + if latents_mean is None: + latents_mean = config.get("latents_mean", [0.0] * latents.shape[-1]) + if latents_std is None: + latents_std = config.get("latents_std", [1.0] * latents.shape[-1]) + + mean = mx.array(latents_mean, dtype=latents.dtype) + std = mx.array(latents_std, dtype=latents.dtype) + z = latents * std + mean + + # Track input temporal dimension for T=1 vs T>1 behavior differences + input_t = z.shape[1] + + # post_quant_conv: 1x1x1 Conv3D that transforms latent channels before decoder + if pqc_weight is not None: + z = _conv3d_forward(z, pqc_weight, pqc_bias, + stride=(1, 1, 1), padding=(0, 0, 0), causal=False) + mx.eval(z) + + # conv_in + x = _conv3d_forward(z, weights["conv_in.weight"], weights.get("conv_in.bias")) + mx.eval(x) + + # mid_block: resnet[0] -> attention[0] -> resnet[1] + x = _resnet_block(x, weights, "mid_block.resnets.0") + mx.eval(x) + + # Mid-block attention (if weights exist) + attn_key = "mid_block.attentions.0.to_qkv.weight" + if attn_key in weights: + x = _attention_block(x, weights, "mid_block.attentions.0") + mx.eval(x) + + x = _resnet_block(x, weights, "mid_block.resnets.1") + mx.eval(x) + + # up_blocks (WanResidualUpBlock) + temporal_upsample = list(reversed(config.get("temperal_downsample", [False, True, True]))) + is_residual = config.get("is_residual", False) + + # Count up_blocks from weights + up_block_ids = set() + for k in weights: + if k.startswith("up_blocks."): + idx = int(k.split(".")[1]) + up_block_ids.add(idx) + num_up_blocks = len(up_block_ids) + + for block_idx in range(num_up_blocks): + prefix = f"up_blocks.{block_idx}" + + # Save input for residual shortcut + x_copy = x + + # Count resnets in this block + resnet_ids = set() + for k in weights: + if k.startswith(f"{prefix}.resnets."): + ri = int(k.split(".")[3]) + resnet_ids.add(ri) + num_resnets = len(resnet_ids) + + # Run resnets + for ri in range(num_resnets): + x = _resnet_block(x, weights, f"{prefix}.resnets.{ri}") + mx.eval(x) + + # Upsampling (WanResample) if upsampler weights exist + has_upsampler = any(k.startswith(f"{prefix}.upsampler.") for k in weights) + t_up = temporal_upsample[block_idx] if block_idx < len(temporal_upsample) else False + + if has_upsampler: + has_time_conv = f"{prefix}.upsampler.time_conv.weight" in weights + if has_time_conv and t_up and input_t > 1: + # Multi-frame: full 3D upsample (spatial + temporal) + x = _wan_resample_upsample3d(x, weights, f"{prefix}.upsampler") + else: + # Single-frame or no temporal: spatial-only 2D upsample. + # HF's cached path skips time_conv when feat_cache[idx] is None + # (first call); for T=1 single-pass this is the correct behavior. + x = _wan_resample_upsample2d(x, weights, f"{prefix}.upsampler") + mx.eval(x) + + # DupUp3D residual shortcut + add + # DupUp3D has no learned weights — it's a pure channel-reshape operation. + # For T=1: HF uses first_chunk=True which applies factor_t=2 then trims + # the first frame. For T>1: full temporal expansion, no trim. + if is_residual and has_upsampler: + in_c = x_copy.shape[-1] + out_c = x.shape[-1] + factor_t = 2 if t_up else 1 + factor_s = 2 + shortcut = _dup_up_3d_residual(x_copy, in_c, out_c, factor_t, factor_s) + # first_chunk trim for single-frame decode + if factor_t > 1 and input_t == 1: + shortcut = shortcut[:, factor_t - 1:, :, :, :] + x = x + shortcut + mx.eval(x) + + # norm_out + silu + conv_out + x = _rms_norm(x, weights["norm_out.gamma"]) + x = nn.silu(x) + x = _conv3d_forward(x, weights["conv_out.weight"], weights.get("conv_out.bias")) + mx.eval(x) + + # Unpatchify if needed (patch_size=2) + # HF packs channels as [C, p1, p2] and interleaves H with p2, W with p1. + # In channels-last: reshape to [B, T, H, W, C, p1, p2], then permute so + # p2 (dim6) interleaves with H and p1 (dim5) interleaves with W. + patch_size = config.get("patch_size", 2) + if patch_size > 1: + b, t, h, w, c_patch = x.shape + c = c_patch // (patch_size * patch_size) + x = x.reshape(b, t, h, w, c, patch_size, patch_size) + x = mx.transpose(x, (0, 1, 2, 6, 3, 5, 4)) # [B, T, H, p2, W, p1, C] + x = x.reshape(b, t, h * patch_size, w * patch_size, c) + + # Clamp to [0, 1] + x = (mx.clip(x, -1.0, 1.0) + 1.0) / 2.0 + + return x diff --git a/video/cosmos3/cosmos3/encode_vae.py b/video/cosmos3/cosmos3/encode_vae.py new file mode 100644 index 000000000..92465a44f --- /dev/null +++ b/video/cosmos3/cosmos3/encode_vae.py @@ -0,0 +1,628 @@ +"""Standalone VAE encode: load HuggingFace VAE weights and encode image/video to latents. + +Implements the Wan2.2 encoder architecture (inverse of decoder): +- Spatial patchification (pixel-space → patched input) +- WanResidualDownBlock: resnets + downsample + AvgDown3D residual shortcut +- Mid-block self-attention +- quant_conv → mean extraction (argmax mode) +- Per-channel normalization + +Multi-frame encoding uses chunked processing with feat_cache to match HF's +causal temporal convolution behavior: frame 0 alone, then 4 frames at a time. +""" + +import json +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from .decode_vae import ( + _transpose_conv3d_weight, + _transpose_conv2d_weight, + _conv3d_forward, + _rms_norm, + _resnet_block, + _attention_block, +) + +# Match HF CACHE_T = 2: each causal conv caches the last 2 frames of its input +CACHE_T = 2 + + +def _conv3d_forward_cached( + x: mx.array, + weight: mx.array, + bias: mx.array, + feat_cache: list, + feat_idx: list, + stride=(1, 1, 1), + padding=(1, 1, 1), +) -> mx.array: + """Conv3D with feat_cache for chunked temporal processing. + + Mirrors HF's WanCausalConv3d.forward(x, cache_x) + the cache bookkeeping + done in WanEncoder3d.forward / WanResidualBlock.forward. + + Cache protocol (matching HF exactly): + - Before the conv, save the last CACHE_T frames of the input as new cache + - If previous cache exists and current input has <2 temporal frames, + prepend the last frame of the previous cache + - Use the previous cache as temporal context instead of zero-padding + + Args: + x: [B, T, H, W, C] channels-last input + weight: [O, kD, kH, kW, I] conv weight + bias: [O] or None + feat_cache: mutable list of cached activations + feat_idx: mutable [int] index into feat_cache + stride: (sD, sH, sW) + padding: (pD, pH, pW) + """ + idx = feat_idx[0] + causal_pad_t = 2 * padding[0] + + # Cache bookkeeping: save last CACHE_T frames of input before conv + cache_x = x[:, -CACHE_T:] if x.shape[1] >= CACHE_T else x + if cache_x.shape[1] < 2 and feat_cache[idx] is not None: + # Prepend last frame from previous chunk's cache + cache_x = mx.concatenate([feat_cache[idx][:, -1:], cache_x], axis=1) + + # Build temporal context for the convolution + if feat_cache[idx] is not None and causal_pad_t > 0: + # Use cached frames instead of zero padding + prev_cache = feat_cache[idx] + # HF: x = torch.cat([cache_x, x], dim=2) where cache_x is the previous cache + x = mx.concatenate([prev_cache, x], axis=1) + # Reduce the zero padding needed + remaining_pad = causal_pad_t - prev_cache.shape[1] + if remaining_pad > 0: + b, _, h, w, c = x.shape + pad_zeros = mx.zeros((b, remaining_pad, h, w, c), dtype=x.dtype) + x = mx.concatenate([pad_zeros, x], axis=1) + elif causal_pad_t > 0: + # First chunk: standard causal zero-padding + b, _, h, w, c = x.shape + pad_zeros = mx.zeros((b, causal_pad_t, h, w, c), dtype=x.dtype) + x = mx.concatenate([pad_zeros, x], axis=1) + + # Update cache for next chunk + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + # Run the actual convolution (no additional causal padding - already applied) + b, t_padded, h, w, c = x.shape + o_ch, kd, kh, kw, i_ch = weight.shape + sd, sh, sw = stride + + t_out = (t_padded - kd) // sd + 1 + + outputs = [] + for ti in range(t_out): + t_start = ti * sd + accum = None + for d in range(kd): + frame = x[:, t_start + d] + w_2d = weight[:, d, :, :, :] + conv_out = mx.conv2d(frame, w_2d, stride=(sh, sw), + padding=(padding[1], padding[2])) + accum = conv_out if accum is None else accum + conv_out + if bias is not None: + accum = accum + bias + outputs.append(accum) + + return mx.stack(outputs, axis=1) + + +def _resnet_block_cached( + x: mx.array, + weights: dict, + prefix: str, + feat_cache: list, + feat_idx: list, +) -> mx.array: + """Residual block with feat_cache for chunked processing. + + Mirrors HF's WanResidualBlock.forward with feat_cache. + Cache slots: conv1 uses one slot, conv2 uses one slot. + conv_shortcut (1x1) doesn't need temporal caching (kernel_size=1). + """ + residual = x + + # First: norm -> silu -> conv1 (cached) + x = _rms_norm(x, weights[f"{prefix}.norm1.gamma"]) + x = nn.silu(x) + x = _conv3d_forward_cached( + x, weights[f"{prefix}.conv1.weight"], weights.get(f"{prefix}.conv1.bias"), + feat_cache, feat_idx, + ) + mx.eval(x) + + # Second: norm -> silu -> conv2 (cached) + x = _rms_norm(x, weights[f"{prefix}.norm2.gamma"]) + x = nn.silu(x) + x = _conv3d_forward_cached( + x, weights[f"{prefix}.conv2.weight"], weights.get(f"{prefix}.conv2.bias"), + feat_cache, feat_idx, + ) + mx.eval(x) + + # Skip connection (conv_shortcut for channel changes — 1x1, no temporal cache needed) + skip_key = f"{prefix}.conv_shortcut.weight" + if skip_key in weights: + residual = _conv3d_forward(residual, + weights[skip_key], + weights.get(f"{prefix}.conv_shortcut.bias"), + padding=(0, 0, 0)) + + return x + residual + + +def _wan_resample_downsample2d(x: mx.array, weights: dict, prefix: str) -> mx.array: + """WanResample downsample2d: learned Conv2d with stride 2 per frame. + + Args: + x: [B, T, H, W, C] channels-last + weights: dict with keys like "{prefix}.resample.1.weight", "{prefix}.resample.1.bias" + """ + b, t, h, w, c = x.shape + + frames = [] + for ti in range(t): + frame = x[:, ti] # [B, H, W, C] + # HF uses ZeroPad2d((0,1,0,1)) → asymmetric padding: right and bottom + # In channels-last: pad W (right) and H (bottom) by 1 + frame = mx.pad(frame, [(0, 0), (0, 1), (0, 1), (0, 0)]) + conv_w = weights[f"{prefix}.resample.1.weight"] + conv_b = weights.get(f"{prefix}.resample.1.bias") + frame = mx.conv2d(frame, conv_w, stride=(2, 2), padding=(0, 0)) + if conv_b is not None: + frame = frame + conv_b + frames.append(frame) + + return mx.stack(frames, axis=1) + + +def _wan_resample_downsample3d_cached( + x: mx.array, + weights: dict, + prefix: str, + feat_cache: list, + feat_idx: list, +) -> mx.array: + """WanResample downsample3d with feat_cache: spatial downsample + temporal time_conv. + + Mirrors HF's WanResample.forward for downsample3d mode with feat_cache: + - First run: spatial downsample only. Store the spatially-downsampled output + in feat_cache for next chunk. Skip time_conv. + - Subsequent runs: prepend cached last frame, apply time_conv (stride-2 temporal + downsample), then spatial downsample. + + Args: + x: [B, T, H, W, C] channels-last + """ + # Spatial downsample first (always runs) + x_spatial = _wan_resample_downsample2d(x, weights, prefix) + + idx = feat_idx[0] + if feat_cache[idx] is None: + # First chunk: just store, skip time_conv + feat_cache[idx] = x_spatial + feat_idx[0] += 1 + return x_spatial + else: + # Subsequent chunks: apply time_conv with cached context + cache_x = x_spatial[:, -1:] # save last frame for next chunk + + # Prepend last frame of previous cache + prev_last = feat_cache[idx][:, -1:] + x_with_context = mx.concatenate([prev_last, x_spatial], axis=1) + + # time_conv: WanCausalConv3d(dim, dim, (3,1,1), stride=(2,1,1), padding=(0,0,0)) + # kernel_size=(3,1,1), stride=(2,1,1), padding=(0,0,0) + # With causal padding = 2*0 = 0 (no causal pad for this conv) + # The cache provides the temporal context instead + tc_weight = weights[f"{prefix}.time_conv.weight"] + tc_bias = weights.get(f"{prefix}.time_conv.bias") + + x_out = _conv3d_forward( + x_with_context, tc_weight, tc_bias, + stride=(2, 1, 1), padding=(0, 0, 0), causal=False, + ) + + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x_out + + +def _nearest_downsample_2x(x_2d: mx.array) -> mx.array: + """Strided 2x spatial downsampling for a 2D tensor [B, H, W, C].""" + return x_2d[:, ::2, ::2, :] + + +def _avg_down_3d(x: mx.array, in_channels: int, out_channels: int, + factor_t: int, factor_s: int) -> mx.array: + """AvgDown3D: parameter-free channel-reshaping average-pool residual shortcut. + + Matches HF's AvgDown3D exactly. Reshapes input by interleaving spatial/temporal + factors into channels, groups, and averages to produce the target channel count + at downsampled resolution. + + Args: + x: [B, T, H, W, C] channels-last + in_channels: input channel count + out_channels: output channel count + factor_t: temporal downsampling factor (1 or 2) + factor_s: spatial downsampling factor (1 or 2) + + Returns: + [B, T//factor_t, H//factor_s, W//factor_s, out_channels] + """ + factor = factor_t * factor_s * factor_s + group_size = in_channels * factor // out_channels + + b, t, h, w, c = x.shape + + # Pad temporal if needed + pad_t = (factor_t - t % factor_t) % factor_t + if pad_t > 0: + x = mx.pad(x, [(0, 0), (pad_t, 0), (0, 0), (0, 0), (0, 0)]) + t = t + pad_t + + t_out = t // factor_t + h_out = h // factor_s + w_out = w // factor_s + + # [B, T, H, W, C] -> [B, T//ft, ft, H//fs, fs, W//fs, fs, C] + x = x.reshape(b, t_out, factor_t, h_out, factor_s, w_out, factor_s, c) + # Move factors next to C: [B, T', H', W', C, ft, fs, fs] + x = mx.transpose(x, (0, 1, 3, 5, 7, 2, 4, 6)) + # Collapse factors into channels: [B, T', H', W', C*factor] + x = x.reshape(b, t_out, h_out, w_out, c * factor) + # Group and average: [B, T', H', W', out_channels, group_size] + x = x.reshape(b, t_out, h_out, w_out, out_channels, group_size) + x = mx.mean(x, axis=-1) + + return x + + +def _patchify_input(x: mx.array, patch_size: int = 2) -> mx.array: + """Patchify pixel-space input: pack spatial patches into channels. + + Input: [B, T, H, W, 3] + Output: [B, T, H//p, W//p, 3*p*p] + + Matches HF's patchify() exactly: + HF: [B,C,T,H,W] -> view [B,C,T,H//p,p,W//p,p] -> permute [B,C,p,p,T,H//p,W//p] + -> view [B, C*p*p, T, H//p, W//p] + + In channels-last: [B,T,H,W,C] -> [B,T,H//p,p,W//p,p,C] + -> permute to [B,T,H//p,W//p,C,p,p] -> [B,T,H//p,W//p,C*p*p] + """ + b, t, h, w, c = x.shape + p = patch_size + # [B, T, H, W, C] -> [B, T, H//p, p, W//p, p, C] + x = x.reshape(b, t, h // p, p, w // p, p, c) + # From [B,T,H//p,p_h,W//p,p_w,C] -> [B,T,H//p,W//p,C,p_w,p_h] + x = mx.transpose(x, (0, 1, 2, 4, 6, 5, 3)) + x = x.reshape(b, t, h // p, w // p, c * p * p) + return x + + +def _count_encoder_cache_slots(weights: dict, config: dict) -> int: + """Count the number of feat_cache slots needed for the encoder. + + Each WanCausalConv3d in the encoder gets one cache slot. + Matches HF's sum(isinstance(m, WanCausalConv3d) for m in encoder.modules()). + """ + count = 0 + is_residual = config.get("is_residual", False) + dim_mult = config.get("dim_mult", [1, 2, 4, 4]) + num_res_blocks = config.get("num_res_blocks", 2) + temporal_downsample = config.get("temperal_downsample", [False, True, True]) + + # conv_in: 1 CausalConv3d + count += 1 + + # down_blocks + base_dim = config.get("base_dim", 160) + dims = [base_dim * u for u in [1] + dim_mult] + + for block_idx in range(len(dim_mult)): + in_dim = dims[block_idx] + out_dim = dims[block_idx + 1] + down_flag = block_idx != len(dim_mult) - 1 + + if is_residual: + # WanResidualDownBlock: num_res_blocks resnets + for _ in range(num_res_blocks): + # Each WanResidualBlock: conv1 + conv2 = 2 CausalConv3d + count += 2 + # conv_shortcut (1x1) exists when in_dim != out_dim but uses + # non-cached _conv3d_forward (kernel_size=1, no temporal context + # needed), so it does NOT consume a cache slot. + in_dim = out_dim + + # downsampler + if down_flag: + t_down = temporal_downsample[block_idx] if block_idx < len(temporal_downsample) else False + if t_down: + # downsample3d: time_conv is a CausalConv3d + count += 1 + # Note: the spatial Conv2d in the resampler is NOT a CausalConv3d + else: + # Non-residual path (shouldn't apply for Cosmos3) + for _ in range(num_res_blocks): + count += 2 + if in_dim != out_dim: + count += 1 + in_dim = out_dim + if down_flag: + t_down = temporal_downsample[block_idx] if block_idx < len(temporal_downsample) else False + if t_down: + count += 1 + + # mid_block: num_layers=1 means 2 resnets (initial + 1 per layer) + # Each resnet: conv1 + conv2 = 2 CausalConv3d + count += 2 # resnets[0] + count += 2 # resnets[1] + + # conv_out: 1 CausalConv3d + count += 1 + + return count + + +def _encoder_forward( + x: mx.array, + weights: dict, + config: dict, + feat_cache: list | None = None, + feat_idx: list | None = None, +) -> mx.array: + """Run the encoder forward pass, optionally with feat_cache for chunked processing. + + Args: + x: [B, T, H, W, C] already-patchified input + weights: encoder weights dict (keys without 'encoder.' prefix) + config: VAE config dict + feat_cache: list of cached activations (None for single-pass mode) + feat_idx: mutable [int] index into feat_cache (None for single-pass mode) + """ + cached = feat_cache is not None + + # conv_in + if cached: + x = _conv3d_forward_cached( + x, weights["conv_in.weight"], weights.get("conv_in.bias"), + feat_cache, feat_idx, + ) + else: + x = _conv3d_forward(x, weights["conv_in.weight"], weights.get("conv_in.bias")) + mx.eval(x) + + # Architecture config + is_residual = config.get("is_residual", False) + temporal_downsample = config.get("temperal_downsample", [False, True, True]) + dim_mult = config.get("dim_mult", [1, 2, 4, 4]) + base_dim = config.get("base_dim", 160) + + dims = [base_dim * u for u in [1] + dim_mult] + num_down_blocks = len(dim_mult) + + for block_idx in range(num_down_blocks): + prefix = f"down_blocks.{block_idx}" + in_dim = dims[block_idx] + out_dim = dims[block_idx + 1] + + # Save input for residual shortcut (WanResidualDownBlock) + x_copy = x if is_residual else None + + # Count resnets + resnet_ids = set() + for k in weights: + if k.startswith(f"{prefix}.resnets."): + ri = int(k.split(".")[3]) + resnet_ids.add(ri) + num_resnets = len(resnet_ids) + + for ri in range(num_resnets): + if cached: + x = _resnet_block_cached(x, weights, f"{prefix}.resnets.{ri}", + feat_cache, feat_idx) + else: + x = _resnet_block(x, weights, f"{prefix}.resnets.{ri}") + mx.eval(x) + + # Downsampling + has_downsampler = any(k.startswith(f"{prefix}.downsampler.") for k in weights) + down_flag = block_idx != len(dim_mult) - 1 + t_down = temporal_downsample[block_idx] if block_idx < len(temporal_downsample) and down_flag else False + + if has_downsampler: + has_time_conv = f"{prefix}.downsampler.time_conv.weight" in weights + if has_time_conv and t_down: + if cached: + x = _wan_resample_downsample3d_cached( + x, weights, f"{prefix}.downsampler", feat_cache, feat_idx) + else: + # Single-pass: skip time_conv (HF behavior when feat_cache is None) + x = _wan_resample_downsample2d(x, weights, f"{prefix}.downsampler") + else: + x = _wan_resample_downsample2d(x, weights, f"{prefix}.downsampler") + mx.eval(x) + + # AvgDown3D residual shortcut (WanResidualDownBlock) + if is_residual: + factor_t = 2 if t_down else 1 + factor_s = 2 if down_flag else 1 + shortcut = _avg_down_3d(x_copy, in_dim, out_dim, factor_t, factor_s) + x = x + shortcut + mx.eval(x) + + # mid_block: resnet[0] -> attention[0] -> resnet[1] + if cached: + x = _resnet_block_cached(x, weights, "mid_block.resnets.0", feat_cache, feat_idx) + else: + x = _resnet_block(x, weights, "mid_block.resnets.0") + mx.eval(x) + + attn_key = "mid_block.attentions.0.to_qkv.weight" + if attn_key in weights: + x = _attention_block(x, weights, "mid_block.attentions.0") + mx.eval(x) + + if cached: + x = _resnet_block_cached(x, weights, "mid_block.resnets.1", feat_cache, feat_idx) + else: + x = _resnet_block(x, weights, "mid_block.resnets.1") + mx.eval(x) + + # norm_out + silu + conv_out + x = _rms_norm(x, weights["norm_out.gamma"]) + x = nn.silu(x) + if cached: + x = _conv3d_forward_cached( + x, weights["conv_out.weight"], weights.get("conv_out.bias"), + feat_cache, feat_idx, + ) + else: + x = _conv3d_forward(x, weights["conv_out.weight"], weights.get("conv_out.bias")) + mx.eval(x) + + return x + + +def _load_encoder_weights(vae_dir: str): + """Load and prepare encoder weights and config from VAE directory. + + Returns: + (weights, qc_weight, qc_bias, config) tuple + """ + vae_path = Path(vae_dir) + + with open(vae_path / "config.json") as f: + config = json.load(f) + + raw_weights = mx.load(str(vae_path / "diffusion_pytorch_model.safetensors")) + + weights = {} + qc_weight = None + qc_bias = None + for k, v in raw_weights.items(): + if k == "quant_conv.weight": + qc_weight = _transpose_conv3d_weight(v).astype(mx.bfloat16) + continue + if k == "quant_conv.bias": + qc_bias = v.astype(mx.bfloat16) + continue + if not k.startswith("encoder."): + continue + name = k[len("encoder."):] + + if "conv" in name and name.endswith(".weight") and v.ndim == 5: + v = _transpose_conv3d_weight(v) + elif name.endswith(".weight") and v.ndim == 4: + v = _transpose_conv2d_weight(v) + + weights[name] = v.astype(mx.bfloat16) + + return weights, qc_weight, qc_bias, config + + +def _prepare_input(video: np.ndarray | mx.array) -> mx.array: + """Prepare video input: normalize to [-1, 1], ensure [1, T, H, W, 3] shape.""" + if isinstance(video, np.ndarray): + if video.dtype == np.uint8: + video = video.astype(np.float32) / 255.0 + x = mx.array(video) + else: + x = video + + if x.ndim == 3: + x = mx.expand_dims(mx.expand_dims(x, 0), 0) # [H,W,3] -> [1, 1, H, W, 3] + elif x.ndim == 4: + x = mx.expand_dims(x, 0) # [T,H,W,3] -> [1, T, H, W, 3] + x = x * 2.0 - 1.0 # [0,1] -> [-1,1] + x = x.astype(mx.bfloat16) + return x + + +def encode_video( + video: np.ndarray | mx.array, + vae_dir: str, +) -> mx.array: + """Encode a video (or single image) to normalized VAE latents. + + Uses chunked encoding matching HF's _encode(): frame 0 processed alone, + then 4 frames at a time, with feat_cache propagating causal temporal + convolution state between chunks. + + For single-frame input, runs without chunking (no feat_cache overhead). + + Args: + video: [T, H, W, 3] or [H, W, 3] uint8/float32 in [0, 1]. + vae_dir: path to vae/ directory with config.json and safetensors + + Returns: + [1, T_lat, H//16, W//16, z_dim] normalized latents (channels-last) + """ + weights, qc_weight, qc_bias, config = _load_encoder_weights(vae_dir) + x = _prepare_input(video) + + num_frames = x.shape[1] + + # Patchify + patch_size = config.get("patch_size", 2) + if patch_size is not None and patch_size > 1: + x = _patchify_input(x, patch_size) + if num_frames == 1: + # Single frame: no chunking needed, no feat_cache overhead + out = _encoder_forward(x, weights, config) + else: + # Multi-frame: chunked encoding matching HF's _encode() + # iter_ = 1 + (num_frame - 1) // 4 + # Chunk 0: frame 0 alone + # Chunk i (i>0): frames 1+4*(i-1) : 1+4*i + num_cache_slots = _count_encoder_cache_slots(weights, config) + feat_cache = [None] * num_cache_slots + feat_idx = [0] + + # Chunk 0: first frame + feat_idx[0] = 0 + out = _encoder_forward(x[:, :1], weights, config, feat_cache, feat_idx) + + # Subsequent chunks: 4 frames at a time + num_iter = 1 + (num_frames - 1 + 3) // 4 # ceil division: process all frames + for i in range(1, num_iter): + feat_idx[0] = 0 + start = 1 + 4 * (i - 1) + end = min(1 + 4 * i, num_frames) + chunk = x[:, start:end] + chunk_out = _encoder_forward(chunk, weights, config, feat_cache, feat_idx) + out = mx.concatenate([out, chunk_out], axis=1) + mx.eval(out) + + # quant_conv + if qc_weight is not None: + out = _conv3d_forward(out, qc_weight, qc_bias, + stride=(1, 1, 1), padding=(0, 0, 0), causal=False) + mx.eval(out) + + # Extract mean (argmax mode): first z_dim channels + z_dim = config.get("z_dim", 48) + mu = out[..., :z_dim] + + # Normalize: z_norm = (mu - mean) * inv_std + latents_mean = config.get("latents_mean", [0.0] * z_dim) + latents_std = config.get("latents_std", [1.0] * z_dim) + mean = mx.array(latents_mean, dtype=mu.dtype) + inv_std = 1.0 / mx.array(latents_std, dtype=mu.dtype) + z_norm = (mu - mean) * inv_std + + return z_norm + + +# Backward-compatible alias +encode_image = encode_video diff --git a/video/cosmos3/cosmos3/load.py b/video/cosmos3/cosmos3/load.py new file mode 100644 index 000000000..76cc5e854 --- /dev/null +++ b/video/cosmos3/cosmos3/load.py @@ -0,0 +1,136 @@ +"""Load Cosmos 3 Nano weights from HuggingFace format into MLX models. + +Handles: +- Loading config.json to build model configs +- Loading sharded safetensors into MLX arrays +- Filtering for reasoner-only weights +- Mapping weights into Cosmos3Model +""" + +import json +from pathlib import Path +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten + +from .convert import convert_weights +from .model import Cosmos3Config, Cosmos3Model + + +def load_transformer_config(model_dir: str | Path) -> Cosmos3Config: + """Load transformer config from HuggingFace model directory.""" + config_path = Path(model_dir) / "transformer" / "config.json" + if not config_path.exists(): + raise FileNotFoundError( + f"Model config not found at {config_path}. " + f"Download the model first: hf download nvidia/Cosmos3-Nano --local-dir {model_dir}" + ) + with open(config_path) as f: + cfg = json.load(f) + + rope_scaling = cfg.get("rope_scaling", {}) + mrope_section = rope_scaling.get("mrope_section", [24, 20, 20]) + + return Cosmos3Config( + hidden_size=cfg["hidden_size"], + num_hidden_layers=cfg["num_hidden_layers"], + num_attention_heads=cfg["num_attention_heads"], + num_key_value_heads=cfg["num_key_value_heads"], + head_dim=cfg["head_dim"], + intermediate_size=cfg["intermediate_size"], + vocab_size=cfg["vocab_size"], + rms_norm_eps=cfg["rms_norm_eps"], + rope_theta=cfg["rope_theta"], + mrope_section=mrope_section, + max_position_embeddings=cfg["max_position_embeddings"], + ) + + +def _load_safetensors_shards(directory: Path) -> dict[str, mx.array]: + """Load all safetensors shards from a directory.""" + weights = {} + for shard in sorted(directory.glob("*.safetensors")): + shard_weights = mx.load(str(shard)) + weights.update(shard_weights) + if not weights: + raise FileNotFoundError( + f"No safetensors files found in {directory}. " + f"Download the model first: hf download nvidia/Cosmos3-Nano --local-dir weights/Cosmos3-Nano" + ) + return weights + + +def load_transformer( + model_dir: str | Path, + reasoner_only: bool = False, + dtype: mx.Dtype = mx.bfloat16, +) -> Cosmos3Model: + """Load Cosmos 3 transformer with weights. + + Args: + model_dir: path to HuggingFace model directory + reasoner_only: strip generation weights for smaller memory + dtype: target dtype (default bfloat16) + + Returns: + Cosmos3Model with loaded weights + """ + model_dir = Path(model_dir) + + # Load config and create model + config = load_transformer_config(model_dir) + model = Cosmos3Model(config) + + # Load and convert weights + raw_weights = _load_safetensors_shards(model_dir / "transformer") + weights = convert_weights(raw_weights, reasoner_only=reasoner_only) + + # Cast to target dtype + weights = {k: v.astype(dtype) for k, v in weights.items()} + + # Weight keys that are in the checkpoint but have no matching nn.Module + # parameter (action modality projections use a non-standard structure) + EXPECTED_EXTRA_KEYS = {"action_proj_in.fc.weight", "action_proj_out.weight"} + + model_params = set(k for k, _ in tree_flatten(model.parameters())) + weight_keys = set(weights.keys()) + skipped = weight_keys - model_params + missing = model_params - weight_keys + unexpected_extra = skipped - EXPECTED_EXTRA_KEYS + + model.load_weights(list(weights.items()), strict=False) + + if unexpected_extra: + print(f" Note: {len(unexpected_extra)} unexpected extra weight keys") + if missing: + raise RuntimeError( + f"Missing {len(missing)} required model parameters in weights. " + f"First 5: {sorted(missing)[:5]}. " + f"The checkpoint may be incomplete or from an incompatible model version." + ) + + return model + + +def load_tokenizer(model_dir: str | Path): + """Load the Qwen2 tokenizer. + + Returns a HuggingFace tokenizer — we use it directly since + tokenization is CPU-only and doesn't need MLX. + """ + from transformers import AutoTokenizer + + model_dir = Path(model_dir) + + # Try text_tokenizer subdirectory first (HF Cosmos3 layout), + # then fall back to root directory + tokenizer_dir = model_dir / "text_tokenizer" + if not tokenizer_dir.exists(): + tokenizer_dir = model_dir + + return AutoTokenizer.from_pretrained( + str(tokenizer_dir), + trust_remote_code=False, + ) diff --git a/video/cosmos3/cosmos3/model.py b/video/cosmos3/cosmos3/model.py new file mode 100644 index 000000000..48b65d1d9 --- /dev/null +++ b/video/cosmos3/cosmos3/model.py @@ -0,0 +1,449 @@ +"""Cosmos 3 Mixture-of-Transformers model for MLX. + +Dual-pathway architecture: AR reasoner (text understanding) and +bidirectional diffuser (video/audio generation). +""" + +from dataclasses import dataclass, field +from typing import Optional, Tuple + +import mlx.core as mx +import mlx.nn as nn + +from .attention import Cosmos3Attention + + +@dataclass +class Cosmos3Config: + """Configuration for Cosmos 3 Nano.""" + + hidden_size: int = 4096 + num_hidden_layers: int = 36 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int = 128 + intermediate_size: int = 12288 + vocab_size: int = 151936 + rms_norm_eps: float = 1e-6 + rope_theta: float = 5_000_000.0 + mrope_section: list[int] = field(default_factory=lambda: [24, 20, 20]) + max_position_embeddings: int = 262144 + + +class Cosmos3MLP(nn.Module): + """GLU/SiLU gated feed-forward network.""" + + def __init__(self, hidden_size: int, intermediate_size: int): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def __call__(self, x: mx.array) -> mx.array: + return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class Cosmos3DecoderLayer(nn.Module): + """Single MoT decoder layer with understanding + generation pathways.""" + + def __init__(self, config: Cosmos3Config): + super().__init__() + + # Understanding pathway norms + self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + # Dual-pathway attention + self.self_attn = Cosmos3Attention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + head_dim=config.head_dim, + mrope_section=config.mrope_section, + rope_theta=config.rope_theta, + rms_norm_eps=config.rms_norm_eps, + ) + + # Shared MLP (understanding pathway) + self.mlp = Cosmos3MLP(config.hidden_size, config.intermediate_size) + + # Generation pathway norms and MLP + self.input_layernorm_moe_gen = nn.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm_moe_gen = nn.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.mlp_moe_gen = Cosmos3MLP(config.hidden_size, config.intermediate_size) + + def __call__( + self, + hidden_states: mx.array, + position_ids: mx.array, + cache: Optional[Tuple[mx.array, mx.array]] = None, + ) -> Tuple[mx.array, Optional[Tuple[mx.array, mx.array]]]: + """Forward pass for understanding pathway only. + + Args: + hidden_states: [batch, seq_len, hidden_size] + position_ids: [3, batch, seq_len] + cache: optional KV cache + + Returns: + (output, updated_cache) + """ + # Pre-norm + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # Self-attention (understanding pathway only) + attn_out, _, new_cache, _ = self.self_attn( + hidden_states=hidden_states, + position_ids=position_ids, + understanding_mask=None, + generation_tokens=None, + cache=cache, + ) + + # Residual + hidden_states = residual + attn_out + + # Post-attention norm + MLP + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states, new_cache + + def forward_with_generation( + self, + und_hidden: mx.array, + gen_hidden: mx.array, + position_ids: mx.array, + ) -> Tuple[mx.array, mx.array, Tuple[mx.array, mx.array]]: + """Forward pass with both understanding and generation pathways. + + Args: + und_hidden: [batch, und_len, hidden_size] understanding tokens + gen_hidden: [batch, gen_len, hidden_size] generation tokens + position_ids: [3, batch, und_len + gen_len] + + Returns: + (und_output, gen_output, und_kv) where und_kv is cached (keys, values) + """ + # Understanding pre-norm + und_residual = und_hidden + und_normed = self.input_layernorm(und_hidden) + + # Generation pre-norm + gen_residual = gen_hidden + gen_normed = self.input_layernorm_moe_gen(gen_hidden) + + # Dual-pathway attention + und_attn, gen_attn, _, und_kv = self.self_attn( + hidden_states=und_normed, + position_ids=position_ids, + understanding_mask=None, + generation_tokens=gen_normed, + ) + + # Understanding residual + MLP + und_hidden = und_residual + und_attn + und_residual = und_hidden + und_hidden = self.post_attention_layernorm(und_hidden) + und_hidden = self.mlp(und_hidden) + und_hidden = und_residual + und_hidden + + # Generation residual + MLP + gen_hidden = gen_residual + gen_attn + gen_residual = gen_hidden + gen_hidden = self.post_attention_layernorm_moe_gen(gen_hidden) + gen_hidden = self.mlp_moe_gen(gen_hidden) + gen_hidden = gen_residual + gen_hidden + + return und_hidden, gen_hidden, und_kv + + def forward_generation_cached( + self, + gen_hidden: mx.array, + und_kv: Tuple[mx.array, mx.array], + position_ids: mx.array, + und_len: int, + ) -> mx.array: + """Forward pass for generation pathway only, using cached understanding K/V. + + Skips the entire understanding pathway. Used for denoising steps 1+ when + the text tokens haven't changed. + """ + gen_residual = gen_hidden + gen_normed = self.input_layernorm_moe_gen(gen_hidden) + + gen_attn = self.self_attn.generation_only_forward( + gen_normed, und_kv, position_ids, und_len, + ) + + gen_hidden = gen_residual + gen_attn + gen_residual = gen_hidden + gen_hidden = self.post_attention_layernorm_moe_gen(gen_hidden) + gen_hidden = self.mlp_moe_gen(gen_hidden) + gen_hidden = gen_residual + gen_hidden + + return gen_hidden + + +class Cosmos3Model(nn.Module): + """Cosmos 3 Mixture-of-Transformers model. + + Dual-pathway: AR reasoner for text understanding, bidirectional + diffuser for video/audio generation. + """ + + def __init__(self, config: Cosmos3Config): + super().__init__() + self.config = config + + # Token embeddings + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + + # Transformer layers + self.layers = [ + Cosmos3DecoderLayer(config) for _ in range(config.num_hidden_layers) + ] + + # Final norms (understanding + generation) + self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.norm_moe_gen = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # LM head + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Generation pathway: latent projections + # patch_latent_dim = latent_channel(48) * latent_patch_size(2)^2 = 192 + # These are loaded from weights; default matches Cosmos3-Nano + self.proj_in = nn.Linear(192, config.hidden_size, bias=True) + self.proj_out = nn.Linear(config.hidden_size, 192, bias=True) + + # Audio projections + sound_dim = 64 + self.audio_proj_in = nn.Linear(sound_dim, config.hidden_size, bias=True) + self.audio_proj_out = nn.Linear(config.hidden_size, sound_dim, bias=True) + + # Timestep embedder + from .timestep import TimestepEmbedding + self.time_embedder = TimestepEmbedding(config.hidden_size) + + # Modality embeddings + self.audio_modality_embed = mx.zeros((config.hidden_size,)) + self.action_modality_embed = mx.zeros((config.hidden_size,)) + + def diffusion_forward( + self, + input_ids: mx.array, + gen_tokens: mx.array, + timestep: mx.array, + grid_t: int = 1, + grid_h: int = 1, + grid_w: int = 1, + audio_tokens: Optional[mx.array] = None, + noisy_frame_indexes: Optional[list[int]] = None, + ) -> mx.array | tuple[mx.array, mx.array]: + """Forward pass for diffusion generation (one denoising step). + + Runs both understanding (text) and generation (latent) pathways + through the dual-pathway MoT transformer. Returns velocity prediction + for the generation tokens, and optionally for audio tokens. + + Args: + input_ids: [batch, text_len] text token IDs + gen_tokens: [batch, num_patches, patch_latent_dim] patchified latents + timestep: [batch] current diffusion timestep + grid_t: temporal grid size (number of latent frames) + grid_h: height grid size (latent height / patch_size) + grid_w: width grid size (latent width / patch_size) + audio_tokens: optional [batch, sound_len, sound_dim] audio latents + noisy_frame_indexes: which temporal frames are noisy (get timestep + embedding). None = all frames are noisy (t2v default). For i2v + with frame 0 conditioned: [1, 2, 3, ...]. + + Returns: + If audio_tokens is None: + [batch, num_patches, patch_latent_dim] velocity prediction + If audio_tokens is provided: + (vision_velocity, audio_velocity) tuple + """ + batch = input_ids.shape[0] + text_len = input_ids.shape[1] + num_patches = gen_tokens.shape[1] + + # Embed text tokens + und_h = self.embed_tokens(input_ids) + + # Project generation tokens into hidden space + gen_h = self.proj_in(gen_tokens) + + # Add timestep embedding only to noisy frame tokens. + # HF Cosmos3OmniTransformer._apply_timestep_embeds_to_noisy_tokens uses + # scatter_add to selectively add timestep to noisy positions only. + # Conditioned frames (e.g. frame 0 in i2v) get no timestep signal. + scaled_t = timestep * 0.001 + t_emb = self.time_embedder(scaled_t) # [batch, hidden_size] + + if noisy_frame_indexes is None: + # All frames noisy (t2v): add timestep to everything + gen_h = gen_h + mx.expand_dims(t_emb, 1) + else: + # Selective: build a mask for noisy token positions + spatial_tokens = grid_h * grid_w + noisy_mask = mx.zeros((num_patches,), dtype=gen_h.dtype) + for fi in noisy_frame_indexes: + start = fi * spatial_tokens + end = start + spatial_tokens + noisy_mask = noisy_mask.at[start:end].add(mx.ones((spatial_tokens,), dtype=gen_h.dtype)) + # [1, num_patches, 1] mask * [1, 1, hidden_size] timestep + gen_h = gen_h + noisy_mask.reshape(1, num_patches, 1) * mx.expand_dims(t_emb, 1) + + # Build 3D mRoPE position IDs + # Text tokens: all 3 axes share monotonically increasing IDs + # HF uses get_3d_mrope_ids_text_tokens which produces float positions + text_pos = mx.arange(text_len, dtype=mx.float32)[None, :] # [1, text_len] + text_position_ids = mx.stack([text_pos, text_pos, text_pos]) # [3, 1, text_len] + + # Generation tokens: FPS-modulated temporal positions + # HF: scaled_t = frame_index / tps * base_tps + temporal_offset + # where tps = fps / temporal_compression_factor, base_tps = base_fps / base_tcf + temporal_margin = 15000 + temporal_offset = float(text_len + temporal_margin) + + # Video FPS modulation: fps=24, temporal_compression_factor=4, base_fps=24 + fps = 24.0 + video_tcf = 4 # temporal compression factor for video VAE + base_fps = 24.0 + tps = fps / video_tcf # 6.0 tokens per second + base_tps = base_fps / video_tcf # 6.0 + frame_indices = mx.arange(grid_t, dtype=mx.float32) + scaled_t = frame_indices / tps * base_tps + temporal_offset + t_idx = mx.broadcast_to(scaled_t.reshape(-1, 1), (grid_t, grid_h * grid_w)).reshape(1, -1) + + h_idx = mx.arange(grid_h, dtype=mx.float32).reshape(1, -1, 1) + h_idx = mx.broadcast_to(h_idx, (grid_t, grid_h, grid_w)).reshape(1, -1) + + w_idx = mx.arange(grid_w, dtype=mx.float32).reshape(1, 1, -1) + w_idx = mx.broadcast_to(w_idx, (grid_t, grid_h, grid_w)).reshape(1, -1) + + gen_position_ids = mx.stack([t_idx, h_idx, w_idx]) # [3, 1, num_patches] + + # Handle audio tokens + if audio_tokens is not None: + sound_len = audio_tokens.shape[1] + + # Project audio tokens + add modality embedding + timestep + audio_h = self.audio_proj_in(audio_tokens) + audio_h = audio_h + self.audio_modality_embed + audio_h = audio_h + mx.expand_dims(t_emb, 1) + + # Audio mRoPE: temporal siblings with video, grid_h=1, grid_w=1 + # Audio: temporal_compression_factor=1, so tps = fps/1 = 24 + # base_tps for audio uses the audio compression factor (1), not video's (4) + audio_tps = fps / 1.0 # 24 tokens per second + audio_base_tps = base_fps / 1.0 # 24.0 (base_fps / audio_tcf) + audio_frame_indices = mx.arange(sound_len, dtype=mx.float32) + audio_scaled_t = audio_frame_indices / audio_tps * audio_base_tps + temporal_offset + audio_t_idx = audio_scaled_t.reshape(1, -1) + audio_h_idx = mx.zeros((1, sound_len), dtype=mx.float32) + audio_w_idx = mx.zeros((1, sound_len), dtype=mx.float32) + audio_position_ids = mx.stack([audio_t_idx, audio_h_idx, audio_w_idx]) + + # Concatenate video + audio in generation pathway + gen_h = mx.concatenate([gen_h, audio_h], axis=1) + gen_position_ids = mx.concatenate([gen_position_ids, audio_position_ids], axis=2) + + # Concatenate text + generation position IDs + position_ids = mx.concatenate([text_position_ids, gen_position_ids], axis=2) + + # Forward through all layers with both pathways + und_kv_cache = [] + for layer in self.layers: + und_h, gen_h, und_kv = layer.forward_with_generation( + und_h, gen_h, position_ids + ) + und_kv_cache.append(und_kv) + + # Apply generation final norm + gen_h = self.norm_moe_gen(gen_h) + + # Split and project back + if audio_tokens is not None: + vision_h = gen_h[:, :num_patches, :] + audio_h = gen_h[:, num_patches:, :] + vision_velocity = self.proj_out(vision_h) + audio_velocity = self.audio_proj_out(audio_h) + return (vision_velocity, audio_velocity), und_kv_cache + else: + velocity = self.proj_out(gen_h) + return velocity, und_kv_cache + + def diffusion_forward_cached( + self, + gen_tokens: mx.array, + timestep: mx.array, + und_kv_cache: list, + position_ids: mx.array, + text_len: int, + audio_tokens: Optional[mx.array] = None, + grid_t: int = 1, + grid_h: int = 1, + grid_w: int = 1, + noisy_frame_indexes: Optional[list[int]] = None, + ) -> mx.array | tuple[mx.array, mx.array]: + """Cached diffusion forward — generation pathway only. + + Reuses precomputed understanding K/V from step 0. + Only recomputes generation pathway (proj_in, timestep, attention, MLP, proj_out). + """ + num_patches = gen_tokens.shape[1] + + # Project generation tokens + selective timestep embedding + gen_h = self.proj_in(gen_tokens) + scaled_t = timestep * 0.001 + t_emb = self.time_embedder(scaled_t) + + if noisy_frame_indexes is None: + gen_h = gen_h + mx.expand_dims(t_emb, 1) + else: + spatial_tokens = grid_h * grid_w + noisy_mask = mx.zeros((num_patches,), dtype=gen_h.dtype) + for fi in noisy_frame_indexes: + start = fi * spatial_tokens + end = start + spatial_tokens + noisy_mask = noisy_mask.at[start:end].add(mx.ones((spatial_tokens,), dtype=gen_h.dtype)) + gen_h = gen_h + noisy_mask.reshape(1, num_patches, 1) * mx.expand_dims(t_emb, 1) + + # Handle audio tokens + if audio_tokens is not None: + audio_h = self.audio_proj_in(audio_tokens) + audio_h = audio_h + self.audio_modality_embed + audio_h = audio_h + mx.expand_dims(t_emb, 1) + gen_h = mx.concatenate([gen_h, audio_h], axis=1) + + # Forward through layers using cached understanding K/V + for layer, und_kv in zip(self.layers, und_kv_cache): + gen_h = layer.forward_generation_cached( + gen_h, und_kv, position_ids, text_len, + ) + + # Apply generation final norm + gen_h = self.norm_moe_gen(gen_h) + + # Split and project back + if audio_tokens is not None: + vision_h = gen_h[:, :num_patches, :] + audio_h = gen_h[:, num_patches:, :] + vision_velocity = self.proj_out(vision_h) + audio_velocity = self.audio_proj_out(audio_h) + return vision_velocity, audio_velocity + else: + velocity = self.proj_out(gen_h) + return velocity diff --git a/video/cosmos3/cosmos3/pipeline.py b/video/cosmos3/cosmos3/pipeline.py new file mode 100644 index 000000000..511f42ea5 --- /dev/null +++ b/video/cosmos3/cosmos3/pipeline.py @@ -0,0 +1,670 @@ +"""Generation pipeline for Cosmos 3 Nano on MLX. + +Wires the MoT transformer (generation path), scheduler, timestep +embedding, video VAE decoder, and audio decoder into a complete +text-to-image/video and image-to-video generation flow. +""" + +import json +import subprocess +import tempfile +import time +import wave +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Union + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from .model import Cosmos3Config, Cosmos3Model +from .scheduler import UniPCScheduler +from .timestep import TimestepEmbedding, apply_timestep_to_noisy_tokens + + +# Cosmos3 Nano VAE latent config (Wan2.2 AutoencoderKL) +Z_DIM = 48 +PATCH_SIZE = 2 +LATENTS_MEAN = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.157, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.123, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.052, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667, +] +LATENTS_STD = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.499, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.06, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744, +] + +_SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a given prompt." +_SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a given prompt." + + +class Cosmos3GenerationPipeline: + """End-to-end generation pipeline: text → image/video (+ audio). + + Orchestrates: + 1. Text tokenization and embedding + 2. Noise latent preparation + 3. Denoising loop (MoT generation path + scheduler) + 4. VAE decode to pixels + 5. Optional audio decode + """ + + def __init__( + self, + model: Cosmos3Model, + tokenizer, + model_dir: Optional[str | Path] = None, + ): + self.model = model + self.tokenizer = tokenizer + self.scheduler = UniPCScheduler() + self._model_dir = Path(model_dir) if model_dir is not None else None + + # Load reference negative prompt if available + self._negative_prompt_text = None + if model_dir is not None: + neg_path = Path(model_dir) / "assets" / "negative_prompt.json" + if neg_path.exists(): + with open(neg_path) as f: + self._negative_prompt_text = json.dumps(json.load(f)) + + # Pre-compute latent normalization tensors + self._latents_mean = mx.array(LATENTS_MEAN) + self._latents_std = mx.array(LATENTS_STD) + + def _build_position_ids( + self, + text_len: int, + grid_t: int, + grid_h: int, + grid_w: int, + audio_tokens: Optional[mx.array] = None, + ) -> mx.array: + """Build 3D mRoPE position IDs for text + generation tokens. + + Matches the position ID construction inside diffusion_forward. + """ + # Text positions + text_pos = mx.arange(text_len, dtype=mx.float32)[None, :] + text_position_ids = mx.stack([text_pos, text_pos, text_pos]) + + # Generation positions (FPS-modulated) + # Large temporal margin separates text and vision in mRoPE space. + # HF pipeline chains text_mrope → vision_mrope with offset = text_len, + # but the model may have been trained with a larger separation. + temporal_margin = 15000 + temporal_offset = float(text_len + temporal_margin) + fps = 24.0 + video_tcf = 4 + base_fps = 24.0 + tps = fps / video_tcf + base_tps = base_fps / video_tcf + + frame_indices = mx.arange(grid_t, dtype=mx.float32) + scaled_t = frame_indices / tps * base_tps + temporal_offset + t_idx = mx.broadcast_to(scaled_t.reshape(-1, 1), (grid_t, grid_h * grid_w)).reshape(1, -1) + + h_idx = mx.arange(grid_h, dtype=mx.float32).reshape(1, -1, 1) + h_idx = mx.broadcast_to(h_idx, (grid_t, grid_h, grid_w)).reshape(1, -1) + + w_idx = mx.arange(grid_w, dtype=mx.float32).reshape(1, 1, -1) + w_idx = mx.broadcast_to(w_idx, (grid_t, grid_h, grid_w)).reshape(1, -1) + + gen_position_ids = mx.stack([t_idx, h_idx, w_idx]) + + # Audio positions + if audio_tokens is not None: + sound_len = audio_tokens.shape[1] + audio_tps = fps / 1.0 + audio_base_tps = base_fps / 1.0 + audio_frame_indices = mx.arange(sound_len, dtype=mx.float32) + audio_scaled_t = audio_frame_indices / audio_tps * audio_base_tps + temporal_offset + audio_t_idx = audio_scaled_t.reshape(1, -1) + audio_h_idx = mx.zeros((1, sound_len), dtype=mx.float32) + audio_w_idx = mx.zeros((1, sound_len), dtype=mx.float32) + audio_position_ids = mx.stack([audio_t_idx, audio_h_idx, audio_w_idx]) + gen_position_ids = mx.concatenate([gen_position_ids, audio_position_ids], axis=2) + + return mx.concatenate([text_position_ids, gen_position_ids], axis=2) + + def _prepare_noise_latents( + self, + num_frames: int = 1, + height: int = 512, + width: int = 512, + z_dim: int = 48, + dtype: mx.Dtype = mx.bfloat16, + ) -> mx.array: + """Prepare initial noise latents. + + Returns: + [1, T_lat, H_lat, W_lat, z_dim] noise tensor (channels-last) + """ + # Compute latent dimensions + t_lat = max(1, num_frames // 4) # 4x temporal compression + h_lat = height // 16 # 16x spatial compression + w_lat = width // 16 + + noise = mx.random.normal((1, t_lat, h_lat, w_lat, z_dim)).astype(dtype) + return noise + + def _patchify_latents(self, latents: mx.array) -> mx.array: + """Convert VAE latents to patch tokens for the transformer. + + Input: [batch, T, H, W, z_dim] + Output: [batch, num_patches, patch_latent_dim] + + With patch_size=2: each 2x2 spatial region becomes one token. + patch_latent_dim = z_dim * patch_size * patch_size = 48 * 4 = 192 + + If H or W are not divisible by patch_size, zero-pads to the next + multiple (matching HF's _patchify_and_pack_latents). The padding + only exists during the transformer forward; _unpatchify_latents + crops back to original dims. + """ + batch, t, h, w, z = latents.shape + p = PATCH_SIZE + + h_p = (h + p - 1) // p # ceil division + w_p = (w + p - 1) // p + h_padded = h_p * p + w_padded = w_p * p + + # Zero-pad if needed (use concatenation, not .at[].add()) + if h_padded != h: + pad_h = mx.zeros((batch, t, h_padded - h, w, z), dtype=latents.dtype) + latents = mx.concatenate([latents, pad_h], axis=2) + if w_padded != w: + pad_w = mx.zeros((batch, t, latents.shape[2], w_padded - w, z), dtype=latents.dtype) + latents = mx.concatenate([latents, pad_w], axis=3) + + # [B, T, H_pad//p, p, W_pad//p, p, z] -> [B, T*H_p*W_p, p*p*z] + x = latents.reshape(batch, t, h_p, p, w_p, p, z) + x = mx.transpose(x, (0, 1, 2, 4, 3, 5, 6)) # [B, T, H_p, W_p, p, p, z] + x = x.reshape(batch, t * h_p * w_p, p * p * z) + + return x + + def _unpatchify_latents( + self, tokens: mx.array, t: int, h_p: int, w_p: int, + h_orig: int = 0, w_orig: int = 0, + ) -> mx.array: + """Convert patch tokens back to VAE latent shape. + + Input: [batch, num_patches, patch_latent_dim] + Output: [batch, T, H_orig, W_orig, z_dim] + + Crops padded dimensions back to h_orig × w_orig when provided. + """ + batch = tokens.shape[0] + p = PATCH_SIZE + z = Z_DIM + + x = tokens.reshape(batch, t, h_p, w_p, p, p, z) + x = mx.transpose(x, (0, 1, 2, 4, 3, 5, 6)) # [B, T, H_p, p, W_p, p, z] + x = x.reshape(batch, t, h_p * p, w_p * p, z) + + # Crop back to original dims (remove padding) + if h_orig > 0 and w_orig > 0: + x = x[:, :, :h_orig, :w_orig, :] + + return x + + def _encode_conditioning_image( + self, + image: np.ndarray, + num_frames: int, + height: int, + width: int, + ) -> mx.array: + """Encode a conditioning image to normalized VAE latents. + + Matches HF Cosmos3OmniPipeline: builds a full video tensor with the + conditioning frame repeated at every temporal position, then encodes + the entire tensor through the VAE. The temporal causal convolutions + produce per-frame latents that depend on preceding frames — frame 0 + sees zero-padded context while frames 1+ see frame 0's features. + This is critical for i2v quality: tiling a single-frame encode + produces identical latents at every position, but the model expects + temporally-varying latents from the causal conv processing. + + Args: + image: [H, W, 3] uint8 or float32 in [0, 1] + num_frames: total video frames + height: target height + width: target width + + Returns: + [1, T_lat, H_lat, W_lat, z_dim] normalized latents (channels-last) + """ + from .encode_vae import encode_video + + vae_dir = str(self._model_dir / "vae") + + # Resize image to target resolution + from PIL import Image as PILImage + if isinstance(image, np.ndarray): + pil_img = PILImage.fromarray( + (image * 255).astype(np.uint8) if image.dtype == np.float32 else image + ) + else: + pil_img = image + pil_img = pil_img.resize((width, height), PILImage.LANCZOS) + image_np = np.array(pil_img).astype(np.float32) / 255.0 + + if num_frames == 1: + # Single image: encode directly + return encode_video(image_np, vae_dir) + + # Build full video tensor: conditioning frame repeated at all positions. + # HF: vision_tensor[:,:,0] = frame; vision_tensor[:,:,1:] = frame.repeat() + # [T, H, W, 3] in [0, 1] + video_tensor = np.stack([image_np] * num_frames, axis=0) + return encode_video(video_tensor, vae_dir) + + def generate( + self, + prompt: str, + num_frames: int = 1, + height: int = 512, + width: int = 512, + num_inference_steps: int = 30, + guidance_scale: float = 6.0, + seed: Optional[int] = None, + enable_audio: bool = False, + image: Optional[Union[np.ndarray, "PILImage"]] = None, + negative_prompt: Optional[str] = None, + ) -> dict: + """Generate image/video from text prompt, optionally conditioned on an image. + + Args: + prompt: text description + num_frames: number of video frames (1 = single image) + height: output height in pixels + width: output width in pixels + num_inference_steps: denoising steps + guidance_scale: classifier-free guidance strength + seed: random seed for reproducibility + enable_audio: whether to generate audio alongside video + image: optional conditioning image for i2v. When provided, frame 0 + is anchored to this image and the remaining frames are denoised + freely. Can be numpy array [H,W,3] (uint8 or float32) or PIL Image. + negative_prompt: optional negative prompt text. If None, uses the + model's built-in negative prompt from assets/negative_prompt.json. + + Returns: + dict with 'latents' (normalized, use decode_latents() to decode), + 'video' (if VAE available), 'audio_latents' (if enable_audio), + 'audio' (if audio decoded) + """ + if seed is not None: + mx.random.seed(seed) + + dtype = mx.bfloat16 + + has_image_condition = image is not None and num_frames > 1 + + # 1. Tokenize prompt + is_image = (num_frames == 1) + system_msg = _SYSTEM_PROMPT_IMAGE if is_image else _SYSTEM_PROMPT_VIDEO + + # Build user content with resolution/duration suffix + if is_image: + user_content = prompt + f" This image is of {height}x{width} resolution." + else: + fps = 24 + duration = num_frames / fps + user_content = (prompt + + f" The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." + + f" This video is of {height}x{width} resolution.") + + messages = [ + {"role": "system", "content": system_msg}, + {"role": "user", "content": user_content}, + ] + text = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + tokens = self.tokenizer.encode(text) + # Append EOS + vision_start sentinel (tells model generation follows) + eos_id = self.tokenizer.eos_token_id + vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>") + tokens = tokens + [eos_id, vision_start_id] + cond_ids = mx.array([tokens]) + + # Unconditional prompt for CFG + neg_prompt = negative_prompt if negative_prompt is not None else (self._negative_prompt_text or "") + if is_image: + neg_suffix = f" This image is not of {height}x{width} resolution." + else: + neg_suffix = (f" The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." + + f" This video is not of {height}x{width} resolution.") + neg_content = neg_prompt + neg_suffix + + uncond_messages = [ + {"role": "system", "content": system_msg}, + {"role": "user", "content": neg_content}, + ] + uncond_text = self.tokenizer.apply_chat_template( + uncond_messages, tokenize=False, add_generation_prompt=True, + ) + uncond_tokens = self.tokenizer.encode(uncond_text) + uncond_tokens = uncond_tokens + [eos_id, vision_start_id] + uncond_ids = mx.array([uncond_tokens]) + + # 2. Prepare noise latents + z_dim = Z_DIM + latents = self._prepare_noise_latents( + num_frames, height, width, z_dim, dtype + ) + t_lat = latents.shape[1] + h_lat = latents.shape[2] + w_lat = latents.shape[3] + + # Patchify grid: ceil division to match HF's zero-padding patchify. + # Latents stay at original h_lat × w_lat between denoising steps. + # Padding to patch-aligned dims only happens inside _patchify_latents, + # and _unpatchify_latents crops back to h_lat × w_lat. + p = PATCH_SIZE + h_p = (h_lat + p - 1) // p # ceil: 45→23 at 720p + w_p = (w_lat + p - 1) // p + + # 2a. Image-to-video conditioning: encode image, create mask, mix + # Vision condition mask: [T_lat, 1, 1] — 1.0 for conditioned frames, 0.0 for noisy + vision_condition_mask = mx.zeros((t_lat, 1, 1)) + if has_image_condition: + print(" Encoding conditioning image...") + cond_latents = self._encode_conditioning_image( + image, num_frames, height, width + ).astype(dtype) + # Crop to match noise latent temporal/spatial dims + cond_latents = cond_latents[:, :t_lat, :h_lat, :w_lat, :] + mx.eval(cond_latents) + + # Frame 0 is conditioned + vision_condition_mask = vision_condition_mask.at[0, 0, 0].add(mx.array(1.0)) + + # Mix: conditioned frames get encoded latent, rest get noise + # cond_latents is [1, T_lat, H_lat, W_lat, z_dim], mask broadcasts over spatial+channel + mask_5d = vision_condition_mask.reshape(1, t_lat, 1, 1, 1) + latents = mask_5d * cond_latents + (1.0 - mask_5d) * latents + mx.eval(latents) + + num_patches = t_lat * h_p * w_p + + # 2b. Prepare audio noise latents if enabled + sound_latents = None + sound_len = 0 + if enable_audio and num_frames <= 1: + print(" Warning: --enable-audio ignored for single-frame generation") + if enable_audio and num_frames > 1: + sound_dim = 64 # Cosmos3 audio latent dim + sampling_rate = 48000 + fps = 24 + hop_size = 1920 + n_audio_samples = int(num_frames / fps * sampling_rate) + sound_len = (n_audio_samples + hop_size - 1) // hop_size + sound_latents = mx.random.normal((1, sound_len, sound_dim)).astype(dtype) + print(f" Audio latents: {sound_len} frames ({n_audio_samples} samples @ {sampling_rate}Hz)") + + # 3. Set up schedulers + self.scheduler.set_timesteps(num_inference_steps) + if sound_latents is not None: + audio_scheduler = UniPCScheduler() + audio_scheduler.set_timesteps(num_inference_steps) + + print(f" Latent shape: ({t_lat}, {h_lat}, {w_lat}, {z_dim})") + print(f" Patches: {num_patches} ({t_lat}×{h_p}×{w_p})") + print(f" Denoising steps: {num_inference_steps}") + + # 4. Denoising loop with text KV caching + # Step 0: full forward (both pathways), cache understanding K/V + # Steps 1+: generation pathway only, reuse cached K/V + cond_kv_cache = None + uncond_kv_cache = None + cond_position_ids = None + uncond_position_ids = None + + # Compute noisy frame indexes for selective timestep embedding (i2v). + # For t2v (no conditioning): None = all frames get timestep. + # For i2v: exclude conditioned frames (frame 0). + noisy_fi = None + if has_image_condition: + cond_frame_set = set() + for fi in range(t_lat): + if vision_condition_mask[fi, 0, 0].item() > 0: + cond_frame_set.add(fi) + noisy_fi = [fi for fi in range(t_lat) if fi not in cond_frame_set] + + for i in range(num_inference_steps): + sigma = float(self.scheduler.sigmas[i].item()) + + # Patchify current latents + gen_tokens = self._patchify_latents(latents).astype(dtype) + + # Timestep tensor (sigma * num_train_timesteps) + t_tensor = mx.array([sigma * self.scheduler.num_train_timesteps]).astype(dtype) + + # Audio tokens for this step + audio_tokens = sound_latents if sound_latents is not None else None + + if cond_kv_cache is None: + # Step 0: full forward, build cache + cond_result, cond_kv_cache = self.model.diffusion_forward( + cond_ids, gen_tokens, t_tensor, + grid_t=t_lat, grid_h=h_p, grid_w=w_p, + audio_tokens=audio_tokens, + noisy_frame_indexes=noisy_fi, + ) + cond_text_len = cond_ids.shape[1] + cond_position_ids = self._build_position_ids( + cond_text_len, t_lat, h_p, w_p, audio_tokens, + ) + mx.eval(*[kv[0] for kv in cond_kv_cache], *[kv[1] for kv in cond_kv_cache]) + else: + # Steps 1+: cached forward (generation pathway only) + cond_result = self.model.diffusion_forward_cached( + gen_tokens, t_tensor, cond_kv_cache, + cond_position_ids, cond_ids.shape[1], + audio_tokens=audio_tokens, + grid_t=t_lat, grid_h=h_p, grid_w=w_p, + noisy_frame_indexes=noisy_fi, + ) + + # Classifier-free guidance + if guidance_scale != 1.0: + if audio_tokens is not None: + cond_velocity, cond_audio_vel = cond_result + mx.eval(cond_velocity, cond_audio_vel) + else: + cond_velocity = cond_result + mx.eval(cond_velocity) + + if uncond_kv_cache is None: + # Step 0: full uncond forward, build cache + uncond_result, uncond_kv_cache = self.model.diffusion_forward( + uncond_ids, gen_tokens, t_tensor, + grid_t=t_lat, grid_h=h_p, grid_w=w_p, + audio_tokens=audio_tokens, + noisy_frame_indexes=noisy_fi, + ) + uncond_text_len = uncond_ids.shape[1] + uncond_position_ids = self._build_position_ids( + uncond_text_len, t_lat, h_p, w_p, audio_tokens, + ) + mx.eval(*[kv[0] for kv in uncond_kv_cache], *[kv[1] for kv in uncond_kv_cache]) + else: + uncond_result = self.model.diffusion_forward_cached( + gen_tokens, t_tensor, uncond_kv_cache, + uncond_position_ids, uncond_ids.shape[1], + audio_tokens=audio_tokens, + noisy_frame_indexes=noisy_fi, + grid_t=t_lat, grid_h=h_p, grid_w=w_p, + ) + + if audio_tokens is not None: + uncond_velocity, uncond_audio_vel = uncond_result + velocity_patches = uncond_velocity + guidance_scale * ( + cond_velocity - uncond_velocity + ) + audio_vel = uncond_audio_vel + guidance_scale * ( + cond_audio_vel - uncond_audio_vel + ) + else: + uncond_velocity = uncond_result + velocity_patches = uncond_velocity + guidance_scale * ( + cond_velocity - uncond_velocity + ) + else: + if audio_tokens is not None: + velocity_patches, audio_vel = cond_result + else: + velocity_patches = cond_result + + # Unpatchify velocity back to latent shape, crop to original dims + velocity = self._unpatchify_latents( + velocity_patches, t_lat, h_p, w_p, + h_orig=h_lat, w_orig=w_lat, + ) + + # Zero velocity at conditioned frame positions (i2v). + # With flow-matching: x_{t-1} = x_t + step * velocity. + # Zeroing velocity keeps conditioned frames fixed at their initial value. + if has_image_condition: + mask_5d = vision_condition_mask.reshape(1, t_lat, 1, 1, 1) + velocity = velocity * (1.0 - mask_5d) + + # Scheduler step (uses internal step_index) + latents = self.scheduler.step(velocity, t_tensor, latents) + mx.eval(latents) + + # Audio scheduler step + if sound_latents is not None: + # audio_vel: [1, sound_len, sound_dim] — already in latent shape + sound_latents = audio_scheduler.step(audio_vel, t_tensor, sound_latents) + mx.eval(sound_latents) + + if (i + 1) % 10 == 0 or i == 0: + print(f" Step {i+1}/{num_inference_steps} (σ={sigma:.4f})") + + # 5. Return normalized latents (decode_latents handles denormalization) + result = {"latents": latents} + + if sound_latents is not None: + # Store audio latents as [sound_dim, T] channels-first for audio decoder + result["audio_latents"] = mx.transpose(sound_latents[0], (1, 0)) + + print(" Done!") + return result + + +def save_video( + video_frames: np.ndarray, + output_path: str, + fps: int = 24, + audio_waveform: np.ndarray = None, + audio_sample_rate: int = 48000, +) -> str: + """Save video frames (and optional audio) as MP4 or GIF. + + Args: + video_frames: [T, H, W, 3] uint8 frames + output_path: output file path (.mp4 or .gif) + fps: video frame rate + audio_waveform: optional [2, N] or [N] float audio in [-1, 1] + audio_sample_rate: audio sample rate in Hz + + Returns: + output path + """ + import shutil + + output_path = str(output_path) + is_mp4 = output_path.endswith(".mp4") + + if is_mp4 and shutil.which("ffmpeg") is None: + if audio_waveform is not None: + print("Warning: ffmpeg not found. Audio will be dropped. Saving as GIF.") + else: + print("Warning: ffmpeg not found. Saving as GIF instead.") + output_path = output_path.rsplit(".", 1)[0] + ".gif" + is_mp4 = False + + if is_mp4 and audio_waveform is not None: + # Write frames as PNG sequence + WAV, mux with ffmpeg + with tempfile.TemporaryDirectory() as tmpdir: + from PIL import Image + + # Write frames + for i, frame in enumerate(video_frames): + Image.fromarray(frame).save(f"{tmpdir}/frame_{i:04d}.png") + + # Write WAV + wav_path = f"{tmpdir}/audio.wav" + if audio_waveform.ndim == 1: + audio_waveform = audio_waveform[np.newaxis, :] + n_channels = audio_waveform.shape[0] + audio_int16 = (audio_waveform * 32767).clip(-32768, 32767).astype(np.int16) + with wave.open(wav_path, "w") as wf: + wf.setnchannels(n_channels) + wf.setsampwidth(2) + wf.setframerate(audio_sample_rate) + if n_channels > 1: + interleaved = np.stack( + [audio_int16[c] for c in range(n_channels)], axis=-1 + ).flatten() + else: + interleaved = audio_int16[0] + wf.writeframes(interleaved.tobytes()) + + # Mux with ffmpeg + cmd = [ + "ffmpeg", "-y", + "-framerate", str(fps), + "-i", f"{tmpdir}/frame_%04d.png", + "-i", wav_path, + "-c:v", "libx264", "-pix_fmt", "yuv420p", + "-c:a", "aac", "-b:a", "192k", + "-shortest", + output_path, + ] + subprocess.run(cmd, check=True) + + elif is_mp4: + # Video-only MP4 + with tempfile.TemporaryDirectory() as tmpdir: + from PIL import Image + + for i, frame in enumerate(video_frames): + Image.fromarray(frame).save(f"{tmpdir}/frame_{i:04d}.png") + + cmd = [ + "ffmpeg", "-y", + "-framerate", str(fps), + "-i", f"{tmpdir}/frame_%04d.png", + "-c:v", "libx264", "-pix_fmt", "yuv420p", + output_path, + ] + subprocess.run(cmd, check=True) + + else: + # GIF fallback + from PIL import Image + + frames = [Image.fromarray(f) for f in video_frames] + duration = int(1000 / fps) + frames[0].save( + output_path, save_all=True, append_images=frames[1:], + duration=duration, loop=0, + ) + + return output_path diff --git a/video/cosmos3/cosmos3/rope.py b/video/cosmos3/cosmos3/rope.py new file mode 100644 index 000000000..e08f155b4 --- /dev/null +++ b/video/cosmos3/cosmos3/rope.py @@ -0,0 +1,151 @@ +"""3D Multi-dimensional Rotary Position Embeddings for Cosmos 3. + +Cosmos 3 uses interleaved mRoPE with three axes (temporal, height, width) +following the Qwen3-VL design. All axes share a single set of inverse +frequencies. Position IDs from different axes are interleaved before +frequency computation: + - temporal positions go to indices {0, 3, 6, ...} + - height positions go to indices {1, 4, 7, ...} + - width positions go to indices {2, 5, 8, ...} + +The mrope_section [24, 20, 20] determines how many frequency dimensions +are assigned to each axis. Total = 64 half-dims = 128 head_dim / 2. +""" + +import math +from typing import Tuple + +import mlx.core as mx +import mlx.nn as nn + + +class Cosmos3RotaryEmbedding(nn.Module): + """Compute 3D interleaved mRoPE cos/sin embeddings. + + Matches HuggingFace Cosmos3VLTextRotaryEmbedding: + - Single shared inv_freq for all axes + - Interleaved axis layout (T,H,W at strides of 3) + """ + + def __init__( + self, + head_dim: int = 128, + mrope_section: list[int] | None = None, + rope_theta: float = 5_000_000.0, + ): + super().__init__() + self.head_dim = head_dim + self.mrope_section = mrope_section or [24, 20, 20] + self.rope_theta = rope_theta + + half_dim = head_dim // 2 + assert sum(self.mrope_section) == half_dim, ( + f"mrope_section {self.mrope_section} sums to {sum(self.mrope_section)}, " + f"expected {half_dim}" + ) + + # Single shared inv_freq for all axes (matching HF reference) + inv_freq = 1.0 / (rope_theta ** (mx.arange(0, head_dim, 2, dtype=mx.float32) / head_dim)) + self._inv_freq = inv_freq # [half_dim = 64] + + def __call__( + self, + position_ids: mx.array, + seq_len: int, + ) -> Tuple[mx.array, mx.array]: + """Compute cos/sin embeddings from 3-axis position IDs. + + Args: + position_ids: [3, batch, seq_len] — position IDs per axis + seq_len: sequence length + + Returns: + cos: [batch, seq_len, head_dim] + sin: [batch, seq_len, head_dim] + """ + half_dim = self.head_dim // 2 # 64 + + # Compute freqs for each axis: pos[axis] @ inv_freq + # Each axis gets [batch, seq_len, half_dim] frequencies + # freqs shape: [3, batch, seq_len, half_dim] + freqs_per_axis = [] + for axis_idx in range(3): + pos = position_ids[axis_idx].astype(mx.float32) # [batch, seq_len] + # Outer product: [batch, seq_len, 1] * [1, 1, half_dim] + f = mx.expand_dims(pos, -1) * mx.expand_dims( + mx.expand_dims(self._inv_freq, 0), 0 + ) # [batch, seq_len, half_dim] + freqs_per_axis.append(f) + + # Interleave axes into a single frequency tensor + # HF reference: temporal at {0,3,6,...}, height at {1,4,7,...}, width at {2,5,8,...} + # + # Build a mapping: for each output dim, which axis provides its value + # All three axes computed the same 64 frequencies (shared inv_freq), + # so we just need to pick which axis's position-scaled result goes where. + # + # Start with temporal everywhere, then overwrite H and W at stride-3 + # This matches HF's apply_interleaved_mrope + axis_assignment = [0] * half_dim # default: temporal + for dim_offset, axis_idx in enumerate([1, 2], start=1): + section_len = self.mrope_section[axis_idx] + total_interleaved = section_len * 3 + for freq_idx in range(dim_offset, min(total_interleaved, half_dim), 3): + axis_assignment[freq_idx] = axis_idx + + # Gather from the appropriate axis for each frequency dimension + # Stack all axes: [3, batch, seq_len, half_dim] + all_freqs = mx.stack(freqs_per_axis) # [3, B, N, 64] + # Select per-dim: use axis_assignment to index into axis dimension + axis_indices = mx.array(axis_assignment) # [half_dim] + # Gather: for each dim d, take all_freqs[axis_assignment[d], :, :, d] + freqs_parts = [] + for d in range(half_dim): + a = axis_assignment[d] + freqs_parts.append(freqs_per_axis[a][..., d:d+1]) + freqs = mx.concatenate(freqs_parts, axis=-1) # [batch, seq_len, half_dim] + + # Double-up: [cos(f), cos(f)] and [sin(f), sin(f)] for full head_dim + cos = mx.cos(freqs) + sin = mx.sin(freqs) + + return cos, sin + + +def rotate_half(x: mx.array) -> mx.array: + """Rotate halves: [x1, x2] -> [-x2, x1] where x1, x2 are each half_dim.""" + half = x.shape[-1] // 2 + x1 = x[..., :half] + x2 = x[..., half:] + return mx.concatenate([-x2, x1], axis=-1) + + +def apply_rotary_pos_emb( + q: mx.array, + k: mx.array, + cos: mx.array, + sin: mx.array, +) -> Tuple[mx.array, mx.array]: + """Apply rotary position embeddings to query and key tensors. + + Args: + q: [batch, seq_len, num_heads, head_dim] + k: [batch, seq_len, num_kv_heads, head_dim] + cos: [batch, seq_len, half_dim] + sin: [batch, seq_len, half_dim] + + Returns: + q_rotated, k_rotated with same shapes as inputs + """ + # Duplicate cos/sin to match full head_dim: [cos, cos] for the two halves + cos_full = mx.concatenate([cos, cos], axis=-1) # [batch, seq_len, head_dim] + sin_full = mx.concatenate([sin, sin], axis=-1) # [batch, seq_len, head_dim] + + # Expand for broadcasting over heads: [batch, seq_len, 1, head_dim] + cos_full = mx.expand_dims(cos_full, 2) + sin_full = mx.expand_dims(sin_full, 2) + + q_rot = q * cos_full + rotate_half(q) * sin_full + k_rot = k * cos_full + rotate_half(k) * sin_full + + return q_rot, k_rot diff --git a/video/cosmos3/cosmos3/scheduler.py b/video/cosmos3/cosmos3/scheduler.py new file mode 100644 index 000000000..acdd76b98 --- /dev/null +++ b/video/cosmos3/cosmos3/scheduler.py @@ -0,0 +1,303 @@ +"""UniPC Multistep Scheduler for Cosmos 3 diffusion generation. + +Implements the unified predictor-corrector (UniPC) framework for +iterative denoising with flow matching. Matches the HuggingFace +UniPCMultistepScheduler with flow_prediction type. + +Key formulas (flow matching path): +- sigma_t is the noise level (0 = clean, 1 = pure noise) +- alpha_t = 1 - sigma_t +- x_t = alpha_t * x_0 + sigma_t * noise +- Model predicts velocity v, x_0 recovered as: x_0 = x_t - sigma_t * v +- Stepping uses log-SNR (lambda) space exponential integrators + +The full UniPC algorithm runs a predictor (P) step followed by a corrector (C) +step at each iteration. The corrector uses the current model output to refine +the sample predicted by the previous predictor step. +""" + +from typing import Optional + +import mlx.core as mx +import numpy as np + + +class UniPCScheduler: + """UniPC multi-step scheduler for flow matching denoising. + + Matches HF UniPCMultistepScheduler config: + - prediction_type: flow_prediction + - use_flow_sigmas: True + - predict_x0: True + - solver_order: 2 + - solver_type: bh2 + - lower_order_final: True + """ + + def __init__( + self, + num_train_timesteps: int = 1000, + flow_shift: float = 1.0, + solver_order: int = 2, + use_karras_sigmas: bool = True, + sigma_min: float = 0.147, + sigma_max: float = 200.0, + lower_order_final: bool = True, + ): + self.num_train_timesteps = num_train_timesteps + self.flow_shift = flow_shift + self.solver_order = solver_order + self.use_karras_sigmas = use_karras_sigmas + self.sigma_min = sigma_min + self.sigma_max = sigma_max + self.lower_order_final = lower_order_final + + self.sigmas = None + self.timesteps = None + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.last_sample = None + self.this_order = None + self.step_index = 0 + + def set_timesteps(self, num_inference_steps: int): + """Set the discrete timesteps for inference. + + Uses Karras sigma schedule (matching HF config: use_karras_sigmas=True, + use_flow_sigmas=True) converted to flow-matching space. + """ + self.num_inference_steps = num_inference_steps + + if self.use_karras_sigmas: + rho = 7.0 + ramp = np.linspace(0, 1, num_inference_steps) + min_inv_rho = self.sigma_min ** (1 / rho) + max_inv_rho = self.sigma_max ** (1 / rho) + karras_sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho + sigmas = karras_sigmas / (karras_sigmas + 1) + else: + sigmas = np.linspace(1, 1 / self.num_train_timesteps, num_inference_steps + 1)[:-1] + if self.flow_shift != 1.0: + sigmas = self.flow_shift * sigmas / (1 + (self.flow_shift - 1) * sigmas) + + eps = 1e-6 + if abs(sigmas[0] - 1.0) < eps: + sigmas[0] -= eps + + sigmas = np.append(sigmas, 0.0) + + self.sigmas = mx.array(sigmas.astype(np.float32)) + self.timesteps = mx.array((sigmas[:-1] * self.num_train_timesteps).astype(np.float32)) + + # Reset state + self.model_outputs = [None] * self.solver_order + self.timestep_list = [None] * self.solver_order + self.lower_order_nums = 0 + self.last_sample = None + self.this_order = None + self.step_index = 0 + + def _convert_model_output(self, model_output: mx.array, sample: mx.array) -> mx.array: + """Convert flow velocity prediction to x0 prediction.""" + sigma = float(self.sigmas[self.step_index].item()) + return sample - sigma * model_output + + def _compute_rks_and_D1s(self, order: int, h: float, lambda_s0: float, + m0: mx.array, step_offset: int = 0): + """Compute rk ratios and D1 differences for multistep methods.""" + rks = [] + D1s = [] + for i in range(1, order): + si = self.step_index - i - step_offset + mi = self.model_outputs[-(i + 1)] + if mi is None: + break + sigma_si = float(self.sigmas[si].item()) + alpha_si = 1.0 - sigma_si + lambda_si = np.log(alpha_si / max(sigma_si, 1e-10)) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + return rks, D1s + + def _compute_bh2_coefficients(self, h: float, order: int): + """Compute R matrix and b vector for BH2 solver.""" + hh = -h # predict_x0 mode + h_phi_1 = np.expm1(hh) + h_phi_k = h_phi_1 / hh - 1 + factorial_i = 1 + B_h = np.expm1(hh) + + b = [] + for i in range(1, order + 1): + b.append(float(h_phi_k * factorial_i / B_h)) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + return h_phi_1, B_h, b + + def _uni_p_bh_update(self, sample: mx.array, order: int) -> mx.array: + """UniP predictor step (B(h) version). + + Predicts x_{t+1} from x_t using stored model outputs (x0 predictions). + """ + m0 = self.model_outputs[-1] + + sigma_t = float(self.sigmas[self.step_index + 1].item()) + sigma_s0 = float(self.sigmas[self.step_index].item()) + alpha_t = 1.0 - sigma_t + alpha_s0 = 1.0 - sigma_s0 + + if sigma_t == 0.0: + return m0 + + lambda_t = np.log(max(alpha_t, 1e-10) / max(sigma_t, 1e-10)) + lambda_s0 = np.log(alpha_s0 / max(sigma_s0, 1e-10)) + h = lambda_t - lambda_s0 + + h_phi_1, B_h, b = self._compute_bh2_coefficients(h, order) + + # First-order base step + x_t_ = (sigma_t / sigma_s0) * sample - alpha_t * h_phi_1 * m0 + + if order == 1 or self.model_outputs[-2] is None: + return x_t_ + + rks, D1s = self._compute_rks_and_D1s(order, h, lambda_s0, m0) + if not D1s: + return x_t_ + + # For order 2, HF hardcodes rho_p = 0.5 + rho_p = 0.5 + pred_res = rho_p * D1s[0] + + x_t = x_t_ - alpha_t * B_h * pred_res + return x_t + + def _uni_c_bh_update( + self, this_model_output: mx.array, last_sample: mx.array, + this_sample: mx.array, order: int, + ) -> mx.array: + """UniC corrector step (B(h) version). + + Corrects the predictor output using the model evaluation at the predicted point. + """ + m0 = self.model_outputs[-1] + + sigma_t = float(self.sigmas[self.step_index].item()) + sigma_s0 = float(self.sigmas[self.step_index - 1].item()) + alpha_t = 1.0 - sigma_t + alpha_s0 = 1.0 - sigma_s0 + + lambda_t = np.log(max(alpha_t, 1e-10) / max(sigma_t, 1e-10)) + lambda_s0 = np.log(alpha_s0 / max(sigma_s0, 1e-10)) + h = lambda_t - lambda_s0 + + h_phi_1, B_h, b_list = self._compute_bh2_coefficients(h, order) + + # Base step using last_sample (the sample before the predictor) + x_t_ = (sigma_t / sigma_s0) * last_sample - alpha_t * h_phi_1 * m0 + + # Compute rks and D1s from history (offset by 1 for corrector indexing) + rks, D1s = self._compute_rks_and_D1s(order, h, lambda_s0, m0, step_offset=1) + + # Build R matrix and solve for rhos_c + # rks_full includes the historical rks plus a trailing 1.0 + rks_np = np.array([rk for rk in rks] + [1.0]) + + R = np.stack([np.power(rks_np, i) for i in range(order)]) # [order, order] + b = np.array(b_list[:order]) + + if order == 1: + rhos_c = np.array([0.5]) + else: + rhos_c = np.linalg.solve(R, b) + + # Apply correction + D1_t = this_model_output - m0 + if D1s: + corr_res = sum(rhos_c[k] * D1s[k] for k in range(len(D1s))) + else: + corr_res = 0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + return x_t + + def step( + self, + model_output: mx.array, + timestep: mx.array, + sample: mx.array, + ) -> mx.array: + """Perform one denoising step using UniPC predictor-corrector. + + Matches HF UniPCMultistepScheduler.step() exactly: + 1. Convert model output to x0 prediction + 2. Corrector step (refine previous predictor output using current model eval) + 3. Update history + 4. Predictor step (predict next sample) + """ + # Convert model output to x0 prediction + model_output_x0 = self._convert_model_output(model_output, sample) + + # Corrector step: refine the current sample using the new model output + use_corrector = ( + self.step_index > 0 + and self.last_sample is not None + ) + if use_corrector: + sample = self._uni_c_bh_update( + this_model_output=model_output_x0, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + ) + + # Shift history buffers + for i in range(self.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_x0 + self.timestep_list[-1] = float(timestep.item()) if hasattr(timestep, 'item') else float(timestep) + + # Determine effective order for this step + if self.lower_order_final: + this_order = min(self.solver_order, len(self.timesteps) - self.step_index) + else: + this_order = self.solver_order + + # Warmup: don't use higher order until we have enough history + this_order = min(this_order, self.lower_order_nums + 1) + assert this_order > 0 + self.this_order = this_order + + # Save sample for next corrector step + self.last_sample = sample + + # Predictor step + prev_sample = self._uni_p_bh_update(sample=sample, order=this_order) + + if self.lower_order_nums < self.solver_order: + self.lower_order_nums += 1 + + self.step_index += 1 + return prev_sample + + def add_noise( + self, + original_samples: mx.array, + noise: mx.array, + timestep: mx.array, + ) -> mx.array: + """Add noise to clean samples at given timestep. + + Flow matching: X_t = (1 - sigma) * data + sigma * noise + """ + sigma = timestep / self.num_train_timesteps + if sigma.ndim == 0: + sigma = sigma.reshape(1) + while sigma.ndim < original_samples.ndim: + sigma = mx.expand_dims(sigma, -1) + + return (1 - sigma) * original_samples + sigma * noise diff --git a/video/cosmos3/cosmos3/timestep.py b/video/cosmos3/cosmos3/timestep.py new file mode 100644 index 000000000..035759f54 --- /dev/null +++ b/video/cosmos3/cosmos3/timestep.py @@ -0,0 +1,80 @@ +"""Timestep embedding for Cosmos 3 diffusion generation. + +Sinusoidal timestep embedding → MLP projection, applied to noisy tokens +via scatter-add (only noisy frames get timestep conditioning). +""" + +import math + +import mlx.core as mx +import mlx.nn as nn + + +class TimestepEmbedding(nn.Module): + """Sinusoidal timestep embedding with MLP projection. + + Maps scalar timesteps to hidden_size embeddings. + """ + + def __init__(self, hidden_size: int = 4096, freq_dim: int = 256): + super().__init__() + self.freq_dim = freq_dim + self.linear_1 = nn.Linear(freq_dim, hidden_size, bias=True) + self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) + + def _sinusoidal_embedding(self, timesteps: mx.array) -> mx.array: + """Compute sinusoidal embeddings for timesteps. + + Args: + timesteps: [batch] scalar timesteps + + Returns: + [batch, freq_dim] sinusoidal embeddings + """ + half_dim = self.freq_dim // 2 + freqs = mx.exp( + -math.log(10000.0) * mx.arange(half_dim, dtype=mx.float32) / half_dim + ) + # [batch, 1] * [1, half_dim] -> [batch, half_dim] + args = mx.expand_dims(timesteps.astype(mx.float32), -1) * mx.expand_dims(freqs, 0) + embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1) + return embedding + + def __call__(self, timesteps: mx.array) -> mx.array: + """Compute timestep embeddings. + + Args: + timesteps: [batch] scalar timesteps + + Returns: + [batch, hidden_size] timestep embeddings + """ + emb = self._sinusoidal_embedding(timesteps) + emb = nn.silu(self.linear_1(emb)) + emb = self.linear_2(emb) + return emb + + +def apply_timestep_to_noisy_tokens( + hidden_states: mx.array, + timestep_emb: mx.array, + noisy_mask: mx.array, +) -> mx.array: + """Apply timestep embeddings only to noisy (generation) tokens. + + Instead of scatter_add, we use broadcasting with a mask. + + Args: + hidden_states: [batch, seq_len, hidden_size] + timestep_emb: [batch, hidden_size] + noisy_mask: [batch, seq_len] boolean mask (True = noisy token) + + Returns: + hidden_states with timestep embedding added to noisy positions + """ + # Expand timestep_emb: [batch, 1, hidden_size] + emb = mx.expand_dims(timestep_emb, 1) + # Expand mask: [batch, seq_len, 1] + mask = mx.expand_dims(noisy_mask.astype(hidden_states.dtype), -1) + # Add only to masked positions + return hidden_states + emb * mask diff --git a/video/cosmos3/img2video.py b/video/cosmos3/img2video.py new file mode 100644 index 000000000..f6e74b963 --- /dev/null +++ b/video/cosmos3/img2video.py @@ -0,0 +1,159 @@ +"""Generate videos from an input image using Cosmos 3 Nano on MLX.""" + +import argparse +import time +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +from PIL import Image + +from cosmos3.load import load_transformer, load_tokenizer +from cosmos3.pipeline import Cosmos3GenerationPipeline, save_video +from cosmos3.decode_vae import decode_latents + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate videos from an input image using Cosmos 3 Nano" + ) + parser.add_argument("prompt") + parser.add_argument("--image", type=str, required=True, help="Conditioning image path") + parser.add_argument( + "--model-dir", + type=str, + default="weights/Cosmos3-Nano", + help="Path to Cosmos3-Nano weights directory", + ) + parser.add_argument( + "--size", + type=lambda x: tuple(map(int, x.split("x"))), + default=(256, 256), + help="Video size as WxH (default: 256x256)", + ) + parser.add_argument("--frames", type=int, default=16) + parser.add_argument( + "--steps", type=int, default=30, help="Number of denoising steps" + ) + parser.add_argument("--guidance", type=float, default=6.0) + parser.add_argument("--seed", type=int) + parser.add_argument( + "--quantize", + "-q", + type=int, + nargs="?", + const=8, + default=0, + choices=[0, 8], + metavar="{8}", + help="Quantize model weights to 8-bit (default when flag used without value)", + ) + parser.add_argument( + "--n-prompt", + default=None, + help="Negative prompt (default: model's built-in negative prompt)", + ) + parser.add_argument("--output", default="out.mp4") + parser.add_argument("--enable-audio", action="store_true", help="Generate audio") + parser.add_argument( + "--no-cache", + action="store_true", + help="Disable Metal buffer cache (mx.set_cache_limit(0)) to reduce swap pressure", + ) + parser.add_argument("--verbose", "-v", action="store_true") + args = parser.parse_args() + + width, height = args.size + + if args.frames < 8 or args.frames % 4 != 0: + parser.error("--frames must be a multiple of 4 and >= 8 (e.g., 8, 16, 32)") + if width % 16 != 0 or height % 16 != 0: + parser.error(f"--size dimensions must be multiples of 16 (got {width}x{height})") + + mx.set_default_device(mx.gpu) + if args.no_cache: + mx.set_cache_limit(0) + + # Preflight: check required assets before expensive model load + vae_dir = Path(args.model_dir) / "vae" + if not (vae_dir / "config.json").exists(): + parser.error(f"VAE weights not found at {vae_dir}. Download the full model first.") + if args.enable_audio and not (Path(args.model_dir) / "sound_tokenizer" / "config.json").exists(): + parser.error(f"Sound tokenizer not found. Download the full model or disable --enable-audio.") + + # Load conditioning image + img = np.array(Image.open(args.image).convert("RGB")) + print(f"Input image: {args.image} ({img.shape[1]}x{img.shape[0]})") + + # Load model + print(f"Loading model from {args.model_dir}...") + t0 = time.time() + model = load_transformer(args.model_dir, reasoner_only=False) + tokenizer = load_tokenizer(args.model_dir) + + if args.quantize: + nn.quantize(model, bits=args.quantize) + mx.eval(model.parameters()) + print(f"Quantized to {args.quantize}-bit") + + print(f"Model loaded in {time.time() - t0:.1f}s") + + pipeline = Cosmos3GenerationPipeline( + model=model, + tokenizer=tokenizer, + model_dir=args.model_dir, + ) + + # Generate latents + print(f"\nGenerating {width}x{height} {args.frames}-frame i2v...") + print(f"Prompt: {args.prompt}") + t_start = time.time() + + result = pipeline.generate( + prompt=args.prompt, + num_frames=args.frames, + height=height, + width=width, + num_inference_steps=args.steps, + guidance_scale=args.guidance, + seed=args.seed, + image=img, + enable_audio=args.enable_audio, + negative_prompt=args.n_prompt, + ) + t_gen = time.time() - t_start + + # Free transformer memory before VAE decode + del model + del pipeline + mx.clear_cache() + + if args.verbose and hasattr(mx, "get_peak_memory"): + peak_mem_generation = mx.get_peak_memory() / 1024**3 + mx.reset_peak_memory() + + # Decode video + vae_dir = str(Path(args.model_dir) / "vae") + video = decode_latents(result["latents"], vae_dir) + mx.eval(video) + video_np = np.array(video[0].astype(mx.float32)) + video_np = (video_np * 255).clip(0, 255).astype(np.uint8) + + # Decode audio if generated + audio_np = None + if "audio_latents" in result: + from cosmos3.decode_audio import decode_audio + + snd_dir = str(Path(args.model_dir) / "sound_tokenizer") + audio_waveform = decode_audio(result["audio_latents"], snd_dir) + mx.eval(audio_waveform) + audio_np = np.array(audio_waveform[0].astype(mx.float32)) + + # Save + saved_path = save_video(video_np, args.output, fps=24, audio_waveform=audio_np) + print(f"\nSaved to {saved_path} ({t_gen:.1f}s generation)") + + if args.verbose and hasattr(mx, "get_peak_memory"): + peak_mem_decoding = mx.get_peak_memory() / 1024**3 + print(f"Peak memory generation: {peak_mem_generation:.3f}GB") + print(f"Peak memory decoding: {peak_mem_decoding:.3f}GB") diff --git a/video/cosmos3/requirements.txt b/video/cosmos3/requirements.txt new file mode 100644 index 000000000..e0619f737 --- /dev/null +++ b/video/cosmos3/requirements.txt @@ -0,0 +1,5 @@ +huggingface_hub +mlx>=0.31.0 +numpy +Pillow +transformers diff --git a/video/cosmos3/txt2video.py b/video/cosmos3/txt2video.py new file mode 100644 index 000000000..70d0ba4a5 --- /dev/null +++ b/video/cosmos3/txt2video.py @@ -0,0 +1,152 @@ +"""Generate videos from text using Cosmos 3 Nano on MLX.""" + +import argparse +import time +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from cosmos3.load import load_transformer, load_tokenizer +from cosmos3.pipeline import Cosmos3GenerationPipeline, save_video +from cosmos3.decode_vae import decode_latents + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Generate videos from text using Cosmos 3 Nano" + ) + parser.add_argument("prompt") + parser.add_argument( + "--model-dir", + type=str, + default="weights/Cosmos3-Nano", + help="Path to Cosmos3-Nano weights directory", + ) + parser.add_argument( + "--size", + type=lambda x: tuple(map(int, x.split("x"))), + default=(256, 256), + help="Video size as WxH (default: 256x256)", + ) + parser.add_argument("--frames", type=int, default=16) + parser.add_argument( + "--steps", type=int, default=30, help="Number of denoising steps" + ) + parser.add_argument("--guidance", type=float, default=6.0) + parser.add_argument("--seed", type=int) + parser.add_argument( + "--quantize", + "-q", + type=int, + nargs="?", + const=8, + default=0, + choices=[0, 8], + metavar="{8}", + help="Quantize model weights to 8-bit (default when flag used without value)", + ) + parser.add_argument( + "--n-prompt", + default=None, + help="Negative prompt (default: model's built-in negative prompt)", + ) + parser.add_argument("--output", default="out.mp4") + parser.add_argument("--enable-audio", action="store_true", help="Generate audio") + parser.add_argument( + "--no-cache", + action="store_true", + help="Disable Metal buffer cache (mx.set_cache_limit(0)) to reduce swap pressure", + ) + parser.add_argument("--verbose", "-v", action="store_true") + args = parser.parse_args() + + width, height = args.size + + if args.frames < 8 or args.frames % 4 != 0: + parser.error("--frames must be a multiple of 4 and >= 8 (e.g., 8, 16, 32)") + if width % 16 != 0 or height % 16 != 0: + parser.error(f"--size dimensions must be multiples of 16 (got {width}x{height})") + + mx.set_default_device(mx.gpu) + if args.no_cache: + mx.set_cache_limit(0) + + # Preflight: check required assets before expensive model load + vae_dir = Path(args.model_dir) / "vae" + if not (vae_dir / "config.json").exists(): + parser.error(f"VAE weights not found at {vae_dir}. Download the full model first.") + if args.enable_audio and not (Path(args.model_dir) / "sound_tokenizer" / "config.json").exists(): + parser.error(f"Sound tokenizer not found. Download the full model or disable --enable-audio.") + + # Load model + print(f"Loading model from {args.model_dir}...") + t0 = time.time() + model = load_transformer(args.model_dir, reasoner_only=False) + tokenizer = load_tokenizer(args.model_dir) + + if args.quantize: + nn.quantize(model, bits=args.quantize) + mx.eval(model.parameters()) + print(f"Quantized to {args.quantize}-bit") + + print(f"Model loaded in {time.time() - t0:.1f}s") + + pipeline = Cosmos3GenerationPipeline( + model=model, + tokenizer=tokenizer, + model_dir=args.model_dir, + ) + + # Generate latents + print(f"\nGenerating {width}x{height} {args.frames}-frame video...") + print(f"Prompt: {args.prompt}") + t_start = time.time() + + result = pipeline.generate( + prompt=args.prompt, + num_frames=args.frames, + height=height, + width=width, + num_inference_steps=args.steps, + guidance_scale=args.guidance, + seed=args.seed, + enable_audio=args.enable_audio, + negative_prompt=args.n_prompt, + ) + t_gen = time.time() - t_start + + # Free transformer memory before VAE decode + del model + del pipeline + mx.clear_cache() + + if args.verbose and hasattr(mx, "get_peak_memory"): + peak_mem_generation = mx.get_peak_memory() / 1024**3 + mx.reset_peak_memory() + + # Decode video + vae_dir = str(Path(args.model_dir) / "vae") + video = decode_latents(result["latents"], vae_dir) + mx.eval(video) + video_np = np.array(video[0].astype(mx.float32)) + video_np = (video_np * 255).clip(0, 255).astype(np.uint8) + + # Decode audio if generated + audio_np = None + if "audio_latents" in result: + from cosmos3.decode_audio import decode_audio + + snd_dir = str(Path(args.model_dir) / "sound_tokenizer") + audio_waveform = decode_audio(result["audio_latents"], snd_dir) + mx.eval(audio_waveform) + audio_np = np.array(audio_waveform[0].astype(mx.float32)) + + # Save + saved_path = save_video(video_np, args.output, fps=24, audio_waveform=audio_np) + print(f"\nSaved to {saved_path} ({t_gen:.1f}s generation)") + + if args.verbose and hasattr(mx, "get_peak_memory"): + peak_mem_decoding = mx.get_peak_memory() / 1024**3 + print(f"Peak memory generation: {peak_mem_generation:.3f}GB") + print(f"Peak memory decoding: {peak_mem_decoding:.3f}GB")