diff --git a/.gitignore b/.gitignore
index 0734cada6..3b3663ec7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -39,7 +39,6 @@ converted_weights/
# SSIM test outputs
fastvideo/tests/ssim/generated_videos/
-fastvideo/tests/ssim/.reference_videos_download.lock
**/.cache/**
@@ -133,4 +132,3 @@ fastvideo/tests/ssim/reference_videos/**
*.nvimlog
.nvimlog
.python-version
-/DreamX-World/
diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml
index df4a3aaa7..55fc5f158 100644
--- a/docs/design/inference_schema_parity_inventory.yaml
+++ b/docs/design/inference_schema_parity_inventory.yaml
@@ -191,6 +191,9 @@ surfaces:
sources:
- fastvideo.configs.pipelines.gen3c.Gen3CConfig
- fastvideo.configs.pipelines.gen3c.Gen3CInferenceConfig
+ color_correction_strength:
+ sources:
+ - fastvideo.configs.pipelines.dreamx_world.DreamXWorld5BARPipelineConfig
default_camera_rotation:
sources:
- fastvideo.configs.pipelines.gen3c.Gen3CConfig
diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md
index ec4fbf046..7017e11a8 100644
--- a/docs/inference/support_matrix.md
+++ b/docs/inference/support_matrix.md
@@ -58,6 +58,8 @@ pipeline initialization and sampling.
| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
| FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
+| DreamX-World 5B Cam | `GD-ML/DreamX-World-5B-Cam` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| DreamX-World 5B AR | `GD-ML/DreamX-World-5B` | 704px1280p | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P
720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P
720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py
index 4171fdcc2..5ae49b488 100644
--- a/fastvideo/configs/models/dits/__init__.py
+++ b/fastvideo/configs/models/dits/__init__.py
@@ -1,6 +1,6 @@
from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
-from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig
+from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig, DreamXWorldConfig
from fastvideo.configs.models.dits.flux_2 import Flux2Config
from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
@@ -14,7 +14,7 @@ from fastvideo.configs.models.dits.hyworld import HYWorldConfig
from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
__all__ = [
- "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "CosmosVideoConfig",
- "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
- "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
+ "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig",
+ "DreamXWorldARConfig", "CosmosVideoConfig", "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig",
+ "HYWorldConfig", "Kandinsky5VideoConfig", "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
]
diff --git a/fastvideo/configs/models/dits/dreamx_world.py b/fastvideo/configs/models/dits/dreamx_world.py
index 286e2c3ef..7c5839fbc 100644
--- a/fastvideo/configs/models/dits/dreamx_world.py
+++ b/fastvideo/configs/models/dits/dreamx_world.py
@@ -20,3 +20,30 @@ class DreamXWorldConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=DreamXWorldArchConfig)
prefix: str = "Wan"
+
+
+@dataclass
+class DreamXWorldARArchConfig(DreamXWorldArchConfig):
+ """DreamX-World-5B autoregressive causal DiT config."""
+
+ model_type: str = "ti2v"
+ patch_size: tuple[int, int, int] = (1, 2, 2)
+ text_len: int = 512
+ text_dim: int = 4096
+ freq_dim: int = 256
+ attn_compress: int = 4
+ cam_self_attn_layers: tuple[int, ...] | None = tuple(range(30))
+ local_attn_size: int = 12
+ sink_size: int = 3
+ num_frames_per_block: int = 3
+ rope_cache_policy: str = "block_relativistic"
+ param_names_mapping: dict = field(default_factory=dict)
+ reverse_param_names_mapping: dict = field(default_factory=dict)
+ lora_param_names_mapping: dict = field(default_factory=dict)
+
+
+@dataclass
+class DreamXWorldARConfig(DreamXWorldConfig):
+ arch_config: DiTArchConfig = field(default_factory=DreamXWorldARArchConfig)
+
+ prefix: str = "Wan"
diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py
index 68fb8f3e6..b1e2a441e 100644
--- a/fastvideo/configs/pipelines/__init__.py
+++ b/fastvideo/configs/pipelines/__init__.py
@@ -1,7 +1,7 @@
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.configs.pipelines.cosmos import CosmosConfig
from fastvideo.configs.pipelines.cosmos2_5 import Cosmos25Config
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
+from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
@@ -17,5 +17,6 @@ __all__ = [
"HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
"Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
"SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
- "DreamXWorld5BCamPipelineConfig", "HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+ "DreamXWorld5BCamPipelineConfig", "DreamXWorld5BARPipelineConfig", "HYWorldConfig", "MatrixGame2I2V480PConfig",
+ "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
]
diff --git a/fastvideo/configs/pipelines/dreamx_world.py b/fastvideo/configs/pipelines/dreamx_world.py
index 3e6b8d140..6c3c1e8ef 100644
--- a/fastvideo/configs/pipelines/dreamx_world.py
+++ b/fastvideo/configs/pipelines/dreamx_world.py
@@ -7,8 +7,8 @@ from dataclasses import dataclass, field
import torch
from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
-from fastvideo.configs.models.dits.dreamx_world import (
- DreamXWorldArchConfig, DreamXWorldConfig)
+from fastvideo.configs.models.dits.dreamx_world import (DreamXWorldARArchConfig, DreamXWorldARConfig,
+ DreamXWorldArchConfig, DreamXWorldConfig)
from fastvideo.configs.models.encoders import BaseEncoderOutput, T5Config
from fastvideo.configs.models.encoders.t5 import T5ArchConfig
from fastvideo.configs.models.vaes import WanVAEConfig
@@ -18,21 +18,42 @@ from fastvideo.configs.pipelines.wan import LucyEditDevConfig, t5_postprocess_te
def make_dreamx_world_5b_cam_dit_config() -> DreamXWorldConfig:
"""Return the DreamX-World DiT config matching DreamX-World-5B-Cam."""
- return DreamXWorldConfig(
- arch_config=DreamXWorldArchConfig(
- num_attention_heads=24,
- attention_head_dim=128,
- in_channels=48,
- out_channels=48,
- ffn_dim=14336,
- num_layers=30,
- cross_attn_norm=True,
- qk_norm="rms_norm_across_heads",
- add_control_adapter=True,
- cam_method="prope",
- attn_compress=1,
- cam_self_attn_layers=None,
- ))
+ return DreamXWorldConfig(arch_config=DreamXWorldArchConfig(
+ num_attention_heads=24,
+ attention_head_dim=128,
+ in_channels=48,
+ out_channels=48,
+ ffn_dim=14336,
+ num_layers=30,
+ cross_attn_norm=True,
+ qk_norm="rms_norm_across_heads",
+ add_control_adapter=True,
+ cam_method="prope",
+ attn_compress=1,
+ cam_self_attn_layers=None,
+ ))
+
+
+def make_dreamx_world_5b_ar_dit_config() -> DreamXWorldARConfig:
+ """Return the DreamX-World-5B autoregressive causal DiT config."""
+ return DreamXWorldARConfig(arch_config=DreamXWorldARArchConfig(
+ model_type="ti2v",
+ num_attention_heads=24,
+ attention_head_dim=128,
+ in_channels=48,
+ out_channels=48,
+ ffn_dim=14336,
+ num_layers=30,
+ cross_attn_norm=True,
+ qk_norm=True,
+ add_control_adapter=True,
+ cam_method="prope",
+ attn_compress=4,
+ cam_self_attn_layers=tuple(range(30)),
+ local_attn_size=12,
+ sink_size=3,
+ num_frames_per_block=3,
+ ))
def make_dreamx_world_5b_cam_vae_config() -> WanVAEConfig:
@@ -67,10 +88,10 @@ class DreamXWorld5BCamPipelineConfig(PipelineConfig):
dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_cam_dit_config)
vae_config: VAEConfig = field(default_factory=make_dreamx_world_5b_cam_vae_config)
- text_encoder_configs: tuple[EncoderConfig, ...] = field(
- default_factory=lambda: (make_dreamx_world_5b_cam_text_encoder_config(), ))
- postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
- default_factory=lambda: (t5_postprocess_text, ))
+ text_encoder_configs: tuple[EncoderConfig,
+ ...] = field(default_factory=lambda: (make_dreamx_world_5b_cam_text_encoder_config(), ))
+ postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+ ...] = field(default_factory=lambda: (t5_postprocess_text, ))
text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
flow_shift: float | None = 3.0
ti2v_task: bool = True
@@ -85,3 +106,22 @@ class DreamXWorld5BCamPipelineConfig(PipelineConfig):
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
self.dit_config.expand_timesteps = self.expand_timesteps
+
+
+@dataclass
+class DreamXWorld5BARPipelineConfig(DreamXWorld5BCamPipelineConfig):
+ """Pipeline config for DreamX-World-5B autoregressive forcing."""
+
+ dit_config: DiTConfig = field(default_factory=make_dreamx_world_5b_ar_dit_config)
+ flow_shift: float | None = 5.0
+ ti2v_task: bool = True
+ is_causal: bool = True
+ dmd_denoising_steps: tuple[int, ...] = (1000, 750, 500, 250)
+ warp_denoising_step: bool = True
+ context_noise: float = 0.1
+ num_frames_per_block: int = 3
+ color_correction_strength: float = 1.0
+
+ def __post_init__(self) -> None:
+ super().__post_init__()
+ self.dit_config.expand_timesteps = True
diff --git a/fastvideo/models/dits/dreamx_world_ar.py b/fastvideo/models/dits/dreamx_world_ar.py
new file mode 100644
index 000000000..144187bbe
--- /dev/null
+++ b/fastvideo/models/dits/dreamx_world_ar.py
@@ -0,0 +1,872 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamX-World autoregressive causal DiT.
+
+Adapted from DreamX-World's Apache-2.0
+``wan/modules/causal_camera_model_2_2_prope_infinity.py``. The implementation is
+kept native to FastVideo: no production import from DreamX, Diffusers, or
+Transformers is required.
+"""
+
+from typing import Any
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARConfig
+from fastvideo.models.dits.base import BaseDiT
+from fastvideo.models.dits.dreamx_world import (_dreamx_apply_tiled_projmat,
+ _dreamx_prope_qkv)
+
+
+def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+ q_bhld = q.transpose(1, 2)
+ k_bhld = k.transpose(1, 2)
+ v_bhld = v.transpose(1, 2)
+ out = F.scaled_dot_product_attention(q_bhld, k_bhld, v_bhld, dropout_p=0.0)
+ return out.transpose(1, 2)
+
+
+def prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
+ viewmats: torch.Tensor, Ks: torch.Tensor):
+ q, k, v, output_projection = _dreamx_prope_qkv(q, k, v, viewmats, Ks)
+
+ def apply_fn_o(x: torch.Tensor) -> torch.Tensor:
+ return _dreamx_apply_tiled_projmat(x, output_projection)
+
+ return q, k, v, apply_fn_o
+
+
+def sinusoidal_embedding_1d(dim, position):
+ assert dim % 2 == 0
+ half = dim // 2
+ position = position.type(torch.float64)
+ sinusoid = torch.outer(
+ position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
+ return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+
+
+def rope_params(max_seq_len, dim, theta=10000):
+ assert dim % 2 == 0
+ freqs = torch.outer(
+ torch.arange(max_seq_len),
+ 1.0 / torch.pow(theta,
+ torch.arange(0, dim, 2).to(torch.float64).div(dim)))
+ return torch.polar(torch.ones_like(freqs), freqs)
+
+
+class WanRMSNorm(nn.Module):
+ def __init__(self, dim, eps=1e-5):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x):
+ return self._norm(x).type_as(x) * self.weight
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
+
+
+class WanLayerNorm(nn.LayerNorm):
+ def __init__(self, dim, eps=1e-6, elementwise_affine=False):
+ super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
+
+ def forward(self, x):
+ return super().forward(x).type_as(x)
+
+
+class WanCrossAttention(nn.Module):
+ def __init__(self, dim, num_heads, window_size=(-1, -1), qk_norm=True, eps=1e-6):
+ super().__init__()
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.q = nn.Linear(dim, dim)
+ self.k = nn.Linear(dim, dim)
+ self.v = nn.Linear(dim, dim)
+ self.o = nn.Linear(dim, dim)
+ self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, context_lens, crossattn_cache=None):
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+ q = self.norm_q(self.q(x)).view(b, -1, n, d)
+
+ if crossattn_cache is not None:
+ if not crossattn_cache["is_init"]:
+ crossattn_cache["is_init"] = True
+ k = self.norm_k(self.k(context)).view(b, -1, n, d)
+ v = self.v(context).view(b, -1, n, d)
+ crossattn_cache["k"] = k
+ crossattn_cache["v"] = v
+ else:
+ k = crossattn_cache["k"]
+ v = crossattn_cache["v"]
+ else:
+ k = self.norm_k(self.k(context)).view(b, -1, n, d)
+ v = self.v(context).view(b, -1, n, d)
+
+ x = attention(q, k, v)
+ x = x.flatten(2)
+ return self.o(x)
+
+
+def block_relativistic_rope(x, grid_sizes, freqs, start_frame=0, relative_frame_indices=None):
+ """
+ Apply Block-Relativistic RoPE to input tensor.
+ Adapted from Infinity-RoPE (https://arxiv.org/abs/2511.20649).
+
+ Args:
+ x: Input tensor [B, L, num_heads, head_dim]
+ grid_sizes: Tensor [B, 3] containing (F, H, W)
+ freqs: RoPE frequencies
+ start_frame: Starting frame index for sequential RoPE
+ relative_frame_indices: Optional tensor [F] specifying explicit frame indices
+ for Block-Relativistic RoPE. Overrides start_frame if provided.
+ """
+ n, c = x.size(2), x.size(3) // 2
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+ x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(
+ seq_len, n, -1, 2))
+
+ if relative_frame_indices is not None:
+ frame_indices = relative_frame_indices.long()
+ freqs_temporal = freqs[0][frame_indices].view(f, 1, 1, -1).expand(f, h, w, -1)
+ else:
+ freqs_temporal = freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1)
+
+ freqs_i = torch.cat([
+ freqs_temporal,
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ], dim=-1).reshape(seq_len, 1, -1)
+
+ x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+ x_i = torch.cat([x_i, x[i, seq_len:]])
+ output.append(x_i)
+
+ return torch.stack(output).type_as(x)
+
+
+class CausalWanSelfAttention(nn.Module):
+ """Self-attention with KV cache and Block-Relativistic RoPE for causal inference."""
+
+ def __init__(self, dim, num_heads, local_attn_size=6, sink_size=1,
+ qk_norm=True, eps=1e-6):
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.local_attn_size = local_attn_size
+ self.sink_size = sink_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+ self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
+
+ self.q = nn.Linear(dim, dim)
+ self.k = nn.Linear(dim, dim)
+ self.v = nn.Linear(dim, dim)
+ self.o = nn.Linear(dim, dim)
+ self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, seq_lens, grid_sizes, freqs, kv_cache,
+ current_start=0, cache_start=None, sink_recache_after_switch=False):
+ """
+ Args:
+ x: Shape [B, L, C]
+ seq_lens: Shape [B]
+ grid_sizes: Shape [B, 3] containing (F, H, W)
+ freqs: RoPE frequencies [1024, head_dim / 2]
+ kv_cache: Dict with 'k', 'v', 'global_end_index', 'local_end_index'
+ current_start: Current position in the global token sequence
+ """
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+ if cache_start is None:
+ cache_start = current_start
+
+ q = self.norm_q(self.q(x)).view(b, s, n, d)
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ v = self.v(x).view(b, s, n, d)
+
+ frame_seqlen = math.prod(grid_sizes[0][1:]).item()
+ num_new_frames = grid_sizes[0][0].item()
+ current_end = current_start + q.shape[1]
+ sink_tokens = self.sink_size * frame_seqlen
+ kv_cache_size = kv_cache["k"].shape[1]
+ num_new_tokens = q.shape[1]
+
+ cache_update_info = None
+ is_recompute = current_end <= kv_cache["global_end_index"].item() and current_start > 0
+
+ if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
+ num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
+ # === ROLLING MODE: cache full, evict oldest non-sink tokens ===
+ num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+ num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
+ local_end_index = kv_cache["local_end_index"].item() + current_end - \
+ kv_cache["global_end_index"].item() - num_evicted_tokens
+ local_start_index = local_end_index - num_new_tokens
+
+ temp_k = kv_cache["k"].detach().clone()
+ temp_v = kv_cache["v"].detach().clone()
+ temp_k[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ temp_k[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+ temp_v[:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ temp_v[:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+
+ write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
+ roped_offset = max(0, write_start_index - local_start_index)
+ write_len = max(0, local_end_index - write_start_index)
+ if write_len > 0:
+ temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
+ temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
+
+ # Block-Relativistic RoPE: query uses window-relative indices
+ query_relative_indices = torch.arange(
+ self.local_attn_size - num_new_frames, self.local_attn_size, device=q.device)
+ roped_query = block_relativistic_rope(
+ q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
+
+ # Block-Relativistic RoPE: cached K uses position-in-window indices
+ num_cache_frames = local_end_index // frame_seqlen
+ cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
+ cache_grid_sizes = grid_sizes.clone()
+ cache_grid_sizes[0, 0] = num_cache_frames
+ roped_temp_k = block_relativistic_rope(
+ temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
+ cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
+
+ cache_update_info = {
+ "action": "roll_and_insert",
+ "sink_tokens": sink_tokens,
+ "num_rolled_tokens": num_rolled_tokens,
+ "num_evicted_tokens": num_evicted_tokens,
+ "local_start_index": local_start_index,
+ "local_end_index": local_end_index,
+ "write_start_index": write_start_index,
+ "write_end_index": local_end_index,
+ "new_k": k[:, roped_offset:roped_offset + write_len],
+ "new_v": v[:, roped_offset:roped_offset + write_len],
+ "current_end": current_end,
+ "is_recompute": is_recompute
+ }
+ else:
+ # === DIRECT INSERT MODE: cache not yet full ===
+ local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
+ local_start_index = local_end_index - num_new_tokens
+
+ temp_k = kv_cache["k"].detach().clone()
+ temp_v = kv_cache["v"].detach().clone()
+
+ write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
+ if sink_recache_after_switch:
+ write_start_index = local_start_index
+ roped_offset = max(0, write_start_index - local_start_index)
+ write_len = max(0, local_end_index - write_start_index)
+ if write_len > 0:
+ temp_k[:, write_start_index:local_end_index] = k[:, roped_offset:roped_offset + write_len]
+ temp_v[:, write_start_index:local_end_index] = v[:, roped_offset:roped_offset + write_len]
+
+ # RoPE with relative indices (growing sequentially before cache fills)
+ current_frame_in_window = local_start_index // frame_seqlen
+ query_relative_indices = torch.arange(
+ current_frame_in_window, current_frame_in_window + num_new_frames, device=q.device)
+ roped_query = block_relativistic_rope(
+ q, grid_sizes, freqs, relative_frame_indices=query_relative_indices).type_as(v)
+
+ num_cache_frames = local_end_index // frame_seqlen
+ cache_relative_indices = torch.arange(0, num_cache_frames, device=k.device)
+ cache_grid_sizes = grid_sizes.clone()
+ cache_grid_sizes[0, 0] = num_cache_frames
+ roped_temp_k = block_relativistic_rope(
+ temp_k[:, :local_end_index].view(b, num_cache_frames, frame_seqlen, n, d).flatten(1, 2),
+ cache_grid_sizes, freqs, relative_frame_indices=cache_relative_indices).type_as(v)
+
+ cache_update_info = {
+ "action": "direct_insert",
+ "local_start_index": local_start_index,
+ "local_end_index": local_end_index,
+ "write_start_index": write_start_index,
+ "write_end_index": local_end_index,
+ "new_k": k[:, roped_offset:roped_offset + write_len],
+ "new_v": v[:, roped_offset:roped_offset + write_len],
+ "current_end": current_end,
+ "is_recompute": is_recompute
+ }
+
+ # Attention: sink tokens + local window
+ if sink_tokens > 0:
+ local_budget = self.max_attention_size - sink_tokens
+ k_sink = roped_temp_k[:, :sink_tokens]
+ v_sink = temp_v[:, :sink_tokens]
+ if local_budget > 0:
+ local_start_for_window = max(sink_tokens, local_end_index - local_budget)
+ k_local = roped_temp_k[:, local_start_for_window:local_end_index]
+ v_local = temp_v[:, local_start_for_window:local_end_index]
+ k_cat = torch.cat([k_sink, k_local], dim=1)
+ v_cat = torch.cat([v_sink, v_local], dim=1)
+ else:
+ k_cat = k_sink
+ v_cat = v_sink
+ x = attention(roped_query, k_cat, v_cat)
+ else:
+ window_start = max(0, local_end_index - self.max_attention_size)
+ x = attention(
+ roped_query,
+ roped_temp_k[:, window_start:local_end_index],
+ temp_v[:, window_start:local_end_index])
+
+ x = x.flatten(2)
+ x = self.o(x)
+ return x, (current_end, local_end_index, cache_update_info)
+
+
+class CausalPropeSelfAttention(nn.Module):
+ """PRoPE self-attention with optional KV cache for camera-controlled inference."""
+
+ def __init__(self, dim, attn_dim, num_heads, window_size=(-1, -1),
+ local_attn_size=-1, sink_size=0, qk_norm=True, eps=1e-6):
+ assert dim % num_heads == 0
+ assert attn_dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.attn_dim = attn_dim
+ self.num_heads = num_heads
+ self.head_dim = attn_dim // num_heads
+ self.local_attn_size = local_attn_size
+ self.sink_size = sink_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+ self.window_size = window_size
+ self.max_attention_size = 39600 if local_attn_size == -1 else local_attn_size * 880
+
+ self.q_proj = nn.Linear(dim, attn_dim)
+ self.k_proj = nn.Linear(dim, attn_dim)
+ self.v_proj = nn.Linear(dim, attn_dim)
+ self.out_proj = nn.Linear(attn_dim, dim)
+
+ self.norm_q = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(attn_dim, eps=eps) if qk_norm else nn.Identity()
+
+ nn.init.zeros_(self.out_proj.weight)
+ nn.init.zeros_(self.out_proj.bias)
+
+ def forward(self, x, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
+ kv_cache=None, current_start=0, cache_start=None,
+ sink_recache_after_switch=False, cache_update_policy="commit_detached"):
+ """
+ Args:
+ x: Shape [B, L, C]
+ cam_viewmats: Camera view matrices
+ cam_K: Camera intrinsics
+ kv_cache: Optional KV cache dict. When None, runs full attention over current chunk.
+ """
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+ if cache_start is None:
+ cache_start = current_start
+
+ q = self.norm_q(self.q_proj(x)).view(b, s, n, d)
+ k = self.norm_k(self.k_proj(x)).view(b, s, n, d)
+ v = self.v_proj(x).view(b, s, n, d)
+
+ # Apply PRoPE (Positional Rotary Position Embedding from camera parameters)
+ q_t, k_t, v_t, apply_fn_o = prope_qkv(
+ q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
+ viewmats=cam_viewmats, Ks=cam_K)
+ proped_q = q_t.transpose(1, 2)
+ proped_k = k_t.transpose(1, 2)
+ proped_v = v_t.transpose(1, 2)
+
+ if kv_cache is None:
+ # No cache: full attention over current chunk
+ x_out = attention(proped_q, proped_k, proped_v)
+ else:
+ # KV cache mode with rolling cache support
+ frame_seqlen = math.prod(grid_sizes[0][1:]).item()
+ num_new_tokens = s
+ current_end = current_start + num_new_tokens
+ sink_tokens = self.sink_size * frame_seqlen
+ kv_cache_size = kv_cache["k"].shape[1]
+ is_recompute = (current_end <= kv_cache["global_end_index"].item()) and (current_start > 0)
+
+ if self.local_attn_size != -1 and (current_end > kv_cache["global_end_index"].item()) and (
+ num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size):
+ # === ROLLING MODE ===
+ num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size
+ num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens
+ local_end_index = kv_cache["local_end_index"].item() + current_end - \
+ kv_cache["global_end_index"].item() - num_evicted_tokens
+ local_start_index = local_end_index - num_new_tokens
+
+ if cache_update_policy != "none":
+ with torch.no_grad():
+ kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+ kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+
+ write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
+ roped_offset = max(0, write_start_index - local_start_index)
+ write_len = max(0, local_end_index - write_start_index)
+ if write_len > 0:
+ with torch.no_grad():
+ kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
+ kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
+ else:
+ # === DIRECT INSERT MODE ===
+ local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item()
+ local_start_index = local_end_index - num_new_tokens
+
+ if cache_update_policy != "none":
+ write_start_index = max(local_start_index, sink_tokens) if is_recompute else local_start_index
+ if sink_recache_after_switch:
+ write_start_index = local_start_index
+ roped_offset = max(0, write_start_index - local_start_index)
+ write_len = max(0, local_end_index - write_start_index)
+ if write_len > 0:
+ with torch.no_grad():
+ kv_cache["k"][:, write_start_index:local_end_index] = proped_k[:, roped_offset:roped_offset + write_len].detach()
+ kv_cache["v"][:, write_start_index:local_end_index] = proped_v[:, roped_offset:roped_offset + write_len].detach()
+
+ # Attention: sink tokens + local window
+ if sink_tokens > 0:
+ local_budget = self.max_attention_size - sink_tokens
+ k_sink = kv_cache["k"][:, :sink_tokens].detach()
+ v_sink = kv_cache["v"][:, :sink_tokens].detach()
+ if local_budget > 0:
+ local_start_for_window = max(sink_tokens, local_end_index - local_budget)
+ k_local = kv_cache["k"][:, local_start_for_window:local_end_index].detach()
+ v_local = kv_cache["v"][:, local_start_for_window:local_end_index].detach()
+ k_cat = torch.cat([k_sink, k_local], dim=1)
+ v_cat = torch.cat([v_sink, v_local], dim=1)
+ else:
+ k_cat = k_sink
+ v_cat = v_sink
+ x_out = attention(proped_q, k_cat, v_cat)
+ else:
+ window_start = max(0, local_end_index - self.max_attention_size)
+ x_out = attention(
+ proped_q,
+ kv_cache["k"][:, window_start:local_end_index].detach(),
+ kv_cache["v"][:, window_start:local_end_index].detach())
+
+ if not is_recompute and cache_update_policy != "none":
+ kv_cache["global_end_index"].fill_(current_end)
+ kv_cache["local_end_index"].fill_(local_end_index)
+
+ # Apply inverse PRoPE
+ x = apply_fn_o(x_out.transpose(1, 2)).transpose(1, 2)
+ x = x.flatten(2)
+ x = self.out_proj(x)
+ return x
+
+
+class CausalWanAttentionBlock(nn.Module):
+
+ def __init__(self, dim, ffn_dim, num_heads, local_attn_size=-1, sink_size=0,
+ qk_norm=True, cross_attn_norm=False, eps=1e-6, **kwargs):
+ super().__init__()
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.num_heads = num_heads
+ self.local_attn_size = local_attn_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+
+ self.add_control_adapter = kwargs.get('add_control_adapter', False)
+ self.cam_method = kwargs.get('cam_method')
+ self.attn_compress = kwargs.get('attn_compress', 1)
+ self.layer_idx = kwargs.get('layer_idx')
+ cam_self_attn_layers = kwargs.get('cam_self_attn_layers')
+
+ # layers
+ self.norm1 = WanLayerNorm(dim, eps)
+ self.self_attn = CausalWanSelfAttention(
+ dim, num_heads, local_attn_size, sink_size, qk_norm, eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
+ self.cross_attn = WanCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+
+ # PRoPE self-attention branch for camera control
+ add_cam_attn = self.add_control_adapter and self.cam_method == 'prope'
+ if add_cam_attn and cam_self_attn_layers is not None:
+ add_cam_attn = self.layer_idx in cam_self_attn_layers
+ if add_cam_attn:
+ self.cam_self_attn = CausalPropeSelfAttention(
+ dim, dim // self.attn_compress, num_heads,
+ local_attn_size=local_attn_size, sink_size=sink_size,
+ qk_norm=qk_norm, eps=eps)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+ def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens,
+ kv_cache, crossattn_cache=None, current_start=0, cache_start=None,
+ cam_viewmats=None, cam_K=None, sink_recache_after_switch=False,
+ cache_update_policy="commit_detached"):
+ num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+ e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2)
+
+ # self-attention
+ attn_input = (self.norm1(x).unflatten(
+ dim=1, sizes=(num_frames, frame_seqlen)) * (1 + e[1]) + e[0]).flatten(1, 2)
+ y, cache_update_info = self.self_attn(
+ attn_input, seq_lens, grid_sizes, freqs, kv_cache,
+ current_start, cache_start, sink_recache_after_switch)
+
+ # PRoPE camera attention (parallel branch)
+ if hasattr(self, 'cam_self_attn') and cam_viewmats is not None and cam_K is not None:
+ prope_kv_cache = None
+ if kv_cache is not None and "prope_k" in kv_cache:
+ prope_kv_cache = {
+ "k": kv_cache["prope_k"],
+ "v": kv_cache["prope_v"],
+ "global_end_index": kv_cache["prope_global_end_index"],
+ "local_end_index": kv_cache["prope_local_end_index"],
+ }
+ y = y + self.cam_self_attn(
+ attn_input, cam_viewmats, cam_K, seq_lens, grid_sizes, freqs,
+ kv_cache=prope_kv_cache, current_start=current_start,
+ cache_start=cache_start, cache_update_policy=cache_update_policy)
+
+ x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[2]).flatten(1, 2)
+
+ # cross-attention & FFN
+ x = x + self.cross_attn(self.norm3(x), context, context_lens,
+ crossattn_cache=crossattn_cache)
+ y = self.ffn(
+ (self.norm2(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+ * (1 + e[4]) + e[3]).flatten(1, 2))
+ x = x + (y.unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * e[5]).flatten(1, 2)
+
+ return x, cache_update_info
+
+
+class CausalHead(nn.Module):
+
+ def __init__(self, dim, out_dim, patch_size, eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.out_dim = out_dim
+ self.patch_size = patch_size
+ self.eps = eps
+
+ out_dim = math.prod(patch_size) * out_dim
+ self.norm = WanLayerNorm(dim, eps)
+ self.head = nn.Linear(dim, out_dim)
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+ def forward(self, x, e):
+ num_frames, frame_seqlen = e.shape[1], x.shape[1] // e.shape[1]
+ e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2)
+ x = self.head(
+ self.norm(x).unflatten(dim=1, sizes=(num_frames, frame_seqlen))
+ * (1 + e[1]) + e[0])
+ return x
+
+
+class DreamXWorldARTransformer3DModel(BaseDiT):
+ """DreamX-World-5B autoregressive causal transformer."""
+
+ _fsdp_shard_conditions = DreamXWorldARConfig()._fsdp_shard_conditions
+ _compile_conditions = DreamXWorldARConfig()._compile_conditions
+ _supported_attention_backends = DreamXWorldARConfig()._supported_attention_backends
+ param_names_mapping = DreamXWorldARConfig().param_names_mapping
+ reverse_param_names_mapping = DreamXWorldARConfig().reverse_param_names_mapping
+ lora_param_names_mapping = DreamXWorldARConfig().lora_param_names_mapping
+ _no_split_modules = ["CausalWanAttentionBlock"]
+
+ def __init__(self, config: DreamXWorldARConfig, hf_config: dict[str, Any]) -> None:
+ super().__init__(config=config, hf_config=hf_config)
+
+ model_type = config.model_type
+ patch_size = config.patch_size
+ text_len = config.text_len
+ in_dim = config.in_channels
+ dim = config.hidden_size
+ ffn_dim = config.ffn_dim
+ freq_dim = config.freq_dim
+ text_dim = config.text_dim
+ out_dim = config.out_channels
+ num_heads = config.num_attention_heads
+ num_layers = config.num_layers
+ local_attn_size = config.local_attn_size
+ sink_size = config.sink_size
+ qk_norm = bool(config.qk_norm)
+ cross_attn_norm = config.cross_attn_norm
+ eps = config.eps
+ add_control_adapter = config.add_control_adapter
+ cam_method = config.cam_method
+ attn_compress = config.attn_compress
+ cam_self_attn_layers = config.cam_self_attn_layers
+
+ assert model_type in ['t2v', 'i2v', 'ti2v']
+ self.model_type = model_type
+ self.patch_size = patch_size
+ self.text_len = text_len
+ self.in_dim = in_dim
+ self.dim = dim
+ self.ffn_dim = ffn_dim
+ self.freq_dim = freq_dim
+ self.text_dim = text_dim
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.local_attn_size = local_attn_size
+ self.qk_norm = qk_norm
+ self.cross_attn_norm = cross_attn_norm
+ self.eps = eps
+
+ # embeddings
+ self.patch_embedding = nn.Conv3d(
+ in_dim, dim, kernel_size=patch_size, stride=patch_size)
+ self.text_embedding = nn.Sequential(
+ nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+ nn.Linear(dim, dim))
+ self.time_embedding = nn.Sequential(
+ nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+ self.time_projection = nn.Sequential(
+ nn.SiLU(), nn.Linear(dim, dim * 6))
+
+ # transformer blocks
+ self.blocks = nn.ModuleList([
+ CausalWanAttentionBlock(
+ dim, ffn_dim, num_heads, local_attn_size, sink_size,
+ qk_norm, cross_attn_norm, eps,
+ add_control_adapter=add_control_adapter,
+ cam_method=cam_method,
+ attn_compress=attn_compress,
+ layer_idx=layer_idx,
+ cam_self_attn_layers=cam_self_attn_layers)
+ for layer_idx in range(num_layers)
+ ])
+ for layer_idx, block in enumerate(self.blocks):
+ block.self_attn.layer_idx = layer_idx
+ block.self_attn.num_layers = self.num_layers
+
+ # head
+ self.head = CausalHead(dim, out_dim, patch_size, eps)
+
+ # RoPE frequencies
+ assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+ d = dim // num_heads
+ self.freqs = torch.cat([
+ rope_params(1024, d - 4 * (d // 6)),
+ rope_params(1024, 2 * (d // 6)),
+ rope_params(1024, 2 * (d // 6))
+ ], dim=1)
+
+ self.num_attention_heads = num_heads
+ self.attention_head_dim = dim // num_heads
+ self.hidden_size = dim
+ self.in_channels = in_dim
+ self.out_channels = out_dim
+ self.num_channels_latents = out_dim
+ self.init_weights()
+ self.num_frame_per_block = config.arch_config.num_frames_per_block
+ self.__post_init__()
+
+ def forward(self, x=None, t=None, context=None, seq_len=None, y=None, y_camera=None,
+ kv_cache=None, crossattn_cache=None, current_start=0,
+ cache_start=0, cache_update_policy="commit_detached",
+ hidden_states=None, encoder_hidden_states=None, timestep=None, **kwargs):
+ """
+ Causal inference with KV caching.
+ See Algorithm 2 of CausVid (https://arxiv.org/abs/2412.07772).
+
+ Args:
+ x: List of input video tensors [C_in, F, H, W]
+ t: Timestep tensor [B, L]
+ context: List of text embeddings [L, C]
+ seq_len: Maximum sequence length for positional encoding
+ y: Optional conditional video inputs (I2V mode)
+ y_camera: Camera parameters dict {'viewmats': ..., 'K': ...}
+ kv_cache: List of KV cache dicts per transformer block
+ crossattn_cache: List of cross-attention cache dicts
+ current_start: Current position in global token sequence
+ cache_start: Cache start position
+ cache_update_policy: Cache update strategy ('commit_detached' or 'none')
+
+ Returns:
+ Stacked output tensors [B, C_out, F, H/8, W/8]
+ """
+ if x is None and hidden_states is not None:
+ x = [sample for sample in hidden_states]
+ if t is None and timestep is not None:
+ t = timestep
+ if context is None and encoder_hidden_states is not None:
+ if isinstance(encoder_hidden_states, torch.Tensor):
+ context = [sample for sample in encoder_hidden_states]
+ else:
+ context = encoder_hidden_states
+ if seq_len is None:
+ if torch.is_tensor(t):
+ seq_len = int(t.shape[1]) if t.dim() > 1 else int(t.numel())
+ elif x is not None:
+ sample = x[0]
+ seq_len = (sample.shape[1] // self.patch_size[0]) * (sample.shape[2] // self.patch_size[1]) * (sample.shape[3] // self.patch_size[2])
+ if x is None or t is None or context is None or seq_len is None:
+ raise ValueError("DreamXWorldARTransformer3DModel requires x/t/context/seq_len or FastVideo aliases")
+
+ device = self.patch_embedding.weight.device
+ if self.freqs.is_meta or self.freqs.device != device:
+ d = self.dim // self.num_heads
+ self.freqs = torch.cat([
+ rope_params(1024, d - 4 * (d // 6)),
+ rope_params(1024, 2 * (d // 6)),
+ rope_params(1024, 2 * (d // 6)),
+ ], dim=1).to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y, strict=True)]
+
+ # patch embedding
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ grid_sizes = torch.stack(
+ [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+ x = [u.flatten(2).transpose(1, 2) for u in x]
+ seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+ assert seq_lens.max() <= seq_len
+ x = torch.cat(x)
+
+ # time embedding
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t.flatten()).type_as(x))
+ e0 = self.time_projection(e).unflatten(
+ 1, (6, self.dim)).unflatten(dim=0, sizes=t.shape)
+
+ # text embedding
+ context_lens = None
+ context = self.text_embedding(
+ torch.stack([
+ torch.cat(
+ [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+ for u in context
+ ]))
+
+ # camera parameters
+ if y_camera is not None and isinstance(y_camera, dict):
+ cam_viewmats = y_camera['viewmats']
+ cam_K = y_camera['K']
+ else:
+ cam_viewmats = None
+ cam_K = None
+
+ block_kwargs = dict(
+ e=e0, seq_lens=seq_lens, grid_sizes=grid_sizes, freqs=self.freqs,
+ context=context, context_lens=context_lens,
+ cam_viewmats=cam_viewmats, cam_K=cam_K,
+ cache_update_policy=cache_update_policy,
+ )
+
+ cache_update_infos = []
+ for block_index, block in enumerate(self.blocks):
+ block_kwargs.update({
+ "kv_cache": kv_cache[block_index] if kv_cache is not None else None,
+ "crossattn_cache": crossattn_cache[block_index] if crossattn_cache is not None else None,
+ "current_start": current_start,
+ "cache_start": cache_start,
+ })
+ x, block_cache_update_info = block(x, **block_kwargs)
+ if kv_cache is not None:
+ cache_update_infos.append((block_index, block_cache_update_info))
+
+ # Apply deferred cache updates
+ if kv_cache is not None and cache_update_infos and cache_update_policy != "none":
+ self._apply_cache_updates(kv_cache, cache_update_infos)
+
+ # head & unpatchify
+ x = self.head(x, e.unflatten(dim=0, sizes=t.shape).unsqueeze(2))
+ x = self.unpatchify(x, grid_sizes)
+ return torch.stack(x)
+
+ def _apply_cache_updates(self, kv_cache, cache_update_infos):
+ """Apply deferred cache updates collected from all transformer blocks.
+
+ For Block-Relativistic RoPE, this stores un-roped K values in the cache.
+ RoPE is applied dynamically during attention based on each token's current
+ relative position in the sliding window.
+ """
+ with torch.no_grad():
+ for block_index, (current_end, local_end_index, update_info) in cache_update_infos:
+ if update_info is not None:
+ cache = kv_cache[block_index]
+
+ if update_info["action"] == "roll_and_insert":
+ sink_tokens = update_info["sink_tokens"]
+ num_rolled_tokens = update_info["num_rolled_tokens"]
+ num_evicted_tokens = update_info["num_evicted_tokens"]
+ write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
+ write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
+ new_k = update_info["new_k"].detach()
+ new_v = update_info["new_v"].detach()
+
+ cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+ cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
+ cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
+
+ if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
+ cache["k"][:, write_start_index:write_end_index] = new_k
+ cache["v"][:, write_start_index:write_end_index] = new_v
+
+ elif update_info["action"] == "direct_insert":
+ write_start_index = update_info.get("write_start_index", update_info["local_start_index"])
+ write_end_index = update_info.get("write_end_index", update_info["local_end_index"])
+ new_k = update_info["new_k"].detach()
+ new_v = update_info["new_v"].detach()
+
+ if write_end_index > write_start_index and new_k.shape[1] == (write_end_index - write_start_index):
+ cache["k"][:, write_start_index:write_end_index] = new_k
+ cache["v"][:, write_start_index:write_end_index] = new_v
+
+ is_recompute = False if update_info is None else update_info.get("is_recompute", False)
+ if not is_recompute:
+ kv_cache[block_index]["global_end_index"].fill_(current_end)
+ kv_cache[block_index]["local_end_index"].fill_(local_end_index)
+
+ def unpatchify(self, x, grid_sizes):
+ """Reconstruct video tensors from patch embeddings."""
+ c = self.out_dim
+ out = []
+ for u, v in zip(x, grid_sizes.tolist(), strict=True):
+ u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+ u = torch.einsum('fhwpqrc->cfphqwr', u)
+ u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size, strict=True)])
+ out.append(u)
+ return out
+
+ def init_weights(self):
+ """Initialize model parameters using Xavier initialization."""
+ for m in self.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.xavier_uniform_(m.weight)
+ if m.bias is not None:
+ nn.init.zeros_(m.bias)
+
+ nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+ for m in self.text_embedding.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.normal_(m.weight, std=.02)
+ for m in self.time_embedding.modules():
+ if isinstance(m, nn.Linear):
+ nn.init.normal_(m.weight, std=.02)
+
+ nn.init.zeros_(self.head.head.weight)
+
+
+EntryClass = DreamXWorldARTransformer3DModel
diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py
index c105143dd..69b9d6d0a 100644
--- a/fastvideo/models/registry.py
+++ b/fastvideo/models/registry.py
@@ -33,6 +33,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
("dits", "hyworld", "HYWorldTransformer3DModel"),
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
+ "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
"CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
"Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
@@ -50,6 +51,7 @@ _IMAGE_TO_VIDEO_DIT_MODELS = {
# "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
+ "DreamXWorldARTransformer3DModel": ("dits", "dreamx_world_ar", "DreamXWorldARTransformer3DModel"),
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
"MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
"CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
diff --git a/fastvideo/pipelines/basic/dreamx_world/__init__.py b/fastvideo/pipelines/basic/dreamx_world/__init__.py
index 8dbc4bff1..38c7f6aa2 100644
--- a/fastvideo/pipelines/basic/dreamx_world/__init__.py
+++ b/fastvideo/pipelines/basic/dreamx_world/__init__.py
@@ -1,10 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
from fastvideo.configs.pipelines.dreamx_world import (
+ DreamXWorld5BARPipelineConfig,
DreamXWorld5BCamPipelineConfig,
+ make_dreamx_world_5b_ar_dit_config,
make_dreamx_world_5b_cam_dit_config,
make_dreamx_world_5b_cam_text_encoder_config,
make_dreamx_world_5b_cam_vae_config,
)
+from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
from fastvideo.pipelines.basic.dreamx_world.stages import (
DREAMX_Y_CAMERA_KEY,
@@ -13,9 +16,12 @@ from fastvideo.pipelines.basic.dreamx_world.stages import (
__all__ = [
"DREAMX_Y_CAMERA_KEY",
+ "DreamXWorld5BARPipelineConfig",
"DreamXWorld5BCamPipelineConfig",
"DreamXWorldCameraConditioningStage",
+ "DreamXWorldARPipeline",
"DreamXWorldPipeline",
+ "make_dreamx_world_5b_ar_dit_config",
"make_dreamx_world_5b_cam_dit_config",
"make_dreamx_world_5b_cam_text_encoder_config",
"make_dreamx_world_5b_cam_vae_config",
diff --git a/fastvideo/pipelines/basic/dreamx_world/ar_denoising.py b/fastvideo/pipelines/basic/dreamx_world/ar_denoising.py
new file mode 100644
index 000000000..874c44a5a
--- /dev/null
+++ b/fastvideo/pipelines/basic/dreamx_world/ar_denoising.py
@@ -0,0 +1,219 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamX-World autoregressive causal denoising stage."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from tqdm.auto import tqdm
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.basic.dreamx_world.stages import DREAMX_Y_CAMERA_KEY
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.denoising import DenoisingStage
+
+
+class DreamXWorldARCausalDenoisingStage(DenoisingStage):
+ """Official DreamX AR-forcing denoising loop with KV cache."""
+
+ _AR_NOISE_SEED_OFFSET = 1_000_003
+
+ def __init__(self, transformer, scheduler, pipeline=None, vae=None) -> None:
+ super().__init__(transformer=transformer, scheduler=scheduler, pipeline=pipeline, vae=vae)
+ self.num_transformer_blocks = len(self.transformer.blocks)
+ self.num_frame_per_block = int(getattr(self.transformer, "num_frame_per_block", 3))
+ self.local_attn_size = int(getattr(self.transformer, "local_attn_size", 12))
+
+ def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
+ assert batch.latents is not None, "latents must be prepared before DreamX AR denoising"
+ assert batch.prompt_embeds, "prompt embeds must be prepared before DreamX AR denoising"
+ latents = batch.latents
+ device = latents.device
+ target_dtype = torch.bfloat16
+ autocast_enabled = device.type == "cuda" and not fastvideo_args.disable_autocast
+
+ frame_seq_length = (latents.shape[-2] // self.transformer.patch_size[1]) * (latents.shape[-1] //
+ self.transformer.patch_size[2])
+ timesteps = torch.tensor(
+ tuple(getattr(fastvideo_args.pipeline_config, "dmd_denoising_steps", (1000, 750, 500, 250))),
+ dtype=torch.long,
+ ).cpu()
+ if getattr(fastvideo_args.pipeline_config, "warp_denoising_step", True):
+ self.scheduler.set_timesteps(1000)
+ scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
+ timesteps = scheduler_timesteps[1000 - timesteps]
+ timesteps = timesteps.to(device)
+
+ if latents.shape[2] % self.num_frame_per_block != 0:
+ raise ValueError("DreamX AR latent frames must be divisible by num_frame_per_block")
+
+ y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY, batch.extra.get("y_camera"))
+ if isinstance(y_camera, dict):
+ y_camera = {
+ k: v.to(device=device, dtype=target_dtype) if torch.is_tensor(v) else v
+ for k, v in y_camera.items()
+ }
+
+ if batch.image_latent is not None and batch.image_latent.shape[1] == latents.shape[1]:
+ latents[:, :, :batch.image_latent.shape[2]] = batch.image_latent.to(device=device, dtype=latents.dtype)
+
+ kv_cache = self._initialize_kv_cache(latents.shape[0], target_dtype, device, frame_seq_length)
+ crossattn_cache = self._initialize_crossattn_cache(latents.shape[0], target_dtype, device)
+ prompt = batch.prompt_embeds[0]
+ if torch.is_tensor(prompt):
+ prompt = prompt.to(device=device, dtype=target_dtype)
+ context = [sample for sample in prompt]
+ else:
+ context = prompt
+
+ num_blocks = latents.shape[2] // self.num_frame_per_block
+ start = 0
+ first_frame_mask = torch.ones_like(latents)
+ first_frame_mask[:, :, 0] = 0
+ base_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
+ noise_generator = self._make_noise_generator(base_generator, device)
+
+ with tqdm(total=num_blocks * len(timesteps), desc="DreamX AR denoising", leave=False) as progress:
+ for _ in range(num_blocks):
+ current_num_frames = self.num_frame_per_block
+ block_latents = latents[:, :, start:start + current_num_frames]
+ noisy_input = block_latents.clone()
+ mask_block = first_frame_mask[:, :, start:start + current_num_frames]
+ camera_block = self._slice_camera(y_camera, start, current_num_frames)
+
+ for idx, current_timestep in enumerate(timesteps):
+ timestep = torch.full(
+ (latents.shape[0], current_num_frames * frame_seq_length),
+ int(current_timestep.item()),
+ device=device,
+ dtype=torch.long,
+ )
+ if start == 0:
+ timestep[:, :frame_seq_length] = 0
+ with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+ denoised = self.transformer(
+ hidden_states=block_latents.to(target_dtype),
+ encoder_hidden_states=torch.stack(context).to(target_dtype),
+ timestep=timestep,
+ y_camera=camera_block,
+ kv_cache=kv_cache,
+ crossattn_cache=crossattn_cache,
+ current_start=start * frame_seq_length,
+ )
+ denoised = denoised.to(latents.dtype)
+ if idx < len(timesteps) - 1:
+ next_timestep = torch.full((latents.shape[0], current_num_frames),
+ int(timesteps[idx + 1].item()),
+ device=device,
+ dtype=torch.long)
+ noise_kwargs = {"device": device, "dtype": denoised.dtype}
+ if noise_generator is not None:
+ noise_kwargs["generator"] = noise_generator
+ noise = torch.randn(denoised.permute(0, 2, 1, 3, 4).shape, **noise_kwargs)
+ block_btchw = self.scheduler.add_noise(
+ denoised.permute(0, 2, 1, 3, 4).flatten(0, 1),
+ noise.flatten(0, 1),
+ next_timestep.flatten(),
+ ).unflatten(0, (latents.shape[0], current_num_frames))
+ block_latents = block_btchw.permute(0, 2, 1, 3, 4)
+ block_latents = block_latents * mask_block + noisy_input * (1 - mask_block)
+ else:
+ block_latents = denoised * mask_block + noisy_input * (1 - mask_block)
+ progress.update()
+
+ latents[:, :, start:start + current_num_frames] = block_latents
+ self._update_context_cache(block_latents, context, camera_block, kv_cache, crossattn_cache, start,
+ frame_seq_length, target_dtype, autocast_enabled,
+ float(getattr(fastvideo_args.pipeline_config, "context_noise", 0.1)))
+ start += current_num_frames
+
+ batch.latents = latents
+ return batch
+
+ def _make_noise_generator(self, generator: torch.Generator | None, device: torch.device) -> torch.Generator | None:
+ if generator is None:
+ return None
+ if getattr(generator, "device", None) == device:
+ return generator
+ seed = int(generator.initial_seed()) + self._AR_NOISE_SEED_OFFSET
+ return torch.Generator(device=device).manual_seed(seed)
+
+ @staticmethod
+ def _context_noise_timestep(context_noise: float) -> int:
+ if 0.0 < context_noise <= 1.0:
+ return int(context_noise * 1000)
+ return int(context_noise)
+
+ def _slice_camera(self, y_camera: Any, start: int, num_frames: int):
+ if not isinstance(y_camera, dict):
+ return y_camera
+ return {
+ "viewmats": y_camera["viewmats"][:, start:start + num_frames],
+ "K": y_camera["K"][:, start:start + num_frames],
+ }
+
+ def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device,
+ frame_seq_length: int) -> list[dict[str, Any]]:
+ size = self.local_attn_size * frame_seq_length if self.local_attn_size != -1 else 18480
+ heads = self.transformer.num_attention_heads
+ head_dim = self.transformer.attention_head_dim
+ cam_self_attn = next(
+ (getattr(block, "cam_self_attn", None)
+ for block in self.transformer.blocks if getattr(block, "cam_self_attn", None) is not None),
+ None,
+ )
+
+ caches = []
+ for _ in range(self.num_transformer_blocks):
+ cache = {
+ "k": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
+ "v": torch.zeros(batch_size, size, heads, head_dim, dtype=dtype, device=device),
+ "global_end_index": torch.tensor([0], dtype=torch.long, device=device),
+ "local_end_index": torch.tensor([0], dtype=torch.long, device=device),
+ }
+ if cam_self_attn is not None:
+ cam_heads = int(cam_self_attn.num_heads)
+ cam_head_dim = int(cam_self_attn.head_dim)
+ cache.update({
+ "prope_k":
+ torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
+ "prope_v":
+ torch.zeros(batch_size, size, cam_heads, cam_head_dim, dtype=dtype, device=device),
+ "prope_global_end_index":
+ torch.tensor([0], dtype=torch.long, device=device),
+ "prope_local_end_index":
+ torch.tensor([0], dtype=torch.long, device=device),
+ })
+ caches.append(cache)
+ return caches
+
+ def _initialize_crossattn_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device):
+ heads = self.transformer.num_attention_heads
+ head_dim = self.transformer.attention_head_dim
+ return [{
+ "k": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
+ "v": torch.zeros(batch_size, 512, heads, head_dim, dtype=dtype, device=device),
+ "is_init": False,
+ } for _ in range(self.num_transformer_blocks)]
+
+ def _update_context_cache(self, block_latents: torch.Tensor, context: Any, camera_block: Any,
+ kv_cache: list[dict[str, Any]], crossattn_cache: list[dict[str, Any]], start: int,
+ frame_seq_length: int, target_dtype: torch.dtype, autocast_enabled: bool,
+ context_noise: float) -> None:
+ timestep = torch.full(
+ (block_latents.shape[0], block_latents.shape[2] * frame_seq_length),
+ self._context_noise_timestep(context_noise),
+ device=block_latents.device,
+ dtype=torch.long,
+ )
+ with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+ self.transformer(
+ hidden_states=block_latents.to(target_dtype),
+ encoder_hidden_states=torch.stack(context).to(target_dtype),
+ timestep=timestep,
+ y_camera=camera_block,
+ kv_cache=kv_cache,
+ crossattn_cache=crossattn_cache,
+ current_start=start * frame_seq_length,
+ )
diff --git a/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py b/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py
index b330c1125..13b87f154 100644
--- a/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py
+++ b/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py
@@ -121,9 +121,8 @@ def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float
translation_step = np.zeros(3)
rotation_step = np.zeros(3)
for motion_type in motion_types:
- translation_step += _translation_step(
- motion_type, current_pose, float(speed) * _TRANSLATION_BASE_UNIT, duration
- )
+ translation_step += _translation_step(motion_type, current_pose,
+ float(speed) * _TRANSLATION_BASE_UNIT, duration)
rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
segment_positions = []
@@ -136,9 +135,8 @@ def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float
positions.extend(segment_positions)
rotations.extend(segment_rotations)
- rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] + [1.0, 0.0, 0.0, 0.0,
- 0.0, 1.0, 0.0, 0.0,
- 0.0, 0.0, 1.0, 0.0]]
+ rows: list[list[float]] = [[0.0] + _INTRINSIC_ROW + [0.0, 0.0] +
+ [1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
for index, (position, rotation) in enumerate(zip(positions, rotations, strict=False)):
rotation_matrix = _quaternion_to_rotation_matrix(_euler_to_quaternion(rotation))
translation = -rotation_matrix @ position
@@ -152,6 +150,8 @@ def _interpolate_camera_poses(
src_indices: np.ndarray,
tgt_indices: np.ndarray,
) -> list[DreamXCamera]:
+ if len(cameras) <= 1:
+ return [cameras[0]] * len(tgt_indices) if cameras else []
src_rot_mat = np.array([camera.w2c_mat[:3, :3] for camera in cameras])
src_trans_vec = np.array([camera.w2c_mat[:3, 3] for camera in cameras])
@@ -161,9 +161,8 @@ def _interpolate_camera_poses(
flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype)
src_rot_mat = src_rot_mat @ flip_mat
- trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False, fill_value="extrapolate")(
- tgt_indices
- )
+ trans = interp1d(src_indices, src_trans_vec, axis=0, kind="linear", bounds_error=False,
+ fill_value="extrapolate")(tgt_indices)
quats = Rotation.from_matrix(src_rot_mat).as_quat().copy()
for index in range(1, len(quats)):
if np.dot(quats[index], quats[index - 1]) < 0:
diff --git a/fastvideo/pipelines/basic/dreamx_world/config.py b/fastvideo/pipelines/basic/dreamx_world/config.py
index cce9047e8..760f5ed60 100644
--- a/fastvideo/pipelines/basic/dreamx_world/config.py
+++ b/fastvideo/pipelines/basic/dreamx_world/config.py
@@ -2,14 +2,18 @@
"""Compatibility exports for DreamX-World pipeline configs."""
from fastvideo.configs.pipelines.dreamx_world import (
+ DreamXWorld5BARPipelineConfig,
DreamXWorld5BCamPipelineConfig,
+ make_dreamx_world_5b_ar_dit_config,
make_dreamx_world_5b_cam_dit_config,
make_dreamx_world_5b_cam_text_encoder_config,
make_dreamx_world_5b_cam_vae_config,
)
__all__ = [
+ "DreamXWorld5BARPipelineConfig",
"DreamXWorld5BCamPipelineConfig",
+ "make_dreamx_world_5b_ar_dit_config",
"make_dreamx_world_5b_cam_dit_config",
"make_dreamx_world_5b_cam_text_encoder_config",
"make_dreamx_world_5b_cam_vae_config",
diff --git a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py
new file mode 100644
index 000000000..f8af2db0e
--- /dev/null
+++ b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py
@@ -0,0 +1,62 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamX-World-5B autoregressive pipeline entrypoint."""
+
+from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
+from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraConditioningStage
+from fastvideo.pipelines.stages import (
+ ConditioningStage,
+ DecodingStage,
+ InputValidationStage,
+ LatentPreparationStage,
+ TextEncodingStage,
+ TimestepPreparationStage,
+)
+
+logger = init_logger(__name__)
+
+
+class DreamXWorldARPipeline(LoRAPipeline, ComposedPipelineBase):
+ """DreamX-World-5B autoregressive causal camera pipeline."""
+
+ _required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
+ pipeline_config_cls = DreamXWorld5BARPipelineConfig
+ sampling_params_cls = SamplingParam
+
+ def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
+ self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
+ self.modules["scheduler"].set_timesteps(1000)
+
+ def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+ self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+ self.add_stage(stage_name="prompt_encoding_stage",
+ stage=TextEncodingStage(
+ text_encoders=[self.get_module("text_encoder")],
+ tokenizers=[self.get_module("tokenizer")],
+ ))
+ self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+ self.add_stage(stage_name="timestep_preparation_stage",
+ stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")))
+ self.add_stage(stage_name="latent_preparation_stage",
+ stage=LatentPreparationStage(
+ scheduler=self.get_module("scheduler"),
+ transformer=self.get_module("transformer", None),
+ ))
+ self.add_stage(stage_name="dreamx_camera_conditioning_stage", stage=DreamXWorldCameraConditioningStage())
+ self.add_stage(stage_name="denoising_stage",
+ stage=DreamXWorldARCausalDenoisingStage(
+ transformer=self.get_module("transformer"),
+ scheduler=self.get_module("scheduler"),
+ vae=self.get_module("vae"),
+ pipeline=self,
+ ))
+ self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"), pipeline=self))
+ logger.info("DreamXWorldARPipeline initialized with autoregressive causal denoising")
+
+
+EntryClass = DreamXWorldARPipeline
diff --git a/fastvideo/pipelines/basic/dreamx_world/presets.py b/fastvideo/pipelines/basic/dreamx_world/presets.py
index 05fb28aa2..2d9d6e407 100644
--- a/fastvideo/pipelines/basic/dreamx_world/presets.py
+++ b/fastvideo/pipelines/basic/dreamx_world/presets.py
@@ -3,13 +3,11 @@
from fastvideo.api.presets import InferencePreset, PresetStageSpec
-_NEGATIVE_PROMPT_CN = (
- "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，"
- "静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，"
- "多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，"
- "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，"
- "背景人很多，倒着走"
-)
+_NEGATIVE_PROMPT_CN = ("色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，"
+ "静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，"
+ "多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，"
+ "形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，"
+ "背景人很多，倒着走")
_DENOISE_STAGE = PresetStageSpec(
name="denoise",
@@ -39,4 +37,22 @@ DREAMX_WORLD_5B_CAM = InferencePreset(
},
)
-ALL_PRESETS = (DREAMX_WORLD_5B_CAM, )
+DREAMX_WORLD_5B_AR = InferencePreset(
+ name="dreamx_world_5b_ar",
+ version=1,
+ model_family="dreamx_world",
+ description="DreamX-World 5B autoregressive camera-control generation",
+ workload_type="i2v",
+ stage_schemas=(_DENOISE_STAGE, ),
+ defaults={
+ "height": 704,
+ "width": 1280,
+ "num_frames": 1005,
+ "fps": 16,
+ "guidance_scale": 1.0,
+ "num_inference_steps": 4,
+ "negative_prompt": _NEGATIVE_PROMPT_CN,
+ },
+)
+
+ALL_PRESETS = (DREAMX_WORLD_5B_CAM, DREAMX_WORLD_5B_AR)
diff --git a/fastvideo/pipelines/basic/dreamx_world/stages.py b/fastvideo/pipelines/basic/dreamx_world/stages.py
index feb0d0471..00a682ee3 100644
--- a/fastvideo/pipelines/basic/dreamx_world/stages.py
+++ b/fastvideo/pipelines/basic/dreamx_world/stages.py
@@ -11,9 +11,7 @@ from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import VerificationResult
from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (
- build_dreamx_camera_condition,
-)
-
+ build_dreamx_camera_condition, )
DREAMX_Y_CAMERA_KEY = "dreamx_y_camera"
@@ -31,8 +29,7 @@ class DreamXWorldCameraConditioningStage(PipelineStage):
return batch
action_seq = batch.extra.get("dreamx_action_seq", batch.action_list)
- action_speed_list = batch.extra.get(
- "dreamx_action_speed_list", batch.action_speed_list)
+ action_speed_list = batch.extra.get("dreamx_action_speed_list", batch.action_speed_list)
if action_seq is None:
action_seq = ["w"]
if action_speed_list is None:
@@ -40,6 +37,10 @@ class DreamXWorldCameraConditioningStage(PipelineStage):
if isinstance(action_seq, str):
action_seq = [action_seq]
+ if isinstance(action_speed_list, int | float):
+ action_speed_list = [action_speed_list]
+ if len(action_speed_list) == 1 and len(action_seq) > 1:
+ action_speed_list = list(action_speed_list) * len(action_seq)
action_speed_list = [float(speed) for speed in action_speed_list]
height = int(batch.height) if batch.height is not None else 704
@@ -57,10 +58,7 @@ class DreamXWorldCameraConditioningStage(PipelineStage):
dtype=dtype,
device=device,
)
- batch.extra[DREAMX_Y_CAMERA_KEY] = {
- key: value.unsqueeze(0)
- for key, value in y_camera.items()
- }
+ batch.extra[DREAMX_Y_CAMERA_KEY] = {key: value.unsqueeze(0) for key, value in y_camera.items()}
return batch
def verify_output(
diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py
index 806c22f6f..505f7e7f6 100644
--- a/fastvideo/pipelines/stages/denoising.py
+++ b/fastvideo/pipelines/stages/denoising.py
@@ -194,7 +194,7 @@ class DenoisingStage(PipelineStage):
dreamx_y_camera = batch.extra.get("dreamx_y_camera", batch.extra.get("y_camera"))
if isinstance(dreamx_y_camera, dict):
dreamx_y_camera = {
- key: value.to(device=local_device) if torch.is_tensor(value) else value
+ key: value.to(device=local_device, dtype=target_dtype) if torch.is_tensor(value) else value
for key, value in dreamx_y_camera.items()
}
dreamx_camera_kwargs = self.prepare_extra_func_kwargs(
@@ -255,8 +255,11 @@ class DenoisingStage(PipelineStage):
# the image latent instead of appending along the channel dim
assert batch.image_latent is None, "TI2V task should not have image latents"
assert self.vae is not None, "VAE is not provided for TI2V task"
+ vae_device = next(self.vae.parameters()).device
self.vae = self.vae.to(local_device)
z = self.vae.encode(batch.pil_image).mean.float()
+ if getattr(fastvideo_args, "vae_cpu_offload", False):
+ self.vae = self.vae.to(vae_device)
if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
if isinstance(self.vae.shift_factor, torch.Tensor):
z -= self.vae.shift_factor.to(z.device, z.dtype)
diff --git a/fastvideo/registry.py b/fastvideo/registry.py
index c6f1467ca..c1ece178d 100644
--- a/fastvideo/registry.py
+++ b/fastvideo/registry.py
@@ -20,7 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
Cosmos25Config,
Cosmos25_14BConfig,
)
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
+from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
from fastvideo.configs.pipelines.gen3c import Gen3CConfig
@@ -786,6 +786,21 @@ def _register_configs() -> None:
model_family="dreamx_world",
default_preset="dreamx_world_5b_cam",
)
+ register_configs(
+ sampling_param_cls=None,
+ pipeline_config_cls=DreamXWorld5BARPipelineConfig,
+ workload_types=(WorkloadType.I2V, ),
+ hf_model_paths=[
+ "GD-ML/DreamX-World-5B",
+ ],
+ model_detectors=[
+ lambda path:
+ ("dreamx-world-5b" in path.lower() and "cam" not in path.lower()) or "dreamxworldarpipeline" in path.lower(
+ )
+ ],
+ model_family="dreamx_world",
+ default_preset="dreamx_world_5b_ar",
+ )
register_configs(
sampling_param_cls=None,
pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
diff --git a/fastvideo/tests/ssim/inference_similarity_utils.py b/fastvideo/tests/ssim/inference_similarity_utils.py
index e2946d28e..27e3a76b6 100644
--- a/fastvideo/tests/ssim/inference_similarity_utils.py
+++ b/fastvideo/tests/ssim/inference_similarity_utils.py
@@ -3,9 +3,9 @@
from __future__ import annotations
import os
+from collections.abc import Iterator
from contextlib import contextmanager
from logging import Logger
-from typing import Iterator
from fastvideo import VideoGenerator
from fastvideo.tests.ssim.reference_utils import (
@@ -67,6 +67,12 @@ def _find_reference_video(reference_folder: str, prompt: str) -> str:
raise FileNotFoundError("Reference video missing")
+def _remove_stale_generated_video(output_dir: str, output_video_name: str) -> None:
+ stale_path = os.path.join(output_dir, output_video_name)
+ if os.path.exists(stale_path):
+ os.remove(stale_path)
+
+
def _assert_similarity(
*,
logger: Logger,
@@ -214,6 +220,7 @@ def run_text_to_video_similarity_test(
)
output_video_name = f"{prompt[:100].strip()}.mp4"
os.makedirs(output_dir, exist_ok=True)
+ _remove_stale_generated_video(output_dir, output_video_name)
params_map = select_ssim_params(
default_params_map,
@@ -289,6 +296,7 @@ def run_image_to_video_similarity_test(
)
output_video_name = f"{prompt[:100].strip()}.mp4"
os.makedirs(output_dir, exist_ok=True)
+ _remove_stale_generated_video(output_dir, output_video_name)
params_map = select_ssim_params(
default_params_map,
diff --git a/fastvideo/tests/ssim/test_dreamx_world_similarity.py b/fastvideo/tests/ssim/test_dreamx_world_similarity.py
index 9f6861921..c8394bacf 100644
--- a/fastvideo/tests/ssim/test_dreamx_world_similarity.py
+++ b/fastvideo/tests/ssim/test_dreamx_world_similarity.py
@@ -11,6 +11,7 @@ from fastvideo.logger import init_logger
from fastvideo.tests.ssim.inference_similarity_utils import (
resolve_inference_device_reference_folder,
run_image_to_video_similarity_test,
+ run_text_to_video_similarity_test,
)
logger = init_logger(__name__)
@@ -24,6 +25,18 @@ _MODEL_PATH = os.getenv(
"DREAMX_WORLD_SSIM_MODEL_PATH",
str(_LOCAL_CONVERTED_MODEL),
)
+_LOCAL_AR_CANDIDATES = (
+ Path("/tmp/converted_dreamx_world_ar"),
+ Path("/root/data/dreamx_world_ar_converted"),
+)
+_DEFAULT_AR_MODEL_PATH = next(
+ (str(path) for path in _LOCAL_AR_CANDIDATES if path.exists()),
+ str(_LOCAL_AR_CANDIDATES[0]),
+)
+_AR_MODEL_PATH = os.getenv(
+ "DREAMX_WORLD_AR_SSIM_MODEL_PATH",
+ _DEFAULT_AR_MODEL_PATH,
+)
DREAMX_WORLD_PARAMS = {
"num_gpus": 1,
@@ -46,14 +59,42 @@ DREAMX_WORLD_FULL_QUALITY_PARAMS = {
"guidance_scale": 5.0,
}
+
+DREAMX_WORLD_AR_PARAMS = {
+ "num_gpus": 1,
+ "model_path": _AR_MODEL_PATH,
+ "height": 192,
+ "width": 192,
+ "num_frames": 81,
+ "num_inference_steps": 4,
+ "guidance_scale": 1.0,
+ "seed": 2048,
+ "fps": 16,
+}
+
+DREAMX_WORLD_AR_FULL_QUALITY_PARAMS = {
+ **DREAMX_WORLD_AR_PARAMS,
+ "height": 704,
+ "width": 1280,
+ "num_frames": 1005,
+}
+
DREAMX_WORLD_MODEL_TO_PARAMS = {
"DreamX-World-5B-Cam": DREAMX_WORLD_PARAMS,
}
+DREAMX_WORLD_AR_MODEL_TO_PARAMS = {
+ "DreamX-World-5B": DREAMX_WORLD_AR_PARAMS,
+}
+
FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS = {
"DreamX-World-5B-Cam": DREAMX_WORLD_FULL_QUALITY_PARAMS,
}
+FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS = {
+ "DreamX-World-5B": DREAMX_WORLD_AR_FULL_QUALITY_PARAMS,
+}
+
DREAMX_WORLD_TEST_CASES = [
(
"A cinematic first-person drive through a futuristic coastal city at sunrise, "
@@ -64,6 +105,16 @@ DREAMX_WORLD_TEST_CASES = [
]
+DREAMX_WORLD_AR_TEST_CASES = [
+ (
+ "A long autonomous drive through a futuristic coastal city at sunrise, "
+ "smooth forward camera motion, reflective glass towers, clean streets.",
+ ("w", "d", "w", "a"),
+ (2.0, 1.0, 2.0, 1.0),
+ ),
+]
+
+
def _write_deterministic_reference_image(path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
image = Image.new("RGB", (96, 96), color=(42, 76, 112))
@@ -120,3 +171,44 @@ def test_dreamx_world_inference_similarity(
"action_speed_list": list(action_speed_list),
},
)
+
+@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_AR_TEST_CASES)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_AR_MODEL_TO_PARAMS.keys()))
+def test_dreamx_world_ar_inference_similarity(
+ prompt: str,
+ action_list: tuple[str, ...],
+ action_speed_list: tuple[float, ...],
+ attention_backend_name: str,
+ model_id: str,
+) -> None:
+ model_path = Path(str(DREAMX_WORLD_AR_MODEL_TO_PARAMS[model_id]["model_path"]))
+ if not model_path.exists():
+ pytest.skip(
+ f"DreamX-World AR converted model path is missing: {model_path}. "
+ "Set DREAMX_WORLD_AR_SSIM_MODEL_PATH to a FastVideo-loadable converted root."
+ )
+
+ run_text_to_video_similarity_test(
+ logger=logger,
+ script_dir=os.path.dirname(os.path.abspath(__file__)),
+ device_reference_folder=device_reference_folder,
+ prompt=prompt,
+ attention_backend_name=attention_backend_name,
+ model_id=model_id,
+ default_params_map=DREAMX_WORLD_AR_MODEL_TO_PARAMS,
+ full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_AR_MODEL_TO_PARAMS,
+ min_acceptable_ssim=0.98,
+ init_kwargs_override={
+ "use_fsdp_inference": False,
+ "dit_cpu_offload": False,
+ "vae_cpu_offload": True,
+ "text_encoder_cpu_offload": True,
+ "pin_cpu_memory": False,
+ "override_pipeline_cls_name": "DreamXWorldARPipeline",
+ },
+ generation_kwargs_override={
+ "action_list": list(action_list),
+ "action_speed_list": list(action_speed_list),
+ },
+ )
diff --git a/scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py b/scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py
new file mode 100755
index 000000000..a73544181
--- /dev/null
+++ b/scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py
@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+"""Convert DreamX-World-5B autoregressive weights to FastVideo layout.
+
+The HF repository stores one raw official ``model.safetensors`` whose keys match
+FastVideo's native ``DreamXWorldARTransformer3DModel``. The converter writes a
+Diffusers-like root with ``transformer/config.json`` and reusable Wan2.2
+components. Use ``--symlink-transformer`` locally to avoid duplicating the 21GB
+AR tensor file.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+
+TRANSFORMER_CONFIG: dict[str, object] = {
+ "_class_name": "DreamXWorldARTransformer3DModel",
+ "model_type": "ti2v",
+ "patch_size": [1, 2, 2],
+ "text_len": 512,
+ "num_attention_heads": 24,
+ "attention_head_dim": 128,
+ "in_channels": 48,
+ "out_channels": 48,
+ "text_dim": 4096,
+ "freq_dim": 256,
+ "ffn_dim": 14336,
+ "num_layers": 30,
+ "local_attn_size": 12,
+ "sink_size": 3,
+ "cross_attn_norm": True,
+ "qk_norm": True,
+ "eps": 1e-6,
+ "add_control_adapter": True,
+ "cam_method": "prope",
+ "attn_compress": 4,
+ "cam_self_attn_layers": list(range(30)),
+ "num_frames_per_block": 3,
+}
+
+MODEL_INDEX: dict[str, object] = {
+ "_class_name": "DreamXWorldARPipeline",
+ "_diffusers_version": "0.31.0",
+ "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+ "text_encoder": ["transformers", "UMT5EncoderModel"],
+ "tokenizer": ["transformers", "AutoTokenizer"],
+ "transformer": ["diffusers", "DreamXWorldARTransformer3DModel"],
+ "vae": ["diffusers", "AutoencoderKLWan"],
+}
+
+REUSED_COMPONENTS = ("scheduler", "text_encoder", "tokenizer", "vae")
+
+
+def _source_safetensors(source: Path) -> Path:
+ if source.is_file():
+ return source
+ path = source / "model.safetensors"
+ if not path.exists():
+ raise FileNotFoundError(f"Missing AR model.safetensors under {source}")
+ return path
+
+
+def convert_transformer(source: Path, output: Path, symlink_transformer: bool) -> None:
+ src = _source_safetensors(source)
+ transformer_dir = output / "transformer"
+ transformer_dir.mkdir(parents=True, exist_ok=True)
+ (transformer_dir / "config.json").write_text(json.dumps(TRANSFORMER_CONFIG, indent=2) + "\n")
+ dst = transformer_dir / "model.safetensors"
+ if dst.exists() or dst.is_symlink():
+ return
+ if symlink_transformer:
+ dst.symlink_to(src.resolve())
+ else:
+ shutil.copy2(src, dst)
+
+
+def _copy_or_link_component(component: str, component_source: Path, output: Path, symlink: bool) -> None:
+ src = component_source / component
+ dst = output / component
+ if not src.exists():
+ raise FileNotFoundError(f"Missing reused component source: {src}")
+ if dst.exists() or dst.is_symlink():
+ return
+ if symlink:
+ dst.symlink_to(src.resolve(), target_is_directory=src.is_dir())
+ elif src.is_dir():
+ shutil.copytree(src, dst)
+ else:
+ shutil.copy2(src, dst)
+
+
+def write_model_index(output: Path, component_source: Path | None, symlink_components: bool) -> None:
+ if component_source is not None:
+ for component in REUSED_COMPONENTS:
+ _copy_or_link_component(component, component_source, output, symlink_components)
+ missing = [component for component in REUSED_COMPONENTS if not (output / component).exists()]
+ if missing:
+ print("skipping model_index.json; missing reused components: " + ", ".join(missing))
+ return
+ (output / "model_index.json").write_text(json.dumps(MODEL_INDEX, indent=2) + "\n")
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--source", type=Path, required=True)
+ parser.add_argument("--output", type=Path, required=True)
+ parser.add_argument("--component-source", type=Path)
+ parser.add_argument("--symlink-components", action="store_true")
+ parser.add_argument("--symlink-transformer", action="store_true")
+ args = parser.parse_args()
+
+ convert_transformer(args.source, args.output, args.symlink_transformer)
+ write_model_index(args.output, args.component_source, args.symlink_components)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py
index 6297863fd..1abf4ee51 100644
--- a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py
+++ b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py
@@ -146,7 +146,7 @@ def _copy_or_link_component(component: str, component_source: Path, output: Path
dst = output / component
if not src.exists():
raise FileNotFoundError(f"Missing reused component source: {src}")
- if dst.exists():
+ if dst.exists() or dst.is_symlink():
print(f"keeping existing {dst}")
return
if symlink:
diff --git a/tests/local_tests/dreamx_world/PORT_STATUS.md b/tests/local_tests/dreamx_world/PORT_STATUS.md
index e4e7bb20b..012d9c9bf 100644
--- a/tests/local_tests/dreamx_world/PORT_STATUS.md
+++ b/tests/local_tests/dreamx_world/PORT_STATUS.md
@@ -3,7 +3,7 @@
## Summary
- model_family: `dreamx_world`
-- workload_types: `I2V camera-control compatibility shim`
+- workload_types: `I2V camera-control compatibility shim`; `I2V autoregressive camera-control forcing`
- official_ref: `https://github.com/AMAP-ML/DreamX-World`
- official_ref_dir: `DreamX-World/`
- hf_weights_path: `GD-ML/DreamX-World-5B-Cam`
@@ -16,7 +16,7 @@
- phase: `phase_11_post_parity_handoff`
- status: `complete`
- owner: `orchestrator`
-- last_updated: `2026-07-01`
+- last_updated: `2026-07-02`
## Component Matrix
@@ -28,6 +28,8 @@
| scheduler | generic | reuse_proven | Diffusers `FlowMatchEulerDiscreteScheduler` | `DreamX-World/inference_dreamx5b.py::setup_models`, default `sampler_name=Flow` | `fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py` | pass | not_required | non_skip_pass | Q003 |
| camera_conditioning | generic | port_pending | `DreamX-World/utils/inference_utils.py`, `DreamX-World/models/prope_utils.py`, `DreamX-World/wan/modules/camera_prope.py` | `DreamX-World/inference_dreamx5b.py::get_camera_sequence`, `pipeline(... control_camera_video=...)` | `fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py` | pass | not_required | non_skip_pass | none |
| pipeline | pipeline | port_complete | `DreamX-World/pipeline/pipeline_dreamxworld.py` | `DreamX-World/inference_dreamx5b.py::process_inference_from_json` | `fastvideo/pipelines/basic/dreamx_world/` plus config/preset/registry | pipeline_load_generate_smoke_pass | model_index_and_config_consistency_smoke_pass | pipeline_api_vs_worker_forward_parity_pass | none |
+| ar_transformer | dit | port_complete | `DreamX-World/wan/modules/causal_camera_model_2_2_prope_infinity.py::CausalWanModel` | `DreamX-World/inference_ar_forcing.py::load_pipeline` | `fastvideo/models/dits/dreamx_world_ar.py`; `fastvideo/configs/models/dits/dreamx_world.py::DreamXWorldARConfig` | tiny_official_forward_parity_pass | identity_conversion_pass | real_5b_strict_load_pass | none |
+| ar_pipeline | pipeline | port_complete | `DreamX-World/pipeline/pipeline_causal_camera.py` | `DreamX-World/inference_ar_forcing.py::main` | `fastvideo/pipelines/basic/dreamx_world/dreamx_world_ar_pipeline.py`; `fastvideo/pipelines/basic/dreamx_world/ar_denoising.py`; registry/preset/config | config_registry_pass | symlink_model_index_pass | short_full_generation_pass | none |
## Conversion State
@@ -50,6 +52,12 @@
| scheduler | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py -v -s` | non_skip_pass | 2026-06-30: FastVideo FlowMatch scheduler matches official Diffusers timesteps and step output for DreamX default Flow sampler; `DreamXWorldPipeline` initializes FlowMatch with official `shift=3.0`. |
| camera_conditioning | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py -v -s` | non_skip_pass | 2026-06-29: 3 parameterized cases passed against official reference on CPU. |
| pipeline_config | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -v -s` | pipeline_entry_preset_scheduler_modelinfo_and_camera_stage_smoke_pass | DreamX 5B-Cam PipelineConfig wires DiT/VAE/UMT5/Flow/TI2V settings and official `shift=3.0`; default preset is registered for `GD-ML/DreamX-World-5B-Cam`; local converted-style `model_index.json` resolves to `DreamXWorldPipeline`; the pipeline initializes FlowMatch, camera conditioning writes `batch.extra["dreamx_y_camera"]`, and generic denoising can pass it as `y_camera`. |
+| ar_transformer | `PYTHONPATH=/workspace/FastVideo python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py -q -rs` | 6_passed_0_skipped | 2026-07-02: AR converter writes symlinked transformer layout/model_index; tiny official `CausalWanModel` vs FastVideo `DreamXWorldARTransformer3DModel` forward parity passes; real 5B `model.safetensors` strict-loads with zero missing/unexpected keys from `/tmp/converted_dreamx_world_ar`. |
+| ar_pipeline_config | `PYTHONPATH=/workspace/FastVideo python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -q -rs` | 10_passed_0_skipped | 2026-07-02: `DreamXWorld5BARPipelineConfig`, `DreamXWorldARPipeline`, registry config selection for `GD-ML/DreamX-World-5B`, and `dreamx_world_5b_ar` preset pass. |
+| ar_full_generation | `PYTHONPATH=/workspace/FastVideo python /tmp/run_dreamx_ar_video_smoke.py` | generated_video_pass | 2026-07-02: A40 short full-generation smoke passed from `/tmp/converted_dreamx_world_ar` with 64x64, 9 frames, 4 denoise steps, `output_type=pil`, `save_video=True`; MP4 saved at `outputs_video/dreamx_world_ar_video_smoke/a quiet road through a futuristic coastal city at sunrise.mp4` and decoded as 9 frames of `(64, 64, 3)` uint8. |
+| ar_long_horizon | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA python /tmp/run_dreamx_ar_long_horizon.py` | generated_video_pass | 2026-07-02: A40 long-horizon AR generation passed from `/tmp/converted_dreamx_world_ar` with 64x64, 1005 frames, 4 denoise steps, seed 4096; MP4 saved at `outputs_video/dreamx_world_ar_long_horizon/A long autonomous drive through a futuristic coastal city at sunrise, smooth forward camera motion,.mp4` and decoded as 1005 frames of `(64, 64, 3)`; end-to-end generation latency was 231.68s after load. |
+| ar_ssim_default | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download` | 1_passed_0_skipped | 2026-07-02: A40 default AR SSIM reference seeded locally at `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/`; helper removes stale generated base MP4 before generation; default params are 192x192, 81 frames, 4 steps, seed 2048, min SSIM 0.98. |
+| ar_ssim_modal_l40s | `modal run /tmp/modal_dreamx_ar_ssim_git.py` | 1_passed_0_skipped | 2026-07-02: Modal L40S run checked out `post-fix dreamx-world-5b-cam branch commit`, downloaded default `L40S_reference_videos` via `reference_videos_cli.py download`, used cached converted AR weights under `/root/data/dreamx_world_ar_converted`, seeded the missing AR L40S reference from generated output, reran a fresh generated-vs-reference compare successfully (`mean_ssim=1.0`), and exported the reference to Modal volume `hf-model-weights:dreamx_ar_ssim_l40s`. Downloaded local reference decodes as 81 frames of `(192, 192, 3)`. A40-vs-L40S reference spot check after deterministic AR noise fix: mean SSIM 0.7770, min 0.6139, max 0.9944. |
| pipeline_smoke | `python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py -q -rs` | 4_passed_0_skipped | 2026-06-30: combined smoke/parity passed. 2026-07-01: smoke alone passed with real `image_path` TI2V coverage (`3 passed`), validating image load, TI2V preprocessing, VAE first-frame encode under CPU offload, camera conditioning, and 1-step latent generation from `converted_weights/dreamx_world`. Tests force `FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA` to avoid the local FlashAttention-4 cute ABI mismatch. |
| basic_example | `PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA DREAMX_WORLD_MODEL_DIR=converted_weights/dreamx_world DREAMX_WORLD_IMAGE_PATH= DREAMX_WORLD_HEIGHT=64 DREAMX_WORLD_WIDTH=64 DREAMX_WORLD_NUM_FRAMES=9 DREAMX_WORLD_STEPS=1 DREAMX_WORLD_GUIDANCE=1.0 DREAMX_WORLD_OUTPUT_PATH=outputs_video/dreamx_world_example_smoke python examples/inference/basic/basic_dreamx_world.py` | generated_video_pass | 2026-06-30: example saved an MP4 under `outputs_video/dreamx_world_example_smoke`; imageio/ffmpeg decoded frame 0 as `(64, 64, 3)` uint8, fps 16, duration 0.56s. |
@@ -62,6 +70,8 @@
| Q003 | Which sampler is in first-PR scope: official default `Flow` only, or also `Flow_Unipc` and `Flow_DPM++`? | orchestrator | Phase 3 | resolved | First PR should support official default `Flow` only. FastVideo FlowMatch scheduler parity is non-skip PASS; optional `Flow_Unipc` and `Flow_DPM++` are out of first-PR scope. |
| Q004 | Which HF token env var should be used if rate limits or gated Wan2.2 base weights require auth? | user | Phase 5 | resolved | No auth was required for the completed local downloads; keep using env var names only if future gated repos require auth. |
| Q005 | Should native FastVideo production code depend on DreamX reference-only packages such as `xfuser` or OpenCV? | user | Phase 3 | resolved | No. These packages may be used only for official reference/local parity setup; native FastVideo integration must remove that runtime requirement. |
+| Q006 | Should AR handoff require a full generated long-horizon video in this no-HF-token/no-GPU-budget pass? | user/runtime | pipeline | resolved | A40 long-horizon generation passes with 1005 frames at 64x64/4 steps. AR default SSIM passes locally on A40 and on Modal L40S with a 192x192/81-frame reference. HF upload/publication remains a separate operation if the reference dataset should be updated upstream. |
+| Q007 | Can A40 references stand in for L40S CI references? | quality | release | resolved | No. After deterministic AR noise fix, A40-vs-L40S reference mean SSIM is 0.7770, below the 0.98 same-device threshold. Publish the L40S-specific reference for CI. |
## Issues And Blockers
@@ -74,6 +84,8 @@
| I005 | parity | transformer | medium | Transformer full forward parity initially failed in bf16 official harness. | Official CUDA bf16 LayerNorm path was unstable; fp32 small-input harness avoids that dtype issue and compares against FastVideo with single-process SP identity patches. | component:transformer | resolved | Official-vs-FastVideo forward parity now passes on CUDA with `diff_max=0.072533`, `diff_mean=0.008014`. |
| I006 | parity | vae/text_encoder | medium | VAE/text parity initially remained skipped after weights were staged. | Text official import needed a reference-only `xfuser` stub; VAE comparison initially used raw official normalized latents against FastVideo raw mu. | component:vae,component:text_encoder | resolved | Text hidden-state parity passes. VAE encode parity passes after raw key mapping and applying the official latent normalization to FastVideo output. |
| I007 | quality | pipeline_ti2v | medium | Real image-path TI2V smoke initially failed when `vae_cpu_offload=True` because DenoisingStage encoded the first frame while the VAE weights remained on CPU. | DreamX SSIM first run failed with `RuntimeError: Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same` at `fastvideo/pipelines/stages/denoising.py` VAE encode. | pipeline | resolved | DenoisingStage now moves the VAE to `local_device` before TI2V first-frame encode; image-path pipeline smoke and DreamX SSIM both pass. |
+| I008 | quality | ssim_helper | high | SSIM helper could compare against a stale generated base MP4 when a rerun saved the new video as `_1.mp4`. | Existing generated outputs made AR reference seeding appear to pass before a fresh generated-vs-reference compare. | quality | resolved | `run_text_to_video_similarity_test` and `run_image_to_video_similarity_test` now remove the stale generated base MP4 before generation. A40 and Modal L40S AR SSIM were rerun after the fix. |
+| I009 | quality | ar_denoising | high | AR denoising added CUDA noise without using the request seed when the original generator was CPU-backed. | Fresh reruns against old AR references produced mean SSIM near 0.05. | pipeline | resolved | `DreamXWorldARCausalDenoisingStage` now derives a device-local generator from the request seed for AR noise. Fresh same-device A40 and L40S reruns pass with mean SSIM 1.0 after reseeding references. |
## Escape Hatches
@@ -115,7 +127,7 @@
- command: `PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B-Cam python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs`
- result: `1 passed, 0 skipped` on 2026-07-01
- reference: Local A40/TORCH_SDPA reference seeded from the generated candidate under `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B-Cam/TORCH_SDPA/`. The test uses a deterministic generated input image, 64x64 request dimensions, 9 frames, 1 denoise step, seed 1024, and min SSIM 0.98. Full-quality params are present for 480x832/161 frames/30 steps.
-- note: HF upload/Modal L40S seeding remains a release operation requiring `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` with write access; no token values were used or recorded.
+- note: Modal L40S seeding passed using the configured Modal profile and unauthenticated HF public downloads. HF upload/publication still requires `HF_API_KEY`, `HUGGINGFACE_HUB_TOKEN`, or `HF_TOKEN` with write access; no token values were used or recorded.
## Final Handoff
@@ -149,6 +161,16 @@ final_handoff:
parity_test: tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py
parity_status: non_skip_pass
concerns_or_unknowns: none
+ - name: ar_transformer
+ reuse_or_port: ported
+ parity_test: tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py
+ parity_status: non_skip_pass
+ concerns_or_unknowns: none
+ - name: ar_pipeline
+ reuse_or_port: ported
+ parity_test: tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py plus AR smoke/SSIM commands listed above
+ parity_status: non_skip_pass
+ concerns_or_unknowns: none
pipeline_smoke: pass
pipeline_parity: pass
example_status: pass
@@ -160,3 +182,8 @@ final_handoff:
blockers: none
escape_hatch: none
```
+
+| 2026-07-02 | Add DreamX-World-5B autoregressive support. | Official AR repo is raw single-safetensors layout and needs a dedicated causal/KV stage. | Added native AR DiT/config, identity converter, AR pipeline config/preset/registry, and targeted non-skip tests. Raw AR weights are staged at `/tmp/dreamx_world_ar_weights`; converted symlink layout at `/tmp/converted_dreamx_world_ar`. |
+| 2026-07-02 | Validate DreamX-World-5B AR short full generation on A40. | Targeted parity/config tests prove components, but end-to-end runtime can still fail at scheduler, RoPE cache, device, or decode boundaries. | `DreamXWorldARPipeline` generated and saved a 64x64/9-frame/4-step MP4 from `/tmp/converted_dreamx_world_ar`; the saved MP4 decodes to 9 frames. |
+| 2026-07-02 | Validate DreamX-World-5B AR long-horizon and default SSIM on A40. | AR needs coverage beyond the 9-frame smoke to exercise longer KV/cache progression and a quality regression path. | 1005-frame 64x64/4-step generation passes and decodes; default AR SSIM test uses 192x192/81 frames because MS-SSIM requires short side >160. |
+| 2026-07-02 | Validate DreamX-World-5B AR default SSIM on Modal L40S. | CI references are device-specific; A40 alone is not enough for L40S reference coverage. | Modal checked out `post-fix dreamx-world-5b-cam branch commit`, downloaded existing default L40S references, seeded the missing AR reference, reran SSIM successfully (`mean_ssim=1.0`), and exported the L40S reference back to the workspace. |
diff --git a/tests/local_tests/dreamx_world/README.md b/tests/local_tests/dreamx_world/README.md
index c19051482..3febd5215 100644
--- a/tests/local_tests/dreamx_world/README.md
+++ b/tests/local_tests/dreamx_world/README.md
@@ -13,8 +13,8 @@ Port progress, open questions, issues, and handoff notes live in
| Field | Value |
|---|---|
| Model family | `dreamx_world` |
-| First-PR scope | `DreamX-World-5B-Cam` only |
-| Out-of-scope variants | `DreamX-World-5B` autoregressive long-horizon forcing |
+| First-PR scope | `DreamX-World-5B-Cam`; follow-up scope now includes `DreamX-World-5B` autoregressive forcing |
+| Out-of-scope variants | none for the DreamX-World 5B/Cam paths currently ported |
| Workload types | I2V camera-control compatibility shim: image + prompt + action sequence to video |
| Official reference | `https://github.com/AMAP-ML/DreamX-World` |
| Local reference dir | `DreamX-World/` |
@@ -95,6 +95,36 @@ python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \
"official_weights/dreamx_world"
```
+
+
+### DreamX-World-5B Autoregressive Setup
+
+The AR repository `GD-ML/DreamX-World-5B` is also raw official layout: no
+`model_index.json`, root `config.json`, and a single `model.safetensors`. The
+current environment has the raw AR checkpoint staged outside the workspace at
+`/tmp/dreamx_world_ar_weights` to avoid workspace quota pressure. The converted
+FastVideo layout is staged at `/tmp/converted_dreamx_world_ar` with the 21GB
+transformer safetensors symlinked instead of copied.
+
+```bash
+python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \
+ "GD-ML/DreamX-World-5B" \
+ "/tmp/dreamx_world_ar_weights"
+
+python scripts/checkpoint_conversion/dreamx_world_ar_to_diffusers.py \
+ --source /tmp/dreamx_world_ar_weights \
+ --output /tmp/converted_dreamx_world_ar \
+ --component-source official_weights/Wan2.2-TI2V-5B-Diffusers \
+ --symlink-components \
+ --symlink-transformer
+```
+
+AR production code uses `DreamXWorldARTransformer3DModel`,
+`DreamXWorld5BARPipelineConfig`, `DreamXWorldARPipeline`, and
+`DreamXWorldARCausalDenoisingStage`. The AR DiT is a native FastVideo port of
+the official Apache-2.0 `CausalWanModel`; it has no production DreamX, Diffusers
+model-class, Transformers model-class, `xfuser`, or OpenCV import.
+
## Prototype And Conversion Artifacts
State-dict key/shape dumps are generated after FastVideo native prototypes exist
@@ -138,6 +168,8 @@ Planned local tests for this family:
| scheduler | Diffusers `FlowMatchEulerDiscreteScheduler`; selected by default `sampler_name=Flow` | `tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py` | First PR can support official default `Flow`; optional `Flow_Unipc` and `Flow_DPM++` are out of first-PR scope. | non_skip_pass |
| camera_conditioning | `DreamX-World/utils/inference_utils.py`, `models/prope_utils.py`, `wan/modules/camera_prope.py` | `tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py` | Action sequence to PRoPE/control input must match official tensor shapes and values. | non_skip_pass |
| pipeline | `DreamX-World/pipeline/pipeline_dreamxworld.py`; call path in `DreamX-World/inference_dreamx5b.py` | `tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py`; `tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py`; `tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py` | DreamX PipelineConfig wires first-scope DiT/VAE/UMT5/Flow/TI2V settings, official `shift=3.0`, default preset values, FlowMatch scheduler initialization, and local `model_index.json` resolution; independent pipeline smoke covers real CUDA local load + latent generation, and parity compares public API output to worker-side explicit ForwardBatch execution. | pipeline_smoke_and_parity_pass |
+| ar_transformer | `DreamX-World/wan/modules/causal_camera_model_2_2_prope_infinity.py`; instantiated by `DreamX-World/inference_ar_forcing.py` as `CausalWanModel` with `local_attn_size=12`, `sink_size=3`, `attn_compress=4` | `tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py` | Native `DreamXWorldARTransformer3DModel` keeps official identity key layout; tiny official-vs-FastVideo forward parity passes; real 5B AR safetensors strict-load passes from `/tmp/converted_dreamx_world_ar`. | tiny_forward_parity_and_real_strict_load_pass |
+| ar_pipeline | `DreamX-World/pipeline/pipeline_causal_camera.py`; AR block/KV/context-noise loop | `tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py`; A40 short full-generation smoke | Dedicated `DreamXWorldARCausalDenoisingStage` implements blockwise KV forcing; raw HF repo must be converted before `VideoGenerator.from_pretrained` because it has no `model_index.json`; converted AR layout generated a 64x64/9-frame/4-step MP4 on A40. | config_registry_and_short_full_generation_pass |
Include reused components in parity. Reuse is accepted only after the FastVideo
component definition and official instantiation arguments have both been checked
@@ -161,6 +193,30 @@ python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py
```bash
python -m pytest tests/local_tests/dreamx_world/ -v -s
# 2026-07-01: 26 passed, 0 skipped
+# 2026-07-02: AR targeted suite passed: 16 passed, 0 skipped
+
+PYTHONPATH=/workspace/FastVideo python /tmp/run_dreamx_ar_video_smoke.py
+# 2026-07-02: DreamX-World-5B AR short full-generation smoke passed on A40
+# output: outputs_video/dreamx_world_ar_video_smoke/a quiet road through a futuristic coastal city at sunrise.mp4
+# decoded: 9 frames, (64, 64, 3), uint8
+
+PYTHONPATH=/workspace/FastVideo FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA python /tmp/run_dreamx_ar_long_horizon.py
+# 2026-07-02: DreamX-World-5B AR long-horizon generation passed on A40
+# output: outputs_video/dreamx_world_ar_long_horizon/A long autonomous drive through a futuristic coastal city at sunrise, smooth forward camera motion,.mp4
+# decoded: 1005 frames, (64, 64, 3)
+
+PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B DREAMX_WORLD_AR_SSIM_MODEL_PATH=/tmp/converted_dreamx_world_ar python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs --skip-ssim-reference-download
+# 2026-07-02: DreamX-World-5B AR default SSIM passed: 1 passed, 0 skipped
+# local A40 reference: fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B/TORCH_SDPA/
+
+modal run /tmp/modal_dreamx_ar_ssim_git.py
+# 2026-07-02: Modal L40S default SSIM passed: 1 passed, 0 skipped
+# checked out post-fix dreamx-world-5b-cam branch commit
+# first downloaded default L40S references with reference_videos_cli.py download
+# seeded missing AR L40S reference, then reran a fresh generated-vs-reference compare
+# Modal JSON: mean_ssim=1.0, min_ssim=1.0, max_ssim=1.0
+# local L40S reference: fastvideo/tests/ssim/reference_videos/default/L40S_reference_videos/DreamX-World-5B/TORCH_SDPA/
+# A40-vs-L40S reference spot check after deterministic AR noise fix: mean SSIM 0.7770, min 0.6139, max 0.9944
python -m pytest tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py -q -rs
# 2026-06-30: 4 passed, 0 skipped
@@ -170,17 +226,21 @@ PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B-Cam pyth
# 2026-07-01: 1 passed, 0 skipped
```
-`camera_conditioning`, default `Flow` scheduler, DreamX component/pipeline configs, default preset, pipeline entry/registry, FlowMatch scheduler initialization, DreamX camera stage, denoising `y_camera` pass-through, conversion model-index/config-consistency checks, real converted 5B dedicated DreamX transformer strict-load and forward parity on CUDA, VAE encode parity on CUDA, text hidden-state parity on CUDA, and the native DreamX PRoPE branch structure smoke are non-skip PASS results. The full local DreamX component suite, independent pipeline smoke/parity suite, image-path TI2V smoke, and SSIM quality regression currently have zero skips. The basic example was also run against `converted_weights/dreamx_world` with a 64x64/9-frame/1-step saved-video smoke, and imageio decoded the generated MP4 successfully.
+`camera_conditioning`, default `Flow` scheduler, DreamX component/pipeline configs, default preset, pipeline entry/registry, FlowMatch scheduler initialization, DreamX camera stage, denoising `y_camera` pass-through, conversion model-index/config-consistency checks, real converted 5B dedicated DreamX transformer strict-load and forward parity on CUDA, VAE encode parity on CUDA, text hidden-state parity on CUDA, native DreamX PRoPE branch structure smoke, AR transformer parity/strict-load, AR config/registry, and AR short full-generation smoke are non-skip PASS results. The full local DreamX component suite, independent pipeline smoke/parity suite, image-path TI2V smoke, and SSIM quality regression currently have zero skips. The basic 5B-Cam example was run against `converted_weights/dreamx_world` with a 64x64/9-frame/1-step saved-video smoke, and imageio decoded the generated MP4 successfully; the AR pipeline separately saved and decoded a 64x64/9-frame/4-step MP4 from `/tmp/converted_dreamx_world_ar`.
## Review Notes
- Required before handoff: non-skip PASS for each required component parity
test, including reused components that own weights or numerical behavior.
-- First PR scope is `DreamX-World-5B-Cam`; do not pull in
- `inference_ar_forcing.py` or `pipeline_causal_camera.py` unless the scope is
- explicitly expanded.
-- AR long-horizon support is out of first-PR scope because it needs a separate causal/KV cache
- forcing pipeline and likely separate parity gates.
+- First PR scope originally targeted `DreamX-World-5B-Cam`; scope was later
+ expanded to include `DreamX-World-5B` AR support with a separate causal/KV
+ pipeline.
+- AR support has targeted parity/config coverage, short full-generation smoke,
+ 1005-frame A40 long-horizon generation, local A40 default SSIM coverage,
+ and Modal L40S default SSIM coverage. L40S validation was rerun after fixing
+ stale generated-output comparison in the SSIM helper and deterministic AR
+ noise seeding. HF reference publication remains a separate token-gated
+ release operation.
- FastVideo production code must not require `xfuser` or OpenCV just because the
official reference import needed them. Port camera/action preprocessing and
sequence-parallel behavior into existing FastVideo-native utilities or keep
@@ -191,4 +251,4 @@ PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B-Cam pyth
## Quality Regression
-Quality regression is added in `fastvideo/tests/ssim/test_dreamx_world_similarity.py`. The default test uses the workspace converted model root, `TORCH_SDPA`, a deterministic generated conditioning image, 9 frames, 1 denoise step, seed 1024, and min SSIM 0.98. A local A40 reference was seeded under `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B-Cam/TORCH_SDPA/`, and the test passed non-skip locally on 2026-07-01. Full-quality params are present for 480x832, 161 frames, and 30 steps; publishing an L40S/HF reference remains a release operation requiring a write-capable HF token env var, never a raw token value.
+Quality regression is added in `fastvideo/tests/ssim/test_dreamx_world_similarity.py`. The 5B-Cam default test uses the workspace converted model root, `TORCH_SDPA`, a deterministic generated conditioning image, 9 frames, 1 denoise step, seed 1024, and min SSIM 0.98. A local A40 reference was seeded under `fastvideo/tests/ssim/reference_videos/default/A40_reference_videos/DreamX-World-5B-Cam/TORCH_SDPA/`, and the test passed non-skip locally on 2026-07-01. The AR default test uses `/tmp/converted_dreamx_world_ar` or `/root/data/dreamx_world_ar_converted`, `TORCH_SDPA`, 192x192, 81 frames, 4 steps, seed 2048, and min SSIM 0.98; local A40 and Modal L40S references were seeded under `fastvideo/tests/ssim/reference_videos/default/{A40,L40S}_reference_videos/DreamX-World-5B/TORCH_SDPA/`, and both tests passed non-skip on 2026-07-02 after fresh generated-output cleanup was added to the SSIM helper. Modal L40S validation checked out `post-fix dreamx-world-5b-cam branch commit`, seeded the missing AR L40S reference, reran the test, and wrote `mean_ssim=1.0`. A cross-device A40-vs-L40S reference spot check produced mean SSIM 0.7770, so CI should use the L40S-specific reference rather than the A40 artifact. Full-quality params are present for 5B-Cam 480x832/161 frames/30 steps and AR 704x1280/1005 frames/4 steps; publishing references to the HF dataset remains a release operation requiring a write-capable HF token env var, never a raw token value.
diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py b/tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py
new file mode 100644
index 000000000..f80c41917
--- /dev/null
+++ b/tests/local_tests/dreamx_world/test_dreamx_world_ar_conversion.py
@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamX-World-5B autoregressive conversion smoke tests."""
+
+import json
+
+from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_ar_dit_config
+from fastvideo.models.registry import _LEGACY_FAST_VIDEO_MODELS
+from scripts.checkpoint_conversion.dreamx_world_ar_to_diffusers import (
+ MODEL_INDEX,
+ REUSED_COMPONENTS,
+ TRANSFORMER_CONFIG,
+ convert_transformer,
+ write_model_index,
+)
+
+
+def test_dreamx_world_ar_converter_writes_symlinked_transformer_and_model_index(tmp_path):
+ source = tmp_path / "raw"
+ source.mkdir()
+ raw_tensor = source / "model.safetensors"
+ raw_tensor.write_bytes(b"placeholder")
+ component_source = tmp_path / "wan22"
+ output = tmp_path / "dreamx_ar"
+
+ for component in REUSED_COMPONENTS:
+ component_dir = component_source / component
+ component_dir.mkdir(parents=True)
+ (component_dir / "config.json").write_text("{}\n")
+
+ convert_transformer(source, output, symlink_transformer=True)
+ write_model_index(output, component_source, symlink_components=True)
+
+ assert (output / "transformer" / "model.safetensors").is_symlink()
+ model_index = json.loads((output / "model_index.json").read_text())
+ assert model_index == MODEL_INDEX
+ assert model_index["_class_name"] == "DreamXWorldARPipeline"
+ assert model_index["transformer"] == ["diffusers", "DreamXWorldARTransformer3DModel"]
+
+
+def test_dreamx_world_ar_transformer_config_matches_pipeline_dit_config():
+ dit_config = make_dreamx_world_5b_ar_dit_config()
+ assert TRANSFORMER_CONFIG["_class_name"] == "DreamXWorldARTransformer3DModel"
+ assert TRANSFORMER_CONFIG["num_attention_heads"] == dit_config.num_attention_heads
+ assert TRANSFORMER_CONFIG["attention_head_dim"] == dit_config.attention_head_dim
+ assert TRANSFORMER_CONFIG["in_channels"] == dit_config.in_channels
+ assert TRANSFORMER_CONFIG["out_channels"] == dit_config.out_channels
+ assert TRANSFORMER_CONFIG["ffn_dim"] == dit_config.ffn_dim
+ assert TRANSFORMER_CONFIG["num_layers"] == dit_config.num_layers
+ assert TRANSFORMER_CONFIG["local_attn_size"] == dit_config.local_attn_size
+ assert TRANSFORMER_CONFIG["sink_size"] == dit_config.sink_size
+ assert TRANSFORMER_CONFIG["attn_compress"] == dit_config.attn_compress
+ assert tuple(TRANSFORMER_CONFIG["cam_self_attn_layers"]) == dit_config.cam_self_attn_layers
+
+
+def test_dreamx_world_ar_model_index_component_classes_are_registered():
+ for component in ("scheduler", "text_encoder", "transformer", "vae"):
+ class_name = MODEL_INDEX[component][1]
+ assert class_name in _LEGACY_FAST_VIDEO_MODELS
diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py
new file mode 100644
index 000000000..ff0eb34c9
--- /dev/null
+++ b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py
@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DreamX-World-5B autoregressive transformer parity.
+
+Coverage scope: both. The tiny forward parity compares FastVideo's native AR DiT
+against the official DreamX ``CausalWanModel`` implementation with identical
+weights. The real 5B checkpoint gate strict-loads the downloaded safetensors
+through the FastVideo model class.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import sys
+
+import pytest
+import torch
+from torch.testing import assert_close
+
+from fastvideo.configs.models.dits.dreamx_world import DreamXWorldARArchConfig, DreamXWorldARConfig
+from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_ar_dit_config
+from fastvideo.models.dits.dreamx_world_ar import DreamXWorldARTransformer3DModel
+from fastvideo.models.loader.fsdp_load import load_model_from_full_model_state_dict
+from fastvideo.models.loader.utils import get_param_names_mapping
+from fastvideo.models.loader.weight_utils import resolve_safetensors_files, safetensors_weights_iterator
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World"))
+CONVERTED_AR_DIR = Path(os.getenv("DREAMX_WORLD_AR_CONVERTED_DIR", "/tmp/converted_dreamx_world_ar"))
+PARITY_SCOPE = "both"
+
+
+def _tiny_config() -> DreamXWorldARConfig:
+ return DreamXWorldARConfig(
+ arch_config=DreamXWorldARArchConfig(
+ num_attention_heads=1,
+ attention_head_dim=8,
+ in_channels=4,
+ out_channels=4,
+ ffn_dim=16,
+ num_layers=1,
+ text_dim=8,
+ freq_dim=8,
+ text_len=4,
+ local_attn_size=2,
+ sink_size=1,
+ attn_compress=1,
+ cam_self_attn_layers=(0,),
+ ))
+
+
+def _load_official_tiny():
+ if not OFFICIAL_REF_DIR.exists():
+ pytest.skip(f"Official DreamX reference missing: {OFFICIAL_REF_DIR}")
+ sys.path.insert(0, str(OFFICIAL_REF_DIR))
+ try:
+ from wan.modules import attention as official_attention
+ from wan.modules import causal_camera_model_2_2_prope_infinity as causal_module
+ from wan.modules import model_2_2 as official_model_2_2
+ from wan.modules.causal_camera_model_2_2_prope_infinity import CausalWanModel
+ except Exception as exc: # noqa: BLE001
+ pytest.skip(f"Cannot import official AR transformer: {exc}")
+ official_attention.FLASH_ATTN_2_AVAILABLE = False
+ official_attention.FLASH_ATTN_3_AVAILABLE = False
+
+ def _sdpa_same_dtype(q, k, v, **kwargs):
+ del kwargs
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), dropout_p=0.0)
+ return out.transpose(1, 2).contiguous()
+
+ official_attention.attention = _sdpa_same_dtype
+ official_attention.flash_attention = _sdpa_same_dtype
+ official_model_2_2.flash_attention = _sdpa_same_dtype
+ causal_module.attention = _sdpa_same_dtype
+ return CausalWanModel(
+ model_type="ti2v",
+ patch_size=(1, 2, 2),
+ text_len=4,
+ in_dim=4,
+ dim=8,
+ ffn_dim=16,
+ freq_dim=8,
+ text_dim=8,
+ out_dim=4,
+ num_heads=1,
+ num_layers=1,
+ local_attn_size=2,
+ sink_size=1,
+ qk_norm=True,
+ cross_attn_norm=True,
+ add_control_adapter=True,
+ cam_method="prope",
+ attn_compress=1,
+ cam_self_attn_layers=(0,),
+ ).eval()
+
+
+def _make_inputs():
+ torch.manual_seed(123)
+ x = [torch.randn(4, 1, 4, 4)]
+ t = torch.zeros(1, 4, dtype=torch.long)
+ context = [torch.randn(2, 8)]
+ camera = {
+ "viewmats": torch.eye(4).reshape(1, 1, 4, 4).repeat(1, 4, 1, 1),
+ "K": torch.eye(3).reshape(1, 1, 3, 3).repeat(1, 4, 1, 1),
+ }
+ kv_cache = [{
+ "k": torch.zeros(1, 8, 1, 8),
+ "v": torch.zeros(1, 8, 1, 8),
+ "global_end_index": torch.tensor([0]),
+ "local_end_index": torch.tensor([0]),
+ "prope_k": torch.zeros(1, 8, 1, 8),
+ "prope_v": torch.zeros(1, 8, 1, 8),
+ "prope_global_end_index": torch.tensor([0]),
+ "prope_local_end_index": torch.tensor([0]),
+ }]
+ cross_cache = [{
+ "k": torch.zeros(1, 4, 1, 8),
+ "v": torch.zeros(1, 4, 1, 8),
+ "is_init": False,
+ }]
+ return x, t, context, camera, kv_cache, cross_cache
+
+
+def test_dreamx_world_ar_tiny_forward_matches_official():
+ official = _load_official_tiny()
+ fastvideo = DreamXWorldARTransformer3DModel(_tiny_config(), {}).eval()
+ fastvideo.load_state_dict(official.state_dict(), strict=True)
+
+ x, t, context, camera, kv_cache, cross_cache = _make_inputs()
+ official_out = official(x=x, t=t, context=context, seq_len=4, y_camera=camera,
+ kv_cache=kv_cache, crossattn_cache=cross_cache).detach()
+ x, t, context, camera, kv_cache, cross_cache = _make_inputs()
+ fastvideo_out = fastvideo(x=x, t=t, context=context, seq_len=4, y_camera=camera,
+ kv_cache=kv_cache, crossattn_cache=cross_cache).detach()
+ assert_close(fastvideo_out, official_out, atol=1e-5, rtol=1e-5)
+
+
+def test_dreamx_world_ar_5b_config_matches_official_shape():
+ config = make_dreamx_world_5b_ar_dit_config()
+ assert config.num_layers == 30
+ assert config.num_attention_heads == 24
+ assert config.attention_head_dim == 128
+ assert config.hidden_size == 3072
+ assert config.ffn_dim == 14336
+ assert config.local_attn_size == 12
+ assert config.sink_size == 3
+ assert config.attn_compress == 4
+ assert config.cam_self_attn_layers == tuple(range(30))
+
+
+def test_dreamx_world_ar_converted_5b_transformer_strict_loads():
+ transformer_dir = CONVERTED_AR_DIR / "transformer"
+ if not transformer_dir.exists():
+ pytest.skip(f"Converted AR transformer missing: {transformer_dir}")
+ with torch.device("meta"):
+ model = DreamXWorldARTransformer3DModel(make_dreamx_world_5b_ar_dit_config(), {})
+ incompatible = load_model_from_full_model_state_dict(
+ model,
+ safetensors_weights_iterator(resolve_safetensors_files(str(transformer_dir)), to_cpu=True),
+ device=torch.device("cpu"),
+ param_dtype=torch.bfloat16,
+ strict=True,
+ param_names_mapping=get_param_names_mapping(model.param_names_mapping),
+ training_mode=False,
+ )
+ assert incompatible.missing_keys == []
+ assert incompatible.unexpected_keys == []
+ assert not any(param.is_meta for param in model.parameters())
diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py b/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py
index 5815478a0..dbc3a26ab 100644
--- a/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py
+++ b/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py
@@ -10,6 +10,7 @@ from scripts.checkpoint_conversion.dreamx_world_to_diffusers import (
MODEL_INDEX,
REUSED_COMPONENTS,
TRANSFORMER_CONFIG,
+ _copy_or_link_component,
write_model_index,
)
@@ -62,3 +63,17 @@ def test_dreamx_world_model_index_component_classes_are_registered():
for component in ("scheduler", "text_encoder", "transformer", "vae"):
class_name = MODEL_INDEX[component][1]
assert class_name in _LEGACY_FAST_VIDEO_MODELS
+
+
+def test_dreamx_world_copy_or_link_component_keeps_broken_symlink(tmp_path):
+ component_source = tmp_path / "wan22"
+ src = component_source / "scheduler"
+ src.mkdir(parents=True)
+ output = tmp_path / "dreamx"
+ output.mkdir()
+ dst = output / "scheduler"
+ dst.symlink_to(tmp_path / "missing_scheduler", target_is_directory=True)
+
+ _copy_or_link_component("scheduler", component_source, output, symlink=True)
+
+ assert dst.is_symlink()
diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py
index b9993e479..385c1f613 100644
--- a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py
+++ b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py
@@ -4,6 +4,7 @@
from types import SimpleNamespace
import json
+import numpy as np
import torch
from fastvideo.api.presets import get_preset
@@ -11,8 +12,14 @@ from fastvideo.fastvideo_args import WorkloadType
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.pipeline_registry import PipelineType, import_pipeline_classes
-from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import build_dreamx_camera_condition
-from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
+from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import (
+ DreamXCamera,
+ _interpolate_camera_poses,
+ build_dreamx_camera_condition,
+)
+from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BARPipelineConfig, DreamXWorld5BCamPipelineConfig
+from fastvideo.pipelines.basic.dreamx_world.ar_denoising import DreamXWorldARCausalDenoisingStage
+from fastvideo.pipelines.basic.dreamx_world.dreamx_world_ar_pipeline import DreamXWorldARPipeline
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import DreamXWorldPipeline
from fastvideo.pipelines.basic.dreamx_world.stages import (
DREAMX_Y_CAMERA_KEY,
@@ -150,3 +157,150 @@ def test_dreamx_world_denoising_kwargs_filter_for_y_camera():
assert stage.prepare_extra_func_kwargs(accepts_y_camera, {"y_camera": y_camera}) == {"y_camera": y_camera}
assert stage.prepare_extra_func_kwargs(no_y_camera, {"y_camera": y_camera}) == {}
+
+
+def test_dreamx_world_ar_pipeline_config_wires_components():
+ config = DreamXWorld5BARPipelineConfig()
+ assert config.is_causal is True
+ assert config.flow_shift == 5.0
+ assert config.dmd_denoising_steps == (1000, 750, 500, 250)
+ assert config.warp_denoising_step is True
+ assert config.context_noise == 0.1
+ assert config.dit_config.arch_config.local_attn_size == 12
+ assert config.dit_config.arch_config.sink_size == 3
+ assert config.dit_config.arch_config.attn_compress == 4
+
+
+def test_dreamx_world_ar_pipeline_registry_and_preset():
+ from fastvideo.api.presets import get_preset
+ from fastvideo.registry import get_pipeline_config_cls_from_name, get_preset_selection
+
+ assert get_pipeline_config_cls_from_name("GD-ML/DreamX-World-5B") is DreamXWorld5BARPipelineConfig
+ preset_name, family = get_preset_selection("GD-ML/DreamX-World-5B")
+ assert (preset_name, family) == ("dreamx_world_5b_ar", "dreamx_world")
+ preset = get_preset("dreamx_world_5b_ar", "dreamx_world")
+ assert preset.defaults["num_inference_steps"] == 4
+ assert DreamXWorldARPipeline.pipeline_config_cls is DreamXWorld5BARPipelineConfig
+
+
+def test_dreamx_world_camera_conditioning_stage_expands_scalar_speed():
+ batch = ForwardBatch(
+ data_type="t2v",
+ action_list=["w", "d"],
+ action_speed_list=2.0,
+ num_frames=17,
+ height=704,
+ width=1280,
+ latents=torch.zeros(1, 16, 5, 44, 80),
+ )
+ stage = DreamXWorldCameraConditioningStage()
+
+ out = stage.forward(batch, fastvideo_args=object())
+
+ assert set(out.extra[DREAMX_Y_CAMERA_KEY]) == {"viewmats", "K"}
+
+
+def test_dreamx_world_camera_interpolation_handles_single_camera():
+ camera = DreamXCamera(
+ fx=0.8,
+ fy=0.8,
+ cx=0.5,
+ cy=0.5,
+ w2c_mat=np.eye(4, dtype=np.float64),
+ )
+
+ out = _interpolate_camera_poses(
+ [camera],
+ src_indices=np.array([0.0]),
+ tgt_indices=np.array([0.0, 1.0, 2.0]),
+ )
+
+ assert out == [camera, camera, camera]
+
+
+def test_dreamx_world_ar_cache_initializes_camera_self_attention_entries():
+ cam_self_attn = SimpleNamespace(num_heads=3, head_dim=5)
+ transformer = SimpleNamespace(
+ blocks=[SimpleNamespace(cam_self_attn=cam_self_attn), SimpleNamespace(cam_self_attn=cam_self_attn)],
+ num_attention_heads=2,
+ attention_head_dim=4,
+ )
+ stage = DreamXWorldARCausalDenoisingStage.__new__(DreamXWorldARCausalDenoisingStage)
+ stage.transformer = transformer
+ stage.num_transformer_blocks = 2
+ stage.local_attn_size = 6
+
+ caches = stage._initialize_kv_cache(
+ batch_size=1,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ frame_seq_length=7,
+ )
+
+ assert len(caches) == 2
+ assert caches[0]["k"].shape == (1, 42, 2, 4)
+ assert caches[0]["prope_k"].shape == (1, 42, 3, 5)
+ assert caches[0]["prope_v"].shape == (1, 42, 3, 5)
+ assert int(caches[0]["prope_global_end_index"].item()) == 0
+ assert int(caches[0]["prope_local_end_index"].item()) == 0
+
+
+def test_dreamx_world_ar_context_noise_fraction_maps_to_scheduler_timestep():
+ assert DreamXWorldARCausalDenoisingStage._context_noise_timestep(0.1) == 100
+ assert DreamXWorldARCausalDenoisingStage._context_noise_timestep(100) == 100
+
+
+
+def test_dreamx_world_ar_context_update_advances_camera_cache_indices():
+ class DummyTransformer:
+ def __call__(self, *, hidden_states, encoder_hidden_states, timestep, y_camera, kv_cache, crossattn_cache,
+ current_start):
+ del encoder_hidden_states, y_camera, crossattn_cache
+ assert current_start == 0
+ assert timestep.unique().tolist() == [100]
+ new_tokens = timestep.shape[1]
+ for cache in kv_cache:
+ cache["local_end_index"] += new_tokens
+ cache["global_end_index"] += new_tokens
+ cache["prope_local_end_index"] += new_tokens
+ cache["prope_global_end_index"] += new_tokens
+ cache["k"][:, :new_tokens] = 1
+ cache["prope_k"][:, :new_tokens] = 1
+ return hidden_states
+
+ cam_self_attn = SimpleNamespace(num_heads=3, head_dim=5)
+ cache_transformer = SimpleNamespace(
+ blocks=[SimpleNamespace(cam_self_attn=cam_self_attn)],
+ num_attention_heads=2,
+ attention_head_dim=4,
+ )
+ stage = DreamXWorldARCausalDenoisingStage.__new__(DreamXWorldARCausalDenoisingStage)
+ stage.transformer = cache_transformer
+ stage.num_transformer_blocks = 1
+ stage.local_attn_size = 6
+ caches = stage._initialize_kv_cache(
+ batch_size=1,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ frame_seq_length=2,
+ )
+ # Keep the cache allocation source separate from the callable transformer used by _update_context_cache.
+ stage.transformer = DummyTransformer()
+
+ stage._update_context_cache(
+ block_latents=torch.zeros(1, 4, 3, 2, 2),
+ context=[torch.zeros(2, 4)],
+ camera_block={"viewmats": torch.eye(4).reshape(1, 1, 4, 4), "K": torch.eye(3).reshape(1, 1, 3, 3)},
+ kv_cache=caches,
+ crossattn_cache=[{}],
+ start=0,
+ frame_seq_length=2,
+ target_dtype=torch.float32,
+ autocast_enabled=False,
+ context_noise=0.1,
+ )
+
+ assert int(caches[0]["local_end_index"].item()) == 6
+ assert int(caches[0]["prope_local_end_index"].item()) == 6
+ assert caches[0]["k"][:, :6].sum().item() == 48
+ assert caches[0]["prope_k"][:, :6].sum().item() == 90