Add DreamX World 5B Cam pipeline

This commit is contained in:
Suckl
2026-07-05 11:43:52 -07:00
committed by SolitaryThinker
parent 6a32cf3a5e
commit dcf0ce2600
30 changed files with 3224 additions and 5 deletions
+2
View File
@@ -39,6 +39,7 @@ converted_weights/
# SSIM test outputs
fastvideo/tests/ssim/generated_videos/
fastvideo/tests/ssim/.reference_videos_download.lock
**/.cache/**
@@ -132,3 +133,4 @@ fastvideo/tests/ssim/reference_videos/**
*.nvimlog
.nvimlog
.python-version
/DreamX-World/
@@ -0,0 +1,64 @@
import os
from fastvideo import VideoGenerator
OUTPUT_PATH = os.getenv("DREAMX_WORLD_OUTPUT_PATH", "video_samples_dreamx_world")
def _env_int(name: str, default: int) -> int:
return int(os.getenv(name, str(default)))
def _env_float(name: str, default: float) -> float:
return float(os.getenv(name, str(default)))
def main():
model_name = os.getenv("DREAMX_WORLD_MODEL_DIR", "GD-ML/DreamX-World-5B-Cam")
generator = VideoGenerator.from_pretrained(
model_name,
num_gpus=1,
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="DreamXWorldPipeline",
)
prompt = os.getenv(
"DREAMX_WORLD_PROMPT",
"A cinematic first-person drive through a futuristic coastal city at "
"sunrise, reflective glass towers, clean streets, soft volumetric light.",
)
image_path = os.getenv(
"DREAMX_WORLD_IMAGE_PATH",
"https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG",
)
kwargs = {
"output_path": OUTPUT_PATH,
"save_video": os.getenv("DREAMX_WORLD_SAVE_VIDEO", "1") != "0",
"height": _env_int("DREAMX_WORLD_HEIGHT", 480),
"width": _env_int("DREAMX_WORLD_WIDTH", 832),
"num_frames": _env_int("DREAMX_WORLD_NUM_FRAMES", 161),
"num_inference_steps": _env_int("DREAMX_WORLD_STEPS", 30),
"guidance_scale": _env_float("DREAMX_WORLD_GUIDANCE", 5.0),
"action_list": os.getenv("DREAMX_WORLD_ACTIONS", "w,d,w").split(","),
"action_speed_list": [
float(value)
for value in os.getenv("DREAMX_WORLD_ACTION_SPEEDS", "4.0,2.0,4.0").split(",")
],
}
if image_path:
kwargs["image_path"] = image_path
try:
generator.generate_video(prompt, **kwargs)
finally:
generator.shutdown()
if __name__ == "__main__":
main()
+2 -1
View File
@@ -1,5 +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.flux_2 import Flux2Config
from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
@@ -13,7 +14,7 @@ from fastvideo.configs.models.dits.hyworld import HYWorldConfig
from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
__all__ = [
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
"HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "DreamXWorldConfig", "CosmosVideoConfig",
"Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
"MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
]
@@ -0,0 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
@dataclass
class DreamXWorldArchConfig(WanVideoArchConfig):
"""DreamX-World DiT config with camera PRoPE control fields."""
add_control_adapter: bool = True
cam_method: str | None = "prope"
attn_compress: int = 1
cam_self_attn_layers: tuple[int, ...] | None = None
@dataclass
class DreamXWorldConfig(DiTConfig):
arch_config: DiTArchConfig = field(default_factory=DreamXWorldArchConfig)
prefix: str = "Wan"
+2 -1
View File
@@ -1,6 +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.hunyuan import FastHunyuanConfig, HunyuanConfig
from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyuan15T2V720PConfig
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
@@ -16,5 +17,5 @@ __all__ = [
"HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
"Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
"SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
"HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
"DreamXWorld5BCamPipelineConfig", "HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
]
@@ -0,0 +1,87 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World-5B-Cam FastVideo model configuration helpers."""
from collections.abc import Callable
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.encoders import BaseEncoderOutput, T5Config
from fastvideo.configs.models.encoders.t5 import T5ArchConfig
from fastvideo.configs.models.vaes import WanVAEConfig
from fastvideo.configs.pipelines.base import PipelineConfig
from fastvideo.configs.pipelines.wan import LucyEditDevConfig, t5_postprocess_text
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,
))
def make_dreamx_world_5b_cam_vae_config() -> WanVAEConfig:
"""Return the Wan2.2 48-channel VAE config used by DreamX-World-5B-Cam."""
return LucyEditDevConfig().vae_config
def make_dreamx_world_5b_cam_text_encoder_config() -> T5Config:
"""Return the UMT5-XXL text encoder config used by DreamX-World-5B-Cam."""
return T5Config(
arch_config=T5ArchConfig(
vocab_size=256384,
d_model=4096,
d_kv=64,
d_ff=10240,
num_layers=24,
num_decoder_layers=None,
num_heads=64,
relative_attention_num_buckets=32,
dropout_rate=0.0,
text_len=512,
feed_forward_proj="gelu",
is_encoder_decoder=False,
),
prefix="umt5",
)
@dataclass
class DreamXWorld5BCamPipelineConfig(PipelineConfig):
"""Pipeline config for the first-scope DreamX-World-5B-Cam mode."""
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_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
flow_shift: float | None = 3.0
ti2v_task: bool = True
expand_timesteps: bool = True
vae_tiling: bool = False
vae_sp: bool = False
vae_precision: str = "fp32"
vae_decode_precision: str | None = "bf16"
dit_precision: str = "bf16"
def __post_init__(self) -> None:
self.vae_config.load_encoder = True
self.vae_config.load_decoder = True
self.dit_config.expand_timesteps = self.expand_timesteps
+500
View File
@@ -0,0 +1,500 @@
# SPDX-License-Identifier: Apache-2.0
import math
from typing import Any
import torch
import torch.nn as nn
from fastvideo.configs.models.dits.dreamx_world import DreamXWorldConfig
from fastvideo.distributed.communication_op import (
sequence_model_parallel_all_gather_with_unpad,
sequence_model_parallel_shard)
from fastvideo.distributed.parallel_state import get_sp_world_size
from fastvideo.layers.layernorm import RMSNorm
from fastvideo.layers.linear import ReplicatedLinear
from fastvideo.layers.rotary_embedding import get_rotary_pos_embed
from fastvideo.models.dits.wanvideo import (LayerNormScaleShift,
PatchEmbed,
WanTimeTextImageEmbedding,
WanTransformer3DModel,
WanTransformerBlock)
from fastvideo.platforms import AttentionBackendEnum, current_platform
from fastvideo.attention import LocalAttention
from fastvideo.layers.quantization import QuantizationConfig
from fastvideo.models.dits.base import BaseDiT
def _dreamx_invert_se3(transforms: torch.Tensor) -> torch.Tensor:
assert transforms.shape[-2:] == (4, 4)
rot_inv = transforms[..., :3, :3].transpose(-1, -2)
out = torch.zeros_like(transforms)
out[..., :3, :3] = rot_inv
out[..., :3, 3] = -torch.einsum("...ij,...j->...i", rot_inv,
transforms[..., :3, 3])
out[..., 3, 3] = 1.0
return out.to(dtype=transforms.dtype)
def _dreamx_lift_k(intrinsics: torch.Tensor) -> torch.Tensor:
assert intrinsics.shape[-2:] == (3, 3)
out = torch.zeros(intrinsics.shape[:-2] + (4, 4),
device=intrinsics.device,
dtype=intrinsics.dtype)
out[..., :3, :3] = intrinsics
out[..., 3, 3] = 1.0
return out
def _dreamx_invert_k(intrinsics: torch.Tensor) -> torch.Tensor:
assert intrinsics.shape[-2:] == (3, 3)
out = torch.zeros_like(intrinsics)
out[..., 0, 0] = 1.0 / intrinsics[..., 0, 0]
out[..., 1, 1] = 1.0 / intrinsics[..., 1, 1]
out[..., 0, 2] = -intrinsics[..., 0, 2] / intrinsics[..., 0, 0]
out[..., 1, 2] = -intrinsics[..., 1, 2] / intrinsics[..., 1, 1]
out[..., 2, 2] = 1.0
return out.to(dtype=intrinsics.dtype)
def _dreamx_apply_tiled_projmat(feats: torch.Tensor,
matrix: torch.Tensor) -> torch.Tensor:
batch, num_heads, seq_len, feat_dim = feats.shape
proj_dim = matrix.shape[-1]
assert feat_dim % proj_dim == 0
if matrix.shape[1] == seq_len:
feats = feats.view(batch, num_heads, seq_len, feat_dim // proj_dim,
proj_dim)
out = torch.einsum("btij,bntpj->bntpi", matrix, feats)
return out.reshape(batch, num_heads, seq_len, feat_dim)
cameras = matrix.shape[1]
assert seq_len > cameras and seq_len % cameras == 0
feats = feats.reshape(batch, num_heads, cameras, -1,
feat_dim // proj_dim, proj_dim)
out = torch.einsum("bcij,bncpkj->bncpki", matrix, feats)
return out.reshape(batch, num_heads, seq_len, feat_dim)
def _dreamx_prope_qkv(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
viewmats: torch.Tensor, intrinsics: torch.Tensor):
batch, num_heads, seq_len, head_dim = q.shape
cameras = viewmats.shape[1]
assert q.shape == k.shape == v.shape
assert viewmats.shape == (batch, cameras, 4, 4)
assert intrinsics.shape == (batch, cameras, 3, 3)
assert head_dim % 4 == 0
intrinsics_norm = torch.zeros_like(intrinsics)
intrinsics_norm[..., 0, 0] = intrinsics[..., 0, 0]
intrinsics_norm[..., 1, 1] = intrinsics[..., 1, 1]
intrinsics_norm[..., 2, 2] = 1.0
proj = torch.einsum("...ij,...jk->...ik",
_dreamx_lift_k(intrinsics_norm), viewmats)
proj_t = proj.transpose(-1, -2).to(dtype=viewmats.dtype)
proj_inv = torch.einsum(
"...ij,...jk->...ik",
_dreamx_invert_se3(viewmats),
_dreamx_lift_k(_dreamx_invert_k(intrinsics_norm)),
).to(dtype=viewmats.dtype)
q = _dreamx_apply_tiled_projmat(q, proj_t)
k = _dreamx_apply_tiled_projmat(k, proj_inv)
v = _dreamx_apply_tiled_projmat(v, proj_inv)
return q, k, v, proj
class DreamXPropeSelfAttention(nn.Module):
"""DreamX-World parallel PRoPE camera self-attention branch."""
def __init__(self,
dim: int,
attn_dim: int,
num_heads: int,
qk_norm: str | bool = True,
eps: float = 1e-6,
quant_config: QuantizationConfig | None = None,
prefix: str = ""):
super().__init__()
assert attn_dim % num_heads == 0
self.attn_dim = attn_dim
self.num_heads = num_heads
self.head_dim = attn_dim // num_heads
self.qk_norm = qk_norm
self.q_proj = ReplicatedLinear(dim,
attn_dim,
quant_config=quant_config,
prefix=f"{prefix}.q_proj")
self.k_proj = ReplicatedLinear(dim,
attn_dim,
quant_config=quant_config,
prefix=f"{prefix}.k_proj")
self.v_proj = ReplicatedLinear(dim,
attn_dim,
quant_config=quant_config,
prefix=f"{prefix}.v_proj")
self.out_proj = ReplicatedLinear(attn_dim,
dim,
quant_config=quant_config,
prefix=f"{prefix}.out_proj")
if qk_norm == "rms_norm":
self.norm_q = RMSNorm(self.head_dim, eps=eps)
self.norm_k = RMSNorm(self.head_dim, eps=eps)
elif qk_norm in (True, "rms_norm_across_heads"):
self.norm_q = RMSNorm(attn_dim, eps=eps)
self.norm_k = RMSNorm(attn_dim, eps=eps)
elif qk_norm is False:
self.norm_q = nn.Identity()
self.norm_k = nn.Identity()
else:
raise ValueError(f"Unsupported qk_norm for DreamX PRoPE: {qk_norm}")
nn.init.zeros_(self.out_proj.weight)
if self.out_proj.bias is not None:
nn.init.zeros_(self.out_proj.bias)
self.attn = LocalAttention(
num_heads=num_heads,
head_size=self.head_dim,
dropout_rate=0,
softmax_scale=None,
causal=False,
supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
AttentionBackendEnum.TORCH_SDPA))
def forward(self, hidden_states: torch.Tensor,
y_camera: dict[str, torch.Tensor]) -> torch.Tensor:
batch_size, seq_len, _ = hidden_states.shape
query, _ = self.q_proj(hidden_states)
key, _ = self.k_proj(hidden_states)
value, _ = self.v_proj(hidden_states)
if self.qk_norm == "rms_norm":
query = query.view(batch_size, seq_len, self.num_heads,
self.head_dim)
key = key.view(batch_size, seq_len, self.num_heads, self.head_dim)
query = self.norm_q(query)
key = self.norm_k(key)
else:
query = self.norm_q(query).view(batch_size, seq_len,
self.num_heads, self.head_dim)
key = self.norm_k(key).view(batch_size, seq_len, self.num_heads,
self.head_dim)
value = value.view(batch_size, seq_len, self.num_heads, self.head_dim)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
query, key, value, output_projection = _dreamx_prope_qkv(
query,
key,
value,
viewmats=y_camera["viewmats"],
intrinsics=y_camera["K"],
)
out = self.attn(query.transpose(1, 2), key.transpose(1, 2),
value.transpose(1, 2))
out = _dreamx_apply_tiled_projmat(out.transpose(1, 2),
output_projection).transpose(1, 2)
out = out.flatten(2)
out, _ = self.out_proj(out)
return out
class DreamXWorldTransformerBlock(WanTransformerBlock):
def __init__(self,
dim: int,
ffn_dim: int,
num_heads: int,
qk_norm: str = "rms_norm_across_heads",
cross_attn_norm: bool = False,
eps: float = 1e-6,
added_kv_proj_dim: int | None = None,
supported_attention_backends: tuple[AttentionBackendEnum, ...]
| None = None,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
add_control_adapter: bool = True,
cam_method: str | None = "prope",
attn_compress: int = 1,
cam_self_attn_layers: tuple[int, ...] | None = None,
layer_idx: int | None = None):
super().__init__(dim, ffn_dim, num_heads, qk_norm, cross_attn_norm,
eps, added_kv_proj_dim,
supported_attention_backends, quant_config, prefix)
self.cam_self_attn = None
add_cam_attn = add_control_adapter and cam_method == "prope"
if add_cam_attn and cam_self_attn_layers is not None:
add_cam_attn = layer_idx in cam_self_attn_layers
if add_cam_attn:
if num_heads % attn_compress != 0 or dim % attn_compress != 0:
raise ValueError("DreamX attn_compress must divide dim and num_heads")
self.cam_self_attn = DreamXPropeSelfAttention(
dim,
dim // attn_compress,
num_heads // attn_compress,
qk_norm=qk_norm,
eps=eps,
quant_config=quant_config,
prefix=f"{prefix}.cam_self_attn")
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
temb: torch.Tensor,
freqs_cis: tuple[torch.Tensor, torch.Tensor],
original_seq_len: int,
y_camera: dict[str, torch.Tensor] | None = None,
) -> torch.Tensor:
if hidden_states.dim() == 4:
hidden_states = hidden_states.squeeze(1)
orig_dtype = hidden_states.dtype
if temb.dim() == 4:
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = (
self.scale_shift_table.unsqueeze(0) + temb.float()).chunk(
6, dim=2)
shift_msa = shift_msa.squeeze(2)
scale_msa = scale_msa.squeeze(2)
gate_msa = gate_msa.squeeze(2)
c_shift_msa = c_shift_msa.squeeze(2)
c_scale_msa = c_scale_msa.squeeze(2)
c_gate_msa = c_gate_msa.squeeze(2)
else:
e = self.scale_shift_table + temb.float()
shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(
6, dim=1)
assert shift_msa.dtype == torch.float32
norm_hidden_states = (self.norm1(hidden_states.float()) *
(1 + scale_msa) + shift_msa).to(orig_dtype)
query, _ = self.to_q(norm_hidden_states)
key, _ = self.to_k(norm_hidden_states)
value, _ = self.to_v(norm_hidden_states)
if self.norm_q is not None:
query = self.norm_q(query)
if self.norm_k is not None:
key = self.norm_k(key)
query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1))
attn_output, _ = self.attn1(
query,
key,
value,
original_seq_len,
freqs_cis=freqs_cis,
)
attn_output = attn_output.flatten(2)
attn_output, _ = self.to_out(attn_output)
attn_output = attn_output.squeeze(1)
if self.cam_self_attn is not None and y_camera is not None:
attn_output = attn_output + self.cam_self_attn(
norm_hidden_states, y_camera)
null_shift = null_scale = torch.tensor([0], device=hidden_states.device)
norm_hidden_states, hidden_states = self.self_attn_residual_norm(
hidden_states, attn_output, gate_msa, null_shift, null_scale)
norm_hidden_states, hidden_states = norm_hidden_states.to(
orig_dtype), hidden_states.to(orig_dtype)
attn_output = self.attn2(norm_hidden_states,
context=encoder_hidden_states,
context_lens=None)
norm_hidden_states, hidden_states = self.cross_attn_residual_norm(
hidden_states, attn_output, 1, c_shift_msa, c_scale_msa)
norm_hidden_states, hidden_states = norm_hidden_states.to(
orig_dtype), hidden_states.to(orig_dtype)
ff_output = self.ffn(norm_hidden_states)
hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa)
hidden_states = hidden_states.to(orig_dtype)
return hidden_states
class DreamXWorldTransformer3DModel(WanTransformer3DModel):
_fsdp_shard_conditions = DreamXWorldConfig()._fsdp_shard_conditions
_compile_conditions = DreamXWorldConfig()._compile_conditions
_supported_attention_backends = DreamXWorldConfig(
)._supported_attention_backends
param_names_mapping = DreamXWorldConfig().param_names_mapping
reverse_param_names_mapping = DreamXWorldConfig().reverse_param_names_mapping
lora_param_names_mapping = DreamXWorldConfig().lora_param_names_mapping
def __init__(self, config: DreamXWorldConfig, hf_config: dict[str,
Any]) -> None:
BaseDiT.__init__(self, config=config, hf_config=hf_config)
self.quant_config = config.quant_config
inner_dim = config.num_attention_heads * config.attention_head_dim
self.hidden_size = config.hidden_size
self.num_attention_heads = config.num_attention_heads
self.in_channels = config.in_channels
self.out_channels = config.out_channels
self.num_channels_latents = config.num_channels_latents
self.patch_size = config.patch_size
self.text_len = config.text_len
assert config.num_attention_heads % get_sp_world_size() == 0, f"The number of attention heads ({config.num_attention_heads}) must be divisible by the sequence parallel size ({get_sp_world_size()})"
self.patch_embedding = PatchEmbed(in_chans=config.in_channels,
embed_dim=inner_dim,
patch_size=config.patch_size,
flatten=False)
self.condition_embedder = WanTimeTextImageEmbedding(
dim=inner_dim,
time_freq_dim=config.freq_dim,
text_embed_dim=config.text_dim,
image_embed_dim=config.image_dim,
)
self.blocks = nn.ModuleList([
DreamXWorldTransformerBlock(
inner_dim,
config.ffn_dim,
config.num_attention_heads,
config.qk_norm,
config.cross_attn_norm,
config.eps,
config.added_kv_proj_dim,
self._supported_attention_backends,
quant_config=config.quant_config,
prefix=f"{config.prefix}.blocks.{i}",
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,
layer_idx=i)
for i in range(config.num_layers)
])
self.norm_out = LayerNormScaleShift(inner_dim,
norm_type="layer",
eps=config.eps,
elementwise_affine=False,
dtype=torch.float32,
compute_dtype=torch.float32)
self.proj_out = nn.Linear(
inner_dim, config.out_channels * math.prod(config.patch_size))
self.scale_shift_table = nn.Parameter(
torch.randn(1, 2, inner_dim) / inner_dim**0.5)
self.gradient_checkpointing = False
self.__post_init__()
def forward(self,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | list[torch.Tensor],
timestep: torch.LongTensor,
encoder_hidden_states_image: torch.Tensor | list[torch.Tensor]
| None = None,
guidance=None,
y_camera: dict[str, torch.Tensor] | None = None,
**kwargs) -> torch.Tensor:
orig_dtype = hidden_states.dtype
if encoder_hidden_states is not None and not isinstance(
encoder_hidden_states, torch.Tensor):
encoder_hidden_states = encoder_hidden_states[0]
if isinstance(encoder_hidden_states_image,
list) and len(encoder_hidden_states_image) > 0:
encoder_hidden_states_image = encoder_hidden_states_image[0]
else:
encoder_hidden_states_image = None
batch_size, _, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.patch_size
post_patch_num_frames = num_frames // p_t
post_patch_height = height // p_h
post_patch_width = width // p_w
d = self.hidden_size // self.num_attention_heads
rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)]
freqs_cos, freqs_sin = get_rotary_pos_embed(
(post_patch_num_frames, post_patch_height, post_patch_width),
self.hidden_size,
self.num_attention_heads,
rope_dim_list,
dtype=torch.float32 if current_platform.is_mps() else torch.float64,
rope_theta=10000)
freqs_cis = (freqs_cos.to(hidden_states.device).float(),
freqs_sin.to(hidden_states.device).float())
hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
hidden_states, original_seq_len = sequence_model_parallel_shard(
hidden_states, dim=1)
if timestep.dim() == 2:
ts_seq_len = timestep.shape[1]
timestep = timestep.flatten()
else:
ts_seq_len = None
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
timestep,
encoder_hidden_states,
encoder_hidden_states_image,
timestep_seq_len=ts_seq_len)
if ts_seq_len is not None:
timestep_proj = timestep_proj.unflatten(2, (6, -1))
else:
timestep_proj = timestep_proj.unflatten(1, (6, -1))
if encoder_hidden_states_image is not None:
if encoder_hidden_states is not None:
encoder_hidden_states = torch.concat(
[encoder_hidden_states_image, encoder_hidden_states], dim=1)
else:
encoder_hidden_states = encoder_hidden_states_image
if current_platform.is_mps() or current_platform.is_npu():
encoder_hidden_states = encoder_hidden_states.to(orig_dtype)
assert encoder_hidden_states.dtype == orig_dtype
if torch.is_grad_enabled() and self.gradient_checkpointing:
for block in self.blocks:
hidden_states = self._gradient_checkpointing_func(
block, hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, original_seq_len, y_camera)
else:
for block in self.blocks:
hidden_states = block(hidden_states, encoder_hidden_states,
timestep_proj, freqs_cis, original_seq_len,
y_camera=y_camera)
if temb.dim() == 3:
shift, scale = (self.scale_shift_table.unsqueeze(0) +
temb.unsqueeze(2)).chunk(2, dim=2)
shift = shift.squeeze(2)
scale = scale.squeeze(2)
else:
shift, scale = (self.scale_shift_table +
temb.unsqueeze(1)).chunk(2, dim=1)
hidden_states = self.norm_out(hidden_states, shift, scale)
hidden_states = sequence_model_parallel_all_gather_with_unpad(
hidden_states, original_seq_len, dim=1)
hidden_states = self.proj_out(hidden_states)
hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames,
post_patch_height,
post_patch_width, p_t, p_h, p_w,
-1)
hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)
return output
EntryClass = DreamXWorldTransformer3DModel
+5 -3
View File
@@ -11,7 +11,7 @@ import tempfile
from abc import ABC, abstractmethod
from collections.abc import Callable, Set
from dataclasses import dataclass, field
from functools import lru_cache
from functools import cache, lru_cache
from typing import NoReturn, TypeVar, cast
import cloudpickle
@@ -32,6 +32,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
"HYWorldTransformer3DModel":
("dits", "hyworld", "HYWorldTransformer3DModel"),
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
"CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"),
"Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"),
@@ -48,6 +49,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
_IMAGE_TO_VIDEO_DIT_MODELS = {
# "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"),
"WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"),
"DreamXWorldTransformer3DModel": ("dits", "dreamx_world", "DreamXWorldTransformer3DModel"),
"CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"),
"MatrixGame2WanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"),
"CausalMatrixGame2WanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"),
@@ -141,7 +143,7 @@ _LEGACY_FAST_VIDEO_MODELS = {
MODELS_PATH = os.path.dirname(__file__)
@lru_cache(maxsize=None)
@cache
def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
discovered_models: dict[str, tuple[str, str, str]] = {}
for root, dirs, files in os.walk(MODELS_PATH):
@@ -156,7 +158,7 @@ def _discover_and_register_models() -> dict[str, tuple[str, str, str]]:
filepath = os.path.join(root, filename)
try:
with open(filepath, "r", encoding="utf-8") as f:
with open(filepath, encoding="utf-8") as f:
source = f.read()
tree = ast.parse(source, filename=filename)
@@ -0,0 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
from fastvideo.configs.pipelines.dreamx_world import (
DreamXWorld5BCamPipelineConfig,
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_pipeline import DreamXWorldPipeline
from fastvideo.pipelines.basic.dreamx_world.stages import (
DREAMX_Y_CAMERA_KEY,
DreamXWorldCameraConditioningStage,
)
__all__ = [
"DREAMX_Y_CAMERA_KEY",
"DreamXWorld5BCamPipelineConfig",
"DreamXWorldCameraConditioningStage",
"DreamXWorldPipeline",
"make_dreamx_world_5b_cam_dit_config",
"make_dreamx_world_5b_cam_text_encoder_config",
"make_dreamx_world_5b_cam_vae_config",
]
@@ -0,0 +1,229 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import math
from dataclasses import dataclass
import numpy as np
import torch
from scipy.interpolate import interp1d
from scipy.spatial.transform import Rotation, Slerp
_ACTION_TO_MOTION = {
"w": "forward",
"a": "left",
"d": "right",
"s": "backward",
"j": "left_rot",
"l": "right_rot",
"i": "up_rot",
"k": "down_rot",
}
_TRANSLATION_BASE_UNIT = 1.0
_ROTATION_BASE_UNIT = 10.0
_INTRINSIC_ROW = [0.8, 0.5, 0.5, 0.5]
@dataclass
class DreamXCamera:
fx: float
fy: float
cx: float
cy: float
w2c_mat: np.ndarray
@property
def c2w_mat(self) -> np.ndarray:
return np.linalg.inv(self.w2c_mat)
@classmethod
def from_pose_row(cls, row: list[float]) -> DreamXCamera:
w2c_mat = np.eye(4, dtype=np.float64)
w2c_mat[:3, :] = np.asarray(row[7:], dtype=np.float64).reshape(3, 4)
return cls(
fx=float(row[1]),
fy=float(row[2]),
cx=float(row[3]),
cy=float(row[4]),
w2c_mat=w2c_mat,
)
def _translation_step(motion_type: str, current_pose: dict[str, np.ndarray], value: float, duration: int) -> np.ndarray:
if motion_type in ("forward", "backward"):
yaw = np.radians(current_pose["rotation"][1])
pitch = np.radians(current_pose["rotation"][0])
forward = np.array([-math.sin(yaw) * math.cos(pitch), math.sin(pitch), math.cos(yaw) * math.cos(pitch)])
direction = 1 if motion_type == "forward" else -1
return forward * value * direction / duration
if motion_type in ("left", "right"):
yaw = np.radians(current_pose["rotation"][1])
right = np.array([math.cos(yaw), 0.0, math.sin(yaw)])
direction = -1 if motion_type == "left" else 1
return right * value * direction / duration
return np.zeros(3)
def _rotation_step(motion_type: str, value: float, duration: int) -> np.ndarray:
if not motion_type.endswith("rot"):
return np.zeros(3)
axis = motion_type.split("_")[0]
rotation = np.zeros(3)
if axis == "left":
rotation[1] = value
elif axis == "right":
rotation[1] = -value
elif axis == "up":
rotation[0] = -value
elif axis == "down":
rotation[0] = value
return rotation / duration
def _euler_to_quaternion(angles: np.ndarray) -> list[float]:
pitch, yaw, roll = np.radians(angles)
cy = math.cos(yaw * 0.5)
sy = math.sin(yaw * 0.5)
cp = math.cos(pitch * 0.5)
sp = math.sin(pitch * 0.5)
cr = math.cos(roll * 0.5)
sr = math.sin(roll * 0.5)
return [
cy * cp * cr + sy * sp * sr,
cy * sp * cr + sy * cp * sr,
sy * cp * cr - cy * sp * sr,
cy * cp * sr - sy * sp * cr,
]
def _quaternion_to_rotation_matrix(quaternion: list[float]) -> np.ndarray:
qw, qx, qy, qz = quaternion
return np.array([
[1 - 2 * (qy**2 + qz**2), 2 * (qx * qy - qw * qz), 2 * (qx * qz + qw * qy)],
[2 * (qx * qy + qw * qz), 1 - 2 * (qx**2 + qz**2), 2 * (qy * qz - qw * qx)],
[2 * (qx * qz - qw * qy), 2 * (qy * qz + qw * qx), 1 - 2 * (qx**2 + qy**2)],
])
def _pose_rows_from_actions(action_seq: list[str], action_speed_list: list[float], duration: int) -> list[list[float]]:
if len(action_seq) != len(action_speed_list):
raise ValueError("action_seq and action_speed_list must have the same length")
positions: list[np.ndarray] = []
rotations: list[np.ndarray] = []
current_pose = {
"position": np.array([0.0, 0.0, 0.0]),
"rotation": np.array([0.0, 0.0, 0.0]),
}
for action_id, speed in zip(action_seq, action_speed_list, strict=True):
motion_types = [_ACTION_TO_MOTION[key] for key in list(action_id)]
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
)
rotation_step += _rotation_step(motion_type, float(speed) * _ROTATION_BASE_UNIT, duration)
segment_positions = []
segment_rotations = []
for index in range(1, duration + 1):
segment_positions.append(current_pose["position"] + translation_step * index)
segment_rotations.append(current_pose["rotation"] + rotation_step * index)
current_pose["position"] = segment_positions[-1].copy()
current_pose["rotation"] = segment_rotations[-1].copy()
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]]
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
extrinsic = np.hstack([rotation_matrix, translation.reshape(3, 1)])
rows.append([float(index)] + _INTRINSIC_ROW + [0.0, 0.0] + extrinsic.flatten().tolist())
return rows
def _interpolate_camera_poses(
cameras: list[DreamXCamera],
src_indices: np.ndarray,
tgt_indices: np.ndarray,
) -> list[DreamXCamera]:
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])
dets = np.linalg.det(src_rot_mat)
flip_handedness = dets.size > 0 and np.median(dets) < 0.0
if flip_handedness:
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
)
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:
quats[index] = -quats[index]
rot = Slerp(src_indices, Rotation.from_quat(quats))(tgt_indices).as_matrix()
if flip_handedness:
rot = rot @ flip_mat
ref = cameras[0]
result = []
for index in range(len(tgt_indices)):
w2c_mat = np.eye(4, dtype=np.float64)
w2c_mat[:3, :] = np.hstack([rot[index], trans[index].reshape(3, 1)])
result.append(DreamXCamera(ref.fx, ref.fy, ref.cx, ref.cy, w2c_mat))
return result
def _relative_c2w_poses(cameras: list[DreamXCamera]) -> np.ndarray:
abs_w2cs = [camera.w2c_mat for camera in cameras]
abs_c2ws = [camera.c2w_mat for camera in cameras]
target_cam_c2w = np.eye(4, dtype=np.float64)
abs2rel = target_cam_c2w @ abs_w2cs[0]
poses = [target_cam_c2w] + [abs2rel @ c2w for c2w in abs_c2ws[1:]]
return np.asarray(poses, dtype=np.float32)
def _invert_se3(transforms: torch.Tensor) -> torch.Tensor:
rotation_inv = transforms[..., :3, :3].transpose(-1, -2)
output = torch.zeros_like(transforms)
output[..., :3, :3] = rotation_inv
output[..., :3, 3] = -torch.einsum("...ij,...j->...i", rotation_inv, transforms[..., :3, 3])
output[..., 3, 3] = 1.0
return output
def build_dreamx_camera_condition(
action_seq: list[str],
action_speed_list: list[float],
*,
num_frames: int,
height: int,
width: int,
dtype: torch.dtype = torch.float32,
device: torch.device | str = "cpu",
) -> dict[str, torch.Tensor]:
del height, width # DreamX-World-5B-Cam uses fixed normalized intrinsics.
duration = math.ceil(num_frames / len(action_seq))
rows = _pose_rows_from_actions(action_seq, action_speed_list, duration)[:num_frames]
cameras = [DreamXCamera.from_pose_row(row) for row in rows]
latent_frame_count = 1 + (len(cameras) - 1) // 4
src_indices = np.arange(len(cameras), dtype=np.float64)
tgt_indices = np.linspace(0, len(cameras) - 1, latent_frame_count)
cameras = _interpolate_camera_poses(cameras, src_indices, tgt_indices)
c2ws = torch.as_tensor(_relative_c2w_poses(cameras), dtype=dtype, device=device)
viewmats = _invert_se3(c2ws)
intrinsics = torch.zeros((latent_frame_count, 3, 3), dtype=dtype, device=device)
intrinsics[:, 0, 0] = 969.6969696969696 / (960.0 * 2)
intrinsics[:, 1, 1] = 969.6969696969696 / (540.0 * 2)
intrinsics[:, 2, 2] = 1.0
return {"viewmats": viewmats, "K": intrinsics}
@@ -0,0 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
"""Compatibility exports for DreamX-World pipeline configs."""
from fastvideo.configs.pipelines.dreamx_world import (
DreamXWorld5BCamPipelineConfig,
make_dreamx_world_5b_cam_dit_config,
make_dreamx_world_5b_cam_text_encoder_config,
make_dreamx_world_5b_cam_vae_config,
)
__all__ = [
"DreamXWorld5BCamPipelineConfig",
"make_dreamx_world_5b_cam_dit_config",
"make_dreamx_world_5b_cam_text_encoder_config",
"make_dreamx_world_5b_cam_vae_config",
]
@@ -0,0 +1,78 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World video pipeline entrypoint."""
from fastvideo.api.sampling_param import SamplingParam
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.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
from fastvideo.pipelines.basic.dreamx_world.stages import DreamXWorldCameraConditioningStage
from fastvideo.pipelines.stages import (
ConditioningStage,
DecodingStage,
DenoisingStage,
InputValidationStage,
LatentPreparationStage,
TextEncodingStage,
TimestepPreparationStage,
)
logger = init_logger(__name__)
class DreamXWorldPipeline(LoRAPipeline, ComposedPipelineBase):
"""DreamX-World-5B-Cam pipeline with native FastVideo camera conditioning."""
_required_config_modules = ["text_encoder", "tokenizer", "vae", "transformer", "scheduler"]
pipeline_config_cls = DreamXWorld5BCamPipelineConfig
sampling_params_cls = SamplingParam
def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler(shift=fastvideo_args.pipeline_config.flow_shift)
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=DenoisingStage(
transformer=self.get_module("transformer"),
transformer_2=self.get_module("transformer_2", None),
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("DreamXWorldPipeline initialized with native camera conditioning")
EntryClass = DreamXWorldPipeline
@@ -0,0 +1,42 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World model family pipeline presets."""
from fastvideo.api.presets import InferencePreset, PresetStageSpec
_NEGATIVE_PROMPT_CN = (
"色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,"
"静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,"
"多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,"
"形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,"
"背景人很多,倒着走"
)
_DENOISE_STAGE = PresetStageSpec(
name="denoise",
kind="denoising",
description="DreamX-World camera-conditioned denoising pass",
allowed_overrides=frozenset({
"num_inference_steps",
"guidance_scale",
}),
)
DREAMX_WORLD_5B_CAM = InferencePreset(
name="dreamx_world_5b_cam",
version=1,
model_family="dreamx_world",
description="DreamX-World 5B camera-control video generation",
workload_type="i2v",
stage_schemas=(_DENOISE_STAGE, ),
defaults={
"height": 480,
"width": 832,
"num_frames": 161,
"fps": 16,
"guidance_scale": 5.0,
"num_inference_steps": 30,
"negative_prompt": _NEGATIVE_PROMPT_CN,
},
)
ALL_PRESETS = (DREAMX_WORLD_5B_CAM, )
@@ -0,0 +1,78 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World pipeline stages."""
from __future__ import annotations
import torch
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
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,
)
DREAMX_Y_CAMERA_KEY = "dreamx_y_camera"
class DreamXWorldCameraConditioningStage(PipelineStage):
"""Build PRoPE camera conditioning for DreamX-World-5B-Cam."""
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
del fastvideo_args
if DREAMX_Y_CAMERA_KEY in batch.extra:
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)
if action_seq is None:
action_seq = ["w"]
if action_speed_list is None:
action_speed_list = [4]
if isinstance(action_seq, str):
action_seq = [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
width = int(batch.width) if batch.width is not None else 1280
num_frames = int(batch.num_frames)
dtype = batch.latents.dtype if torch.is_tensor(batch.latents) else torch.float32
device = batch.latents.device if torch.is_tensor(batch.latents) else "cpu"
y_camera = build_dreamx_camera_condition(
list(action_seq),
action_speed_list,
num_frames=num_frames,
height=height,
width=width,
dtype=dtype,
device=device,
)
batch.extra[DREAMX_Y_CAMERA_KEY] = {
key: value.unsqueeze(0)
for key, value in y_camera.items()
}
return batch
def verify_output(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> VerificationResult:
del fastvideo_args
result = VerificationResult()
y_camera = batch.extra.get(DREAMX_Y_CAMERA_KEY)
result.add_check("dreamx_y_camera", y_camera, lambda value: isinstance(value, dict))
if isinstance(y_camera, dict):
result.add_check("dreamx_y_camera.viewmats", y_camera.get("viewmats"), torch.is_tensor)
result.add_check("dreamx_y_camera.K", y_camera.get("K"), torch.is_tensor)
return result
+16
View File
@@ -190,6 +190,19 @@ 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
for key, value in dreamx_y_camera.items()
}
dreamx_camera_kwargs = self.prepare_extra_func_kwargs(
self.transformer.forward,
{
"y_camera": dreamx_y_camera,
},
)
for key in ("flux2_txt_ids", "flux2_img_ids"):
value = batch.extra.get(key)
if torch.is_tensor(value):
@@ -241,6 +254,7 @@ 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"
self.vae = self.vae.to(local_device)
z = self.vae.encode(batch.pil_image).mean.float()
if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
if isinstance(self.vae.shift_factor, torch.Tensor):
@@ -494,6 +508,7 @@ class DenoisingStage(PipelineStage):
**pos_cond_kwargs,
**action_kwargs,
**camera_kwargs,
**dreamx_camera_kwargs,
**timesteps_r_kwarg,
**flux2_id_kwargs,
)
@@ -536,6 +551,7 @@ class DenoisingStage(PipelineStage):
**neg_cond_kwargs,
**action_kwargs,
**camera_kwargs,
**dreamx_camera_kwargs,
**timesteps_r_kwarg,
**flux2_id_kwargs,
)
+16
View File
@@ -20,6 +20,7 @@ from fastvideo.configs.pipelines.cosmos2_5 import (
Cosmos25Config,
Cosmos25_14BConfig,
)
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
from fastvideo.configs.pipelines.hunyuan import FastHunyuanConfig, HunyuanConfig
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
from fastvideo.configs.pipelines.gen3c import Gen3CConfig
@@ -773,6 +774,18 @@ def _register_configs() -> None:
model_family="wan",
default_preset="wan_2_2_ti2v_5b",
)
register_configs(
sampling_param_cls=None,
pipeline_config_cls=DreamXWorld5BCamPipelineConfig,
workload_types=(WorkloadType.I2V, ),
hf_model_paths=[
"GD-ML/DreamX-World-5B-Cam",
],
model_detectors=[lambda path: "dreamx-world" in path.lower() or "dreamxworldpipeline" in path.lower()],
model_family="dreamx_world",
default_preset="dreamx_world_5b_cam",
)
register_configs(
sampling_param_cls=None,
pipeline_config_cls=FastWan2_2_TI2V_5B_Config,
@@ -951,6 +964,8 @@ def _register_presets() -> None:
from fastvideo.api.presets import register_preset
from fastvideo.pipelines.basic.cosmos.presets import (
ALL_PRESETS as COSMOS_PRESETS, )
from fastvideo.pipelines.basic.dreamx_world.presets import (
ALL_PRESETS as DREAMX_WORLD_PRESETS, )
from fastvideo.pipelines.basic.gamecraft.presets import (
ALL_PRESETS as GAMECRAFT_PRESETS, )
from fastvideo.pipelines.basic.gen3c.presets import (
@@ -984,6 +999,7 @@ def _register_presets() -> None:
all_preset_groups = (
COSMOS_PRESETS,
DREAMX_WORLD_PRESETS,
FLUX2_PRESETS,
GAMECRAFT_PRESETS,
GEN3C_PRESETS,
@@ -0,0 +1,122 @@
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations
import os
from pathlib import Path
import pytest
from PIL import Image, ImageDraw
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,
)
logger = init_logger(__name__)
REQUIRED_GPUS = 1
device_reference_folder = resolve_inference_device_reference_folder(logger)
_LOCAL_CONVERTED_MODEL = Path("converted_weights/dreamx_world")
_MODEL_PATH = os.getenv(
"DREAMX_WORLD_SSIM_MODEL_PATH",
str(_LOCAL_CONVERTED_MODEL),
)
DREAMX_WORLD_PARAMS = {
"num_gpus": 1,
"model_path": _MODEL_PATH,
"height": 64,
"width": 64,
"num_frames": 9,
"num_inference_steps": 1,
"guidance_scale": 1.0,
"seed": 1024,
"fps": 16,
}
DREAMX_WORLD_FULL_QUALITY_PARAMS = {
**DREAMX_WORLD_PARAMS,
"height": 480,
"width": 832,
"num_frames": 161,
"num_inference_steps": 30,
"guidance_scale": 5.0,
}
DREAMX_WORLD_MODEL_TO_PARAMS = {
"DreamX-World-5B-Cam": DREAMX_WORLD_PARAMS,
}
FULL_QUALITY_DREAMX_WORLD_MODEL_TO_PARAMS = {
"DreamX-World-5B-Cam": DREAMX_WORLD_FULL_QUALITY_PARAMS,
}
DREAMX_WORLD_TEST_CASES = [
(
"A cinematic first-person drive through a futuristic coastal city at sunrise, "
"reflective glass towers, clean streets, soft volumetric light.",
("w", "d", "w"),
(4.0, 2.0, 4.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))
draw = ImageDraw.Draw(image)
draw.rectangle((0, 56, 96, 96), fill=(36, 44, 52))
draw.polygon([(0, 56), (48, 30), (96, 56)], fill=(120, 142, 158))
draw.rectangle((34, 42, 62, 72), fill=(178, 198, 212))
draw.line((0, 80, 96, 66), fill=(238, 209, 124), width=3)
image.save(path)
@pytest.mark.parametrize(("prompt", "action_list", "action_speed_list"), DREAMX_WORLD_TEST_CASES)
@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
@pytest.mark.parametrize("model_id", list(DREAMX_WORLD_MODEL_TO_PARAMS.keys()))
def test_dreamx_world_inference_similarity(
prompt: str,
action_list: tuple[str, ...],
action_speed_list: tuple[float, ...],
attention_backend_name: str,
model_id: str,
tmp_path: Path,
) -> None:
model_path = Path(str(DREAMX_WORLD_MODEL_TO_PARAMS[model_id]["model_path"]))
if not model_path.exists():
pytest.skip(
f"DreamX-World converted model path is missing: {model_path}. "
"Set DREAMX_WORLD_SSIM_MODEL_PATH to a FastVideo-loadable converted root."
)
image_path = tmp_path / "dreamx_world_ssim_input.png"
_write_deterministic_reference_image(image_path)
run_image_to_video_similarity_test(
logger=logger,
script_dir=os.path.dirname(os.path.abspath(__file__)),
device_reference_folder=device_reference_folder,
prompt=prompt,
image_path=str(image_path),
attention_backend_name=attention_backend_name,
model_id=model_id,
default_params_map=DREAMX_WORLD_MODEL_TO_PARAMS,
full_quality_params_map=FULL_QUALITY_DREAMX_WORLD_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": "DreamXWorldPipeline",
},
generation_kwargs_override={
"action_list": list(action_list),
"action_speed_list": list(action_speed_list),
},
)
@@ -0,0 +1,219 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Convert DreamX-World-5B-Cam raw transformer weights to FastVideo-loadable format.
The GD-ML/DreamX-World-5B-Cam repository stores the transformer as raw
DreamX/Wan official shards. FastVideo's TransformerLoader expects a Diffusers-like
transformer folder with a config.json and safetensors whose keys can be mapped by
WanVideoConfig.param_names_mapping. This script performs the raw official ->
Diffusers-like key rename and writes the DreamX 5B-Cam transformer config.
Example:
python scripts/checkpoint_conversion/dreamx_world_to_diffusers.py \
--source official_weights/dreamx_world \
--output converted_weights/dreamx_world
"""
from __future__ import annotations
import argparse
import json
import re
import shutil
from collections import OrderedDict
from pathlib import Path
import torch
from huggingface_hub import save_torch_state_dict
from safetensors import safe_open
from safetensors.torch import load_file
OFFICIAL_TO_DIFFUSERS_MAPPING: dict[str, str] = {
r"^text_embedding\.0\.(.*)$": r"condition_embedder.text_embedder.linear_1.\1",
r"^text_embedding\.2\.(.*)$": r"condition_embedder.text_embedder.linear_2.\1",
r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.linear_1.\1",
r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.linear_2.\1",
r"^time_projection\.1\.(.*)$": r"condition_embedder.time_proj.\1",
r"^img_emb\.proj\.0\.(.*)$": r"condition_embedder.image_embedder.norm1.\1",
r"^img_emb\.proj\.1\.(.*)$": r"condition_embedder.image_embedder.ff.net.0.proj.\1",
r"^img_emb\.proj\.3\.(.*)$": r"condition_embedder.image_embedder.ff.net.2.\1",
r"^img_emb\.proj\.4\.(.*)$": r"condition_embedder.image_embedder.norm2.\1",
r"^head\.modulation": r"scale_shift_table",
r"^head\.head\.(.*)$": r"proj_out.\1",
r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.attn1.to_q.\2",
r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.attn1.to_k.\2",
r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.attn1.to_v.\2",
r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.attn1.to_out.0.\2",
r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.attn1.norm_q.\2",
r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.attn1.norm_k.\2",
r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2",
r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2",
r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.0.\2",
r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
r"^blocks\.(\d+)\.cross_attn\.norm_k_img\.(.*)$": r"blocks.\1.attn2.norm_added_k.\2",
r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.net.0.proj.\2",
r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.net.2.\2",
r"^blocks\.(\d+)\.modulation": r"blocks.\1.scale_shift_table",
r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.norm2.\2",
}
TRANSFORMER_CONFIG: dict[str, object] = {
"_class_name": "DreamXWorldTransformer3DModel",
"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,
"cross_attn_norm": True,
"qk_norm": "rms_norm_across_heads",
"eps": 1e-6,
"image_dim": None,
"added_kv_proj_dim": None,
"rope_max_seq_len": 1024,
"add_control_adapter": True,
"cam_method": "prope",
"attn_compress": 1,
"cam_self_attn_layers": None,
}
MODEL_INDEX: dict[str, object] = {
"_class_name": "DreamXWorldPipeline",
"_diffusers_version": "0.31.0",
"scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
"text_encoder": ["transformers", "UMT5EncoderModel"],
"tokenizer": ["transformers", "AutoTokenizer"],
"transformer": ["diffusers", "DreamXWorldTransformer3DModel"],
"vae": ["diffusers", "AutoencoderKLWan"],
}
REUSED_COMPONENTS = ("scheduler", "text_encoder", "tokenizer", "vae")
def map_transformer_key(key: str) -> str:
for pattern, replacement in OFFICIAL_TO_DIFFUSERS_MAPPING.items():
if re.match(pattern, key):
return re.sub(pattern, replacement, key)
return key
def _safetensor_files(source: Path) -> list[Path]:
if source.is_file():
if source.suffix != ".safetensors":
raise ValueError(f"Only .safetensors files are supported, got {source}")
return [source]
index_path = source / "diffusion_pytorch_model.safetensors.index.json"
if index_path.exists():
index = json.loads(index_path.read_text())
return sorted({source / shard for shard in index["weight_map"].values()})
files = sorted(source.glob("*.safetensors"))
if not files:
raise FileNotFoundError(f"No safetensors files found under {source}")
return files
def convert_transformer(source: Path, output: Path, max_shard_size: str) -> None:
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")
converted: OrderedDict[str, torch.Tensor] = OrderedDict()
for shard in _safetensor_files(source):
print(f"loading {shard}")
for key, tensor in load_file(shard, device="cpu").items():
new_key = map_transformer_key(key)
if new_key in converted:
raise ValueError(f"Duplicate converted key: {new_key}")
converted[new_key] = tensor
print(f"saving {len(converted)} tensors to {transformer_dir}")
save_torch_state_dict(converted, transformer_dir, max_shard_size=max_shard_size)
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():
print(f"keeping existing {dst}")
return
if symlink:
dst.symlink_to(src.resolve(), target_is_directory=src.is_dir())
print(f"linked {dst} -> {src}")
elif src.is_dir():
shutil.copytree(src, dst)
print(f"copied {src} -> {dst}")
else:
shutil.copy2(src, dst)
print(f"copied {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))
print("pass --component-source <Wan2.2 Diffusers root> to copy or link reused components")
return
(output / "model_index.json").write_text(json.dumps(MODEL_INDEX, indent=2) + "\n")
print(f"wrote {output / 'model_index.json'}")
def analyze(source: Path) -> None:
total = 0
unchanged = 0
examples: list[tuple[str, str]] = []
for shard in _safetensor_files(source):
with safe_open(shard, framework="pt", device="cpu") as tensors:
for key in tensors:
total += 1
new_key = map_transformer_key(key)
unchanged += int(new_key == key)
if len(examples) < 20 and new_key != key:
examples.append((key, new_key))
print(f"total_keys={total} unchanged_keys={unchanged}")
for old, new in examples:
print(f"{old} -> {new}")
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--source", type=Path, required=True, help="DreamX raw transformer directory or safetensors file")
parser.add_argument("--output", type=Path, required=True, help="Output model root; transformer/ is created inside it")
parser.add_argument("--max-shard-size", default="10GB")
parser.add_argument(
"--component-source",
type=Path,
help="Optional Wan2.2 Diffusers root whose scheduler/text_encoder/tokenizer/vae components are reused.",
)
parser.add_argument(
"--symlink-components",
action="store_true",
help="Symlink reused components from --component-source instead of copying them.",
)
parser.add_argument("--analyze", action="store_true", help="Only print key mapping summary")
args = parser.parse_args()
if args.analyze:
analyze(args.source)
else:
convert_transformer(args.source, args.output, args.max_shard_size)
write_model_index(args.output, args.component_source, args.symlink_components)
if __name__ == "__main__":
main()
@@ -0,0 +1,162 @@
# DreamX World Port Status
## Summary
- model_family: `dreamx_world`
- workload_types: `I2V camera-control compatibility shim`
- official_ref: `https://github.com/AMAP-ML/DreamX-World`
- official_ref_dir: `DreamX-World/`
- hf_weights_path: `GD-ML/DreamX-World-5B-Cam`
- local_weights_dir: `official_weights/dreamx_world`
- source_layout: `raw_official`
- local_tests_readme: `tests/local_tests/dreamx_world/README.md`
## Current Phase
- phase: `phase_11_post_parity_handoff`
- status: `complete`
- owner: `orchestrator`
- last_updated: `2026-07-01`
## Component Matrix
| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues |
|---|---|---|---|---|---|---|---|---|---|
| transformer | dit | ported_dedicated | `DreamX-World/models/wan_transformer3d.py`; PRoPE helpers in `DreamX-World/models/prope_utils.py` | `DreamX-World/inference_dreamx5b.py::setup_models`, `Wan2_2Transformer3DModel.from_pretrained(... cam_method=prope, add_control_adapter=True)` | `fastvideo/models/dits/dreamx_world.py`; `fastvideo/configs/models/dits/dreamx_world.py`; DreamX pipeline config helper | native_prope_pass | real_conversion_pass | strict_load_and_forward_parity_pass | none |
| vae | vae | reuse_pending | `DreamX-World/models/wan_vae3_8.py` | `DreamX-World/inference_dreamx5b.py::setup_models`, `AutoencoderKLWan3_8.from_pretrained(Wan2.2_VAE.pth)` | `fastvideo/models/vaes/wanvae.py`; DreamX VAE config helper | config_smoke_pass | raw_key_mapping_pass | encode_parity_pass | none |
| text_encoder/tokenizer | encoder | reuse_pending | `DreamX-World/models/wan_text_encoder.py`; tokenizer via Wan2.2 base model | `DreamX-World/inference_dreamx5b.py::setup_models`, `WanT5EncoderModel` + tokenizer subpaths | `fastvideo/models/encoders/t5.py::UMT5EncoderModel`; DreamX UMT5 config helper | config_smoke_pass | staged_weight_load_pass | hidden_state_parity_pass | none |
| 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 |
## Conversion State
- conversion_script: `scripts/checkpoint_conversion/dreamx_world_to_diffusers.py`
- converted_weights_dir: `converted_weights/dreamx_world`
- source_layout: `raw_official`
- strict_load_status: `pass`
- conversion_script_status: `transformer_model_index_and_config_consistency_smoke_pass`
- model_index_status: `smoke_pass`
- passthrough_components: `Wan2.2 Diffusers scheduler, tokenizer, and text encoder are symlinked from official_weights/Wan2.2-TI2V-5B-Diffusers; VAE parity uses raw Wan2.2_VAE.pth with an explicit DreamX raw-to-FastVideo key mapper because official encode returns normalized latents.`
- retry_history: `none`
## Parity Commands
| Scope | Command | Last Result | Notes |
|---|---|---|---|
| transformer | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py -v -s` | strict_load_and_forward_parity_pass | 2026-07-01: converted real 5B-Cam transformer shards strict-load into dedicated `DreamXWorldTransformer3DModel` with 0 shape mismatches; official-vs-FastVideo small-input fp32 forward parity passes on CUDA (`diff_max=0.072533`, `diff_mean=0.008014`). |
| vae | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py -v -s` | encode_parity_pass | 2026-06-30: official DreamX raw `Wan2.2_VAE.pth` maps 196/196 keys into FastVideo Wan VAE and encode parity passes after applying the same official latent normalization (`(mu - mean) / std`). |
| text_encoder | `python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py -v -s` | hidden_state_parity_pass | 2026-06-30: official `WanT5EncoderModel` vs FastVideo `UMT5EncoderModel` hidden-state parity passes on CUDA using staged Wan2.2 text encoder/tokenizer weights and reference-only `xfuser` stubs. |
| 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`. |
| 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. |
## Open Questions
| ID | Question | Owner | Needed By Phase | Status | Resolution |
|---|---|---|---|---|---|
| Q001 | Should first PR expose only the `DreamX-World-5B-Cam` 5s camera-control mode and exclude AR long-horizon forcing? | user | prep | resolved | User approved starting with `DreamX-World-5B-Cam`; AR long-horizon is out of first-PR scope. |
| Q002 | Does FastVideo's existing Wan2.2 TI2V transformer support DreamX PRoPE/control adapter with a small extension, or is a DreamX-specific DiT required? | component:transformer | Phase 3 | resolved | Project guidance prefers a separate DreamX DiT for maintainability. DreamX PRoPE/control adapter now lives in `fastvideo/models/dits/dreamx_world.py`; Wan DiT/config have no DreamX-specific fields or `y_camera` signature. |
| 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. |
## Issues And Blockers
| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution |
|---|---|---|---|---|---|---|---|---|
| I001 | prep | official_env | medium | Official import initially failed because `xfuser` was missing. | `ModuleNotFoundError: No module named 'xfuser'` from `python -c "import sys; sys.path.insert(0, 'DreamX-World'); import inference_dreamx5b"` | prep | resolved | Installed `xfuser==0.4.1`; import progressed. |
| I002 | prep | official_env | medium | Official import then failed because GUI OpenCV required missing system `libxcb.so.1`. | `ImportError: libxcb.so.1: cannot open shared object file` through `cv2` import in Diffusers ConsisID path. | prep | resolved | Installed `opencv-python-headless`; `import inference_dreamx5b` passed. |
| I003 | prep | weights | medium | HF repo has raw official transformer shards and no Diffusers `model_index.json`. | `inspect_hf_layout.py GD-ML/DreamX-World-5B-Cam --json` returned `source_layout=raw_official`, `needs_conversion=yes`, `model_index_class=null`. | conversion | resolved | Downloaded raw DreamX shards to `official_weights/dreamx_world`; converted transformer to `converted_weights/dreamx_world/transformer`; symlinked reusable Wan2.2 Diffusers components; real 5B transformer strict-load passes. |
| I004 | prep | dependencies | high | Official reference import required extra packages in the local environment, but FastVideo native runtime should not inherit those dependencies. | `xfuser==0.4.1` and `opencv-python-headless` were installed only to make `DreamX-World/inference_dreamx5b.py` import for reference/parity. | pipeline | resolved | Production DreamX FastVideo code uses native camera/image/video utilities and has no runtime `xfuser` or OpenCV import requirement; those packages remain reference-only local parity dependencies. |
| 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. |
## Escape Hatches
| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution |
|---|---|---|---|---|---|---|
## Decisions
| Date | Decision | Rationale | Impact |
|---|---|---|---|
| 2026-06-29 | First PR scope is `DreamX-World-5B-Cam` only. | Cam mode is closest to existing Wan2.2 TI2V support; AR forcing needs separate causal/KV pipeline work. | Component inventory and parity focus on `inference_dreamx5b.py` and `pipeline_dreamxworld.py`. |
| 2026-06-29 | Do not install full DreamX requirements during prep. | Full requirements pin core FastVideo stack packages. | Installed only `xfuser==0.4.1` and `opencv-python-headless` to make official imports work. |
| 2026-06-29 | Treat HF DreamX-World-5B-Cam weights as raw official transformer layout requiring conversion. | HF inspection found no `model_index.json`. | Phase 5 must create `scripts/checkpoint_conversion/dreamx_world_to_diffusers.py` after component prototype/key dumps. |
| 2026-06-29 | Do not add DreamX reference-only dependencies to FastVideo production requirements. | The current environment should remain the FastVideo environment; extra packages are only for official reference parity. | Native DreamX integration must avoid runtime `xfuser` and OpenCV requirements unless explicitly approved later. |
| 2026-06-29 | Implement DreamX camera conditioning as native FastVideo utility. | It is weightless and removes the need to import DreamX reference utilities at production runtime. | `fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py` now has non-skip parity against official action-to-PRoPE tensors. |
| 2026-06-29 | First PR supports DreamX default `Flow` sampler only. | FastVideo FlowMatch Euler scheduler matches the official Diffusers scheduler for DreamX defaults. | Pipeline work can use FastVideo native FlowMatch scheduler; UniPC and DPM++ are out of first-PR scope. |
| 2026-07-01 | Keep DreamX PRoPE/control adapter in a dedicated DreamX DiT class. | Project guidance is that putting too much DreamX behavior into Wan makes the model hard to manage. | `fastvideo/models/dits/dreamx_world.py` defines `DreamXWorldTransformer3DModel`, `DreamXWorldTransformerBlock`, and `DreamXPropeSelfAttention`; `fastvideo/configs/models/dits/dreamx_world.py` owns DreamX adapter config fields; Wan DiT/config are unchanged from DreamX. |
| 2026-06-30 | Camera parity test loads official camera functions by file instead of importing the official `utils` package. | Official package initialization pulls unrelated dependencies that can require GUI OpenCV system libraries. | Camera parity remains non-skip without adding DreamX reference-only dependencies to FastVideo production requirements. |
| 2026-06-30 | Add DreamX-World-5B-Cam model and pipeline config helpers plus a conversion script. | Official HF DreamX 5B-Cam transformer config is 30 layers, hidden size 3072, 24 heads, 48 latent channels, plus Wan2.2 48-channel VAE and UMT5-XXL text encoder. | DreamX helpers wire DiT/VAE/UMT5/Flow/TI2V settings; `dreamx_world_to_diffusers.py` writes a FastVideo-loadable transformer config plus renamed safetensors; strict-load smoke passes on a tiny official DreamX transformer and the real 5B converted shards. |
| 2026-06-30 | Pass DreamX camera PRoPE condition through the FastVideo batch/denoising path. | DreamX transformer expects `y_camera={"viewmats", "K"}` at denoising time. | `DreamXWorldPipeline` is registered as a basic pipeline entry and initializes the official default FlowMatch scheduler; `dreamx_world_5b_cam` preset mirrors official 5B-Cam defaults; `DreamXWorldCameraConditioningStage` writes `batch.extra["dreamx_y_camera"]`; generic denoising filters and forwards it as `y_camera` only for compatible transformers. |
## Handoff Notes
- Prep, component parity, pipeline smoke/parity, and the basic example validation are complete for `DreamX-World-5B-Cam`.
- Official reference clone is staged at `DreamX-World/` and ignored by git.
- Workspace-local weights are staged: DreamX raw transformer shards under `official_weights/dreamx_world`, Wan2.2 raw base artifacts under `official_weights/Wan2.2-TI2V-5B`, and Wan2.2 Diffusers reusable components under `official_weights/Wan2.2-TI2V-5B-Diffusers`.
- Camera conditioning parity is active and passing without weights.
- Default Flow scheduler parity is active and passing without weights.
- Transformer has corrected official 5B-Cam architecture in dedicated DreamX DiT/config files, native PRoPE/control-adapter, conversion mapping, real converted 5B strict-load, and official-vs-FastVideo forward parity passing on CUDA. VAE encode parity and text hidden-state parity pass on CUDA. Pipeline entry/registry, local model_info resolution, preset, config, FlowMatch scheduler init, camera stage, denoising `y_camera` kwarg smokes, independent CUDA pipeline smoke/parity, and a small saved-video basic example pass.
- Full local DreamX component suite is non-skip PASS: `python -m pytest tests/local_tests/dreamx_world/ -q -rs` returned `26 passed` on 2026-07-01. Pipeline smoke/parity and SSIM quality regression are also non-skip PASS locally.
- Keep `xfuser` and OpenCV as reference-only parity dependencies. Do not add them
to FastVideo requirements or production imports.
## Quality Regression
- status: `added`
- test: `fastvideo/tests/ssim/test_dreamx_world_similarity.py`
- 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.
## Final Handoff
```text
final_handoff:
prep_handoff_complete: yes
conversion_status: pass
components:
- name: transformer
reuse_or_port: ported_dedicated_dit
parity_test: tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py
parity_status: non_skip_pass
concerns_or_unknowns: none
- name: vae
reuse_or_port: reused
parity_test: tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py
parity_status: non_skip_pass
concerns_or_unknowns: none
- name: text_encoder_tokenizer
reuse_or_port: reused
parity_test: tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py
parity_status: non_skip_pass
concerns_or_unknowns: none
- name: scheduler
reuse_or_port: reused
parity_test: tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py
parity_status: non_skip_pass
concerns_or_unknowns: none
- name: camera_conditioning
reuse_or_port: ported
parity_test: tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py
parity_status: non_skip_pass
concerns_or_unknowns: none
pipeline_smoke: pass
pipeline_parity: pass
example_status: pass
quality_regression: added
local_tests_readme: tests/local_tests/dreamx_world/README.md
port_state_file: tests/local_tests/dreamx_world/PORT_STATUS.md
token_values_committed: no
runtime_third_party_model_imports: none
blockers: none
escape_hatch: none
```
+194
View File
@@ -0,0 +1,194 @@
# DreamX World Local Tests
Local-only parity and smoke tests for the `dreamx_world` FastVideo port. These
tests compare FastVideo against the official DreamX-World reference
implementation and are not expected to run in CI unless explicitly promoted
later.
Port progress, open questions, issues, and handoff notes live in
`tests/local_tests/dreamx_world/PORT_STATUS.md`.
## Reference Assets
| 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 |
| 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/` |
| Official commit/version | `221875811ba31f7eac6c3025b215c09ad2cefd1d` |
| HF weights | `GD-ML/DreamX-World-5B-Cam` |
| HF revision | default |
| Local weights dir | `official_weights/dreamx_world` |
| Source layout | `raw_official` |
| Needs conversion | `yes` |
Do not write token values in this file. Current token env var detected during
prep: `none`.
## Shared Environment Setup
Run from the FastVideo repo root in the same conda/env used for FastVideo. Do
not create a separate upstream environment for parity tests.
```bash
python ".agents/skills/add-model-01-prep/scripts/clone_reference_repo.py" \
"https://github.com/AMAP-ML/DreamX-World.git" \
"DreamX-World" \
--commit "221875811ba31f7eac6c3025b215c09ad2cefd1d" \
--update-gitignore
```
DreamX-World does not expose a packaging file for editable install. During prep
the official import check used `sys.path.insert(0, "DreamX-World")`.
Additional official deps installed into the current environment for imports:
```bash
uv pip install xfuser==0.4.1
uv pip install opencv-python-headless
```
These packages are for running the official DreamX reference during local
parity only. They must not become FastVideo production/runtime dependencies for
the native `dreamx_world` pipeline.
Do not install the full `DreamX-World/requirements.txt` without explicit
approval. It pins core FastVideo stack packages including `torch`, `torchvision`,
`triton`, `flash_attn`, and `diffusers`.
## Official Environment Status
```text
dependency_changes: installed official deps in current env
official_env_status: imports_ok
private_dep_stubs: none
blocked_on: none
```
Import check used during prep:
```bash
python -c "import sys; sys.path.insert(0, 'DreamX-World'); import inference_dreamx5b; print('imports_ok')"
```
## Weight Setup
HF layout inspection found no root `model_index.json`; the repo contains
`config.json`, a safetensors index, and three transformer safetensors shards.
This is a raw official transformer layout and requires conversion before
FastVideo can load it through `VideoGenerator.from_pretrained`.
```bash
python ".agents/skills/add-model-01-prep/scripts/inspect_hf_layout.py" \
"GD-ML/DreamX-World-5B-Cam" \
--json
```
Weights have been staged workspace-locally. The raw DreamX transformer repo lives at `official_weights/dreamx_world`; Wan2.2 raw base artifacts live at `official_weights/Wan2.2-TI2V-5B`; Wan2.2 Diffusers reusable components live at `official_weights/Wan2.2-TI2V-5B-Diffusers`. To reproduce the DreamX download:
```bash
python ".agents/skills/add-model-01-prep/scripts/download_hf_weights.py" \
"GD-ML/DreamX-World-5B-Cam" \
"official_weights/dreamx_world"
```
## Prototype And Conversion Artifacts
State-dict key/shape dumps are generated after FastVideo native prototypes exist
and are used to build the conversion mapping.
```text
official_key_dumps:
transformer: converted_weights/dreamx_world/_mapping/transformer_official_keys.json
fastvideo_key_dumps:
transformer: converted_weights/dreamx_world/_mapping/transformer_fastvideo_keys.json
conversion_script: scripts/checkpoint_conversion/dreamx_world_to_diffusers.py
conversion_script_status: transformer_model_index_and_config_consistency_smoke_pass
conversion_source_layout: raw_official
converted_weights_dir: converted_weights/dreamx_world
model_index_status: smoke_pass
strict_load_status: pass
```
The converter writes `transformer/` from raw DreamX shards. To create a full
FastVideo-loadable diffusers-style root, pass a Wan2.2 Diffusers directory as the
component source so reusable components are copied or symlinked before
`model_index.json` is emitted:
```bash
python scripts/checkpoint_conversion/dreamx_world_to_diffusers.py \
--source official_weights/dreamx_world \
--output converted_weights/dreamx_world \
--component-source /path/to/Wan2.2-TI2V-5B-Diffusers \
--symlink-components
```
## Expected Parity Tests
Planned local tests for this family:
| Component | Official files / args | Test | Concerns | Status |
|---|---|---|---|---|
| transformer | `DreamX-World/models/wan_transformer3d.py`; instantiated in `DreamX-World/inference_dreamx5b.py` with `cam_method=prope`, `add_control_adapter=True` | `tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py` | FastVideo uses dedicated `DreamXWorldTransformer3DModel`/`DreamXWorldConfig` files for DreamX 5B-Cam config, PRoPE, conversion mapping, real converted 5B strict-load PASS, and official-vs-FastVideo small-input forward parity PASS on CUDA. Wan DiT/config have no DreamX-specific adapter fields. | strict_load_and_forward_parity_pass |
| vae | `DreamX-World/models/wan_vae3_8.py`; `vae_type=AutoencoderKLWan3_8`, `vae_subpath=Wan2.2_VAE.pth` | `tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py` | DreamX raw `Wan2.2_VAE.pth` maps 196/196 keys into FastVideo Wan VAE; encode parity passes after applying official latent normalization. | encode_parity_pass |
| text_encoder/tokenizer | `DreamX-World/models/wan_text_encoder.py`; T5 path from Wan2.2 base model | `tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py` | Official `WanT5EncoderModel` vs FastVideo `UMT5EncoderModel` hidden-state parity passes on CUDA with staged Wan2.2 weights/tokenizer. | hidden_state_parity_pass |
| 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 |
Include reused components in parity. Reuse is accepted only after the FastVideo
component definition and official instantiation arguments have both been checked
and the component parity test passes non-skip.
Run the relevant tests with:
```bash
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py -v -s
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py -v -s
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py -v -s
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py -v -s
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py -v -s
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py -v -s
python -m pytest tests/local_tests/dreamx_world/test_dreamx_world_conversion.py -v -s
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
```
## Current Local Results
```bash
python -m pytest tests/local_tests/dreamx_world/ -v -s
# 2026-07-01: 26 passed, 0 skipped
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
# 2026-07-01: smoke image-path TI2V coverage passed separately with 3 passed, 0 skipped
PYTHONPATH=/workspace/FastVideo FASTVIDEO_SSIM_MODEL_ID=DreamX-World-5B-Cam python -m pytest fastvideo/tests/ssim/test_dreamx_world_similarity.py -q -rs
# 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.
## 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.
- 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
reference-only imports inside local parity tests.
- Review agents should verify setup commands still match the PR, then run the
listed parity tests or report the exact blocker.
## 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.
@@ -0,0 +1 @@
# SPDX-License-Identifier: Apache-2.0
@@ -0,0 +1,95 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World camera-conditioning parity against the official reference.
Coverage scope: implementation_subcomponent. This verifies the weightless
action-sequence to PRoPE camera tensor path used by DreamX-World-5B-Cam.
"""
from __future__ import annotations
import os
from pathlib import Path
import pytest
import torch
from torch.testing import assert_close
from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import build_dreamx_camera_condition
REPO_ROOT = Path(__file__).resolve().parents[3]
OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World"))
PARITY_SCOPE = "implementation_subcomponent"
def _load_official_functions():
if not OFFICIAL_REF_DIR.exists():
pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}")
try:
import importlib.util
pose_path = OFFICIAL_REF_DIR / "utils" / "pose_utils.py"
pose_spec = importlib.util.spec_from_file_location(
"dreamx_world_pose_utils", pose_path)
if pose_spec is None or pose_spec.loader is None:
raise RuntimeError(f"Cannot load DreamX pose_utils: {pose_path}")
pose_module = importlib.util.module_from_spec(pose_spec)
pose_spec.loader.exec_module(pose_module)
source = (OFFICIAL_REF_DIR / "utils" / "inference_utils.py").read_text()
source = source.replace(
"from .pose_utils import interpolate_camera_poses\n", "")
namespace = {"interpolate_camera_poses": pose_module.interpolate_camera_poses}
exec(compile(source, str(OFFICIAL_REF_DIR / "utils" / "inference_utils.py"), "exec"), namespace)
except Exception as exc: # noqa: BLE001 - local parity should skip missing reference deps.
pytest.skip(f"Cannot load DreamX camera reference: {exc}")
return namespace["ActionToPoseFromID"], namespace["GetPoseEmbedsFromPosesPrope"]
def _official_camera_condition(
action_seq: list[str],
action_speed_list: list[float],
*,
num_frames: int,
height: int,
width: int,
dtype: torch.dtype,
):
action_to_pose, get_pose_embeds = _load_official_functions()
duration = -(-num_frames // len(action_seq))
poses = action_to_pose(action_seq, action_speed_list, duration=duration)[:num_frames]
condition, _ = get_pose_embeds(poses, height, width, len(poses), False, 0, dtype=dtype, device="cpu")
return condition
@pytest.mark.parametrize(
("action_seq", "action_speed_list", "num_frames"),
[
(["w"], [4], 81),
(["wj", "d"], [4, 6], 121),
(["i", "k", "l"], [3, 5, 2], 85),
],
)
def test_dreamx_world_camera_conditioning_matches_official(action_seq, action_speed_list, num_frames):
dtype = torch.float32
official = _official_camera_condition(
action_seq,
action_speed_list,
num_frames=num_frames,
height=704,
width=1280,
dtype=dtype,
)
fastvideo = build_dreamx_camera_condition(
action_seq,
action_speed_list,
num_frames=num_frames,
height=704,
width=1280,
dtype=dtype,
device="cpu",
)
assert official.keys() == fastvideo.keys() == {"viewmats", "K"}
for key in ("viewmats", "K"):
assert official[key].shape == fastvideo[key].shape
diff = (official[key] - fastvideo[key]).abs()
print(f"{key}: shape={tuple(fastvideo[key].shape)} diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}")
assert_close(fastvideo[key], official[key], atol=1e-5, rtol=1e-5)
@@ -0,0 +1,64 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World conversion script smoke tests."""
import json
from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_dit_config
from fastvideo.models.registry import _LEGACY_FAST_VIDEO_MODELS
from scripts.checkpoint_conversion.dreamx_world_to_diffusers import (
MODEL_INDEX,
REUSED_COMPONENTS,
TRANSFORMER_CONFIG,
write_model_index,
)
def test_dreamx_world_converter_writes_full_model_index_with_reused_components(tmp_path):
component_source = tmp_path / "wan22"
output = tmp_path / "dreamx"
output.mkdir()
for component in REUSED_COMPONENTS:
component_dir = component_source / component
component_dir.mkdir(parents=True)
(component_dir / "config.json").write_text("{}\n")
write_model_index(output, component_source, symlink_components=True)
model_index = json.loads((output / "model_index.json").read_text())
assert model_index == MODEL_INDEX
assert model_index["_class_name"] == "DreamXWorldPipeline"
assert model_index["transformer"] == ["diffusers", "DreamXWorldTransformer3DModel"]
for component in REUSED_COMPONENTS:
assert (output / component).is_symlink()
def test_dreamx_world_transformer_config_has_camera_adapter_enabled():
assert TRANSFORMER_CONFIG["_class_name"] == "DreamXWorldTransformer3DModel"
assert TRANSFORMER_CONFIG["add_control_adapter"] is True
assert TRANSFORMER_CONFIG["cam_method"] == "prope"
assert TRANSFORMER_CONFIG["num_layers"] == 30
def test_dreamx_world_converter_transformer_config_matches_pipeline_dit_config():
dit_config = make_dreamx_world_5b_cam_dit_config()
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["cross_attn_norm"] == dit_config.cross_attn_norm
assert TRANSFORMER_CONFIG["qk_norm"] == dit_config.qk_norm
assert TRANSFORMER_CONFIG["add_control_adapter"] == dit_config.add_control_adapter
assert TRANSFORMER_CONFIG["cam_method"] == dit_config.cam_method
assert TRANSFORMER_CONFIG["attn_compress"] == dit_config.attn_compress
assert TRANSFORMER_CONFIG["cam_self_attn_layers"] == dit_config.cam_self_attn_layers
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
@@ -0,0 +1,152 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World pipeline config and conditioning smoke tests."""
from types import SimpleNamespace
import json
import torch
from fastvideo.api.presets import get_preset
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.dreamx_world_pipeline import DreamXWorldPipeline
from fastvideo.pipelines.basic.dreamx_world.stages import (
DREAMX_Y_CAMERA_KEY,
DreamXWorldCameraConditioningStage,
)
from fastvideo.pipelines.stages.denoising import DenoisingStage
from fastvideo.registry import get_default_preset, get_model_info, get_pipeline_config_cls_from_name
def test_dreamx_world_5b_cam_pipeline_config_wires_first_scope_components():
config = DreamXWorld5BCamPipelineConfig()
assert config.flow_shift == 3.0
assert config.ti2v_task is True
assert config.expand_timesteps is True
assert config.dit_config.expand_timesteps is True
assert config.dit_config.num_layers == 30
assert config.dit_config.add_control_adapter is True
assert config.dit_config.cam_method == "prope"
assert config.vae_config.load_encoder is True
assert config.vae_config.load_decoder is True
assert config.vae_config.z_dim == 48
assert config.vae_config.scale_factor_temporal == 4
assert config.vae_config.scale_factor_spatial == 16
assert len(config.text_encoder_configs) == 1
text_config = config.text_encoder_configs[0]
assert text_config.prefix == "umt5"
assert text_config.vocab_size == 256384
assert text_config.d_model == 4096
assert config.text_encoder_precisions == ("bf16",)
def test_dreamx_world_pipeline_registry_discovers_entrypoint():
pipelines = import_pipeline_classes(PipelineType.BASIC)
assert pipelines["basic"]["DreamXWorldPipeline"] is DreamXWorldPipeline
def test_dreamx_world_local_model_index_resolves_model_info(tmp_path):
model_dir = tmp_path / "DreamX-World-5B-Cam-converted"
model_dir.mkdir()
for component in ("scheduler", "text_encoder", "tokenizer", "transformer", "vae"):
(model_dir / component).mkdir()
(model_dir / "model_index.json").write_text(json.dumps({
"_class_name": "DreamXWorldPipeline",
"_diffusers_version": "0.31.0",
"scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
"text_encoder": ["transformers", "UMT5EncoderModel"],
"tokenizer": ["transformers", "AutoTokenizer"],
"transformer": ["diffusers", "DreamXWorldTransformer3DModel"],
"vae": ["diffusers", "AutoencoderKLWan"],
}) + "\n")
info = get_model_info(str(model_dir), pipeline_type=PipelineType.BASIC, workload_type=WorkloadType.I2V)
assert info.pipeline_cls is DreamXWorldPipeline
assert info.pipeline_config_cls is DreamXWorld5BCamPipelineConfig
def test_dreamx_world_model_path_resolves_pipeline_config():
assert get_pipeline_config_cls_from_name("GD-ML/DreamX-World-5B-Cam") is DreamXWorld5BCamPipelineConfig
def test_dreamx_world_default_preset_is_registered():
preset_name = get_default_preset("GD-ML/DreamX-World-5B-Cam")
preset = get_preset(preset_name, "dreamx_world")
assert preset.name == "dreamx_world_5b_cam"
assert preset.workload_type == "i2v"
assert preset.defaults["height"] == 480
assert preset.defaults["width"] == 832
assert preset.defaults["num_frames"] == 161
assert preset.defaults["num_inference_steps"] == 30
assert preset.defaults["guidance_scale"] == 5.0
def test_dreamx_world_pipeline_initializes_official_flow_scheduler():
pipeline = DreamXWorldPipeline.__new__(DreamXWorldPipeline)
pipeline.modules = {}
fastvideo_args = SimpleNamespace(pipeline_config=DreamXWorld5BCamPipelineConfig())
pipeline.initialize_pipeline(fastvideo_args)
scheduler = pipeline.modules["scheduler"]
assert isinstance(scheduler, FlowMatchEulerDiscreteScheduler)
assert scheduler.config.shift == 3.0
def test_dreamx_world_camera_conditioning_stage_sets_y_camera_extra():
batch = ForwardBatch(
data_type="t2v",
action_list=["wj", "d"],
action_speed_list=[4, 6],
num_frames=17,
height=704,
width=1280,
latents=torch.zeros(1, 16, 5, 44, 80),
)
stage = DreamXWorldCameraConditioningStage()
out = stage.forward(batch, fastvideo_args=object())
y_camera = out.extra[DREAMX_Y_CAMERA_KEY]
expected = build_dreamx_camera_condition(
["wj", "d"],
[4, 6],
num_frames=17,
height=704,
width=1280,
dtype=torch.float32,
device="cpu",
)
assert set(y_camera) == {"viewmats", "K"}
for key, expected_value in expected.items():
assert y_camera[key].shape == (1, *expected_value.shape)
torch.testing.assert_close(y_camera[key][0], expected_value)
assert stage.verify_output(out, fastvideo_args=object()).is_valid()
def test_dreamx_world_denoising_kwargs_filter_for_y_camera():
y_camera = {"viewmats": torch.eye(4).reshape(1, 1, 4, 4), "K": torch.eye(3).reshape(1, 1, 3, 3)}
stage = DenoisingStage.__new__(DenoisingStage)
def accepts_y_camera(hidden_states, encoder_hidden_states, timestep, y_camera=None):
return y_camera
def no_y_camera(hidden_states, encoder_hidden_states, timestep):
return hidden_states
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}) == {}
@@ -0,0 +1,53 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World default Flow scheduler parity.
Coverage scope: implementation_subcomponent. DreamX-World-5B-Cam defaults to
Diffusers FlowMatchEulerDiscreteScheduler for sampler_name=Flow. This test
checks that FastVideo's native FlowMatchEulerDiscreteScheduler matches the
timestep schedule and Euler step used by the official default path.
"""
from __future__ import annotations
from pathlib import Path
import inspect
import torch
from diffusers import FlowMatchEulerDiscreteScheduler as OfficialFlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from torch.testing import assert_close
from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
FlowMatchEulerDiscreteScheduler as FastVideoFlowMatchEulerDiscreteScheduler,
)
REPO_ROOT = Path(__file__).resolve().parents[3]
PARITY_SCOPE = "implementation_subcomponent"
def _scheduler_kwargs(cls):
config_path = REPO_ROOT / "DreamX-World" / "configs" / "wan2.2" / "wan_ti2v_5b.yaml"
config = OmegaConf.load(config_path)
raw_kwargs = OmegaConf.to_container(config["scheduler_kwargs"])
signature = inspect.signature(cls)
return {key: value for key, value in raw_kwargs.items() if key in signature.parameters}
def test_dreamx_world_flow_scheduler_timesteps_and_step_match():
official = OfficialFlowMatchEulerDiscreteScheduler(**_scheduler_kwargs(OfficialFlowMatchEulerDiscreteScheduler))
fastvideo = FastVideoFlowMatchEulerDiscreteScheduler(**_scheduler_kwargs(FastVideoFlowMatchEulerDiscreteScheduler))
official.set_timesteps(50, device="cpu", mu=1)
fastvideo.set_timesteps(50, device="cpu", mu=1)
assert_close(fastvideo.timesteps, official.timesteps, atol=0, rtol=0)
assert_close(fastvideo.sigmas, official.sigmas, atol=0, rtol=0)
torch.manual_seed(7)
sample = torch.randn(1, 4, 2, 8, 8)
model_output = torch.randn_like(sample)
timestep = official.timesteps[3]
official_prev = official.step(model_output, timestep, sample, return_dict=False)[0]
fastvideo_prev = fastvideo.step(model_output, fastvideo.timesteps[3], sample, return_dict=False)[0]
diff = (official_prev - fastvideo_prev).abs()
print(f"scheduler diff_max={diff.max().item():.8f} diff_mean={diff.mean().item():.8f}")
assert_close(fastvideo_prev, official_prev, atol=0, rtol=0)
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World Wan T5 encoder reuse parity scaffold.
Coverage scope: implementation_subcomponent. It records the official
WanT5EncoderModel loading path and FastVideo T5 target for later activation
with staged Wan2.2 base text encoder/tokenizer weights.
"""
from __future__ import annotations
import os
from pathlib import Path
import sys
import types
import pytest
import torch
from omegaconf import OmegaConf
from torch.testing import assert_close
from transformers import AutoTokenizer
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.configs.pipelines.dreamx_world import (
DreamXWorld5BCamPipelineConfig,
make_dreamx_world_5b_cam_text_encoder_config,
)
from fastvideo.models.loader.component_loader import TextEncoderLoader
REPO_ROOT = Path(__file__).resolve().parents[3]
OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World"))
WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B"))
WAN_DIFFUSERS_DIR = Path(os.getenv("DREAMX_WORLD_WAN_DIFFUSERS_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B-Diffusers"))
PARITY_SCOPE = "implementation_subcomponent"
def _add_official_to_path():
if not OFFICIAL_REF_DIR.exists():
pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}")
if str(OFFICIAL_REF_DIR) not in sys.path:
sys.path.insert(0, str(OFFICIAL_REF_DIR))
def _install_xfuser_stub() -> None:
if "xfuser" in sys.modules:
return
xfuser = types.ModuleType("xfuser")
core = types.ModuleType("xfuser.core")
distributed = types.ModuleType("xfuser.core.distributed")
long_ctx = types.ModuleType("xfuser.core.long_ctx_attention")
distributed.get_sequence_parallel_rank = lambda: 0
distributed.get_sequence_parallel_world_size = lambda: 1
distributed.get_sp_group = lambda: None
distributed.get_world_group = lambda: types.SimpleNamespace(local_rank=0, rank=0)
distributed.init_distributed_environment = lambda *args, **kwargs: None
distributed.initialize_model_parallel = lambda *args, **kwargs: None
distributed.model_parallel_is_initialized = lambda: False
class XFuserLongContextAttention:
def __call__(self, *args, **kwargs):
raise RuntimeError("xfuser stub cannot execute attention")
long_ctx.xFuserLongContextAttention = XFuserLongContextAttention
sys.modules.update({
"xfuser": xfuser,
"xfuser.core": core,
"xfuser.core.distributed": distributed,
"xfuser.core.long_ctx_attention": long_ctx,
})
def _text_kwargs():
config = OmegaConf.load(OFFICIAL_REF_DIR / "configs" / "wan2.2" / "wan_ti2v_5b.yaml")
return OmegaConf.to_container(config["text_encoder_kwargs"])
def _patch_single_process_text_parallel(monkeypatch):
import fastvideo.layers.linear as fastvideo_linear
import fastvideo.layers.vocab_parallel_embedding as fastvideo_embedding
import fastvideo.models.encoders.t5 as fastvideo_t5
for module in (fastvideo_t5, fastvideo_embedding, fastvideo_linear):
if hasattr(module, "get_tp_rank"):
monkeypatch.setattr(module, "get_tp_rank", lambda: 0)
if hasattr(module, "get_tp_world_size"):
monkeypatch.setattr(module, "get_tp_world_size", lambda: 1)
monkeypatch.setattr(fastvideo_embedding, "tensor_model_parallel_all_reduce", lambda x: x)
def _load_official_text_encoder(device, dtype):
_add_official_to_path()
_install_xfuser_stub()
text_path = WAN_BASE_DIR / "models_t5_umt5-xxl-enc-bf16.pth"
if not text_path.exists():
pytest.skip(f"Wan2.2 text encoder weights missing: {text_path}")
try:
from models import WanT5EncoderModel
except Exception as exc: # noqa: BLE001
pytest.skip(f"Cannot import official DreamX text encoder: {exc}")
model = WanT5EncoderModel.from_pretrained(
str(text_path), additional_kwargs=_text_kwargs(), low_cpu_mem_usage=True, torch_dtype=dtype
)
return model.to(device=device, dtype=dtype).eval()
def _load_fastvideo_text_encoder(device, dtype, monkeypatch):
text_encoder_path = WAN_DIFFUSERS_DIR / "text_encoder"
if not text_encoder_path.exists():
pytest.skip(f"Wan2.2 Diffusers text encoder missing: {text_encoder_path}")
_patch_single_process_text_parallel(monkeypatch)
pipeline_config = DreamXWorld5BCamPipelineConfig()
pipeline_config.text_encoder_configs[0]._fsdp_shard_conditions = []
args = FastVideoArgs(
model_path=str(text_encoder_path),
pipeline_config=pipeline_config,
text_encoder_cpu_offload=(device.type == "cpu"),
)
args.model_paths = {}
return TextEncoderLoader().load(str(text_encoder_path), args).to(device=device, dtype=dtype).eval()
def test_dreamx_world_text_encoder_config_matches_umt5_xxl_shape():
config = make_dreamx_world_5b_cam_text_encoder_config()
assert config.vocab_size == 256384
assert config.d_model == 4096
assert config.d_kv == 64
assert config.d_ff == 10240
assert config.num_heads == 64
assert config.num_layers == 24
assert config.relative_attention_num_buckets == 32
assert config.dropout_rate == 0.0
assert config.text_len == 512
assert config.prefix == "umt5"
def test_dreamx_world_fastvideo_text_encoder_loads_staged_weights(monkeypatch):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = _load_fastvideo_text_encoder(device, torch.bfloat16, monkeypatch)
assert model.__class__.__name__ == "UMT5EncoderModel"
assert next(model.parameters()).device.type == device.type
assert next(model.parameters()).dtype == torch.bfloat16
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for text encoder parity.")
def test_dreamx_world_text_encoder_parity_scaffold(monkeypatch):
device = torch.device("cuda:0")
dtype = torch.bfloat16
official = _load_official_text_encoder(device, dtype)
fastvideo = _load_fastvideo_text_encoder(device, dtype, monkeypatch)
tokenizer_path = WAN_BASE_DIR / "google" / "umt5-xxl"
if not tokenizer_path.exists():
pytest.skip(f"Wan2.2 tokenizer missing: {tokenizer_path}")
tokenizer = AutoTokenizer.from_pretrained(str(tokenizer_path))
batch = tokenizer(["A quiet forest trail at sunrise."], padding="max_length", max_length=512, return_tensors="pt")
input_ids = batch.input_ids.to(device)
attention_mask = batch.attention_mask.to(device)
with torch.inference_mode():
official_hidden = official(input_ids, attention_mask=attention_mask)[0].float().cpu()
fastvideo_hidden = fastvideo(input_ids, attention_mask=attention_mask).last_hidden_state.float().cpu()
assert official_hidden.shape == fastvideo_hidden.shape
assert_close(fastvideo_hidden, official_hidden, atol=1e-3, rtol=1e-3)
@@ -0,0 +1,331 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World transformer parity scaffold.
Coverage scope: both. The official side loads DreamX-World-5B-Cam through
Wan2_2Transformer3DModel.from_pretrained with PRoPE camera control enabled.
The FastVideo side strict-loads the converted DreamX transformer weights into
the native DreamX-World DiT implementation.
"""
from __future__ import annotations
import math
import os
from pathlib import Path
import sys
import types
import pytest
import torch
from omegaconf import OmegaConf
from torch.testing import assert_close
from fastvideo.forward_context import set_forward_context
from fastvideo.configs.models.dits.dreamx_world import (
DreamXWorldArchConfig, DreamXWorldConfig)
from fastvideo.pipelines.basic.dreamx_world.camera_conditioning import build_dreamx_camera_condition
from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_dit_config
from fastvideo.platforms import AttentionBackendEnum
from fastvideo.models.dits.dreamx_world import (
DreamXPropeSelfAttention, DreamXWorldTransformer3DModel,
DreamXWorldTransformerBlock)
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
from scripts.checkpoint_conversion.dreamx_world_to_diffusers import map_transformer_key
REPO_ROOT = Path(__file__).resolve().parents[3]
OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World"))
LOCAL_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_LOCAL_WEIGHTS_DIR", REPO_ROOT / "official_weights" / "dreamx_world"))
WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B"))
CONVERTED_WEIGHTS_DIR = Path(os.getenv("DREAMX_WORLD_CONVERTED_WEIGHTS_DIR", REPO_ROOT / "converted_weights" / "dreamx_world"))
PARITY_SCOPE = "both"
def _install_xfuser_stub() -> None:
if "xfuser" in sys.modules:
return
xfuser = types.ModuleType("xfuser")
core = types.ModuleType("xfuser.core")
distributed = types.ModuleType("xfuser.core.distributed")
long_ctx = types.ModuleType("xfuser.core.long_ctx_attention")
distributed.get_sequence_parallel_rank = lambda: 0
distributed.get_sequence_parallel_world_size = lambda: 1
distributed.get_sp_group = lambda: None
distributed.get_world_group = lambda: types.SimpleNamespace(local_rank=0, rank=0)
distributed.init_distributed_environment = lambda *args, **kwargs: None
distributed.initialize_model_parallel = lambda *args, **kwargs: None
distributed.model_parallel_is_initialized = lambda: False
class XFuserLongContextAttention:
def __call__(self, *args, **kwargs):
raise RuntimeError("xfuser stub cannot execute attention")
long_ctx.xFuserLongContextAttention = XFuserLongContextAttention
sys.modules.update({
"xfuser": xfuser,
"xfuser.core": core,
"xfuser.core.distributed": distributed,
"xfuser.core.long_ctx_attention": long_ctx,
})
def _make_tiny_dreamx_config() -> DreamXWorldConfig:
return DreamXWorldConfig(
arch_config=DreamXWorldArchConfig(
num_attention_heads=1,
attention_head_dim=8,
in_channels=16,
out_channels=16,
ffn_dim=32,
num_layers=1,
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 _add_official_to_path() -> None:
if not OFFICIAL_REF_DIR.exists():
pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}")
if str(OFFICIAL_REF_DIR) not in sys.path:
sys.path.insert(0, str(OFFICIAL_REF_DIR))
def _official_transformer_kwargs() -> dict:
config_path = OFFICIAL_REF_DIR / "configs" / "wan2.2" / "wan_ti2v_5b.yaml"
if not config_path.exists():
pytest.skip(f"DreamX Wan config missing: {config_path}")
config = OmegaConf.load(config_path)
kwargs = OmegaConf.to_container(config["transformer_additional_kwargs"])
kwargs["cam_method"] = "prope"
kwargs["add_control_adapter"] = True
return kwargs
def _load_official_transformer(device: torch.device, dtype: torch.dtype) -> torch.nn.Module:
_add_official_to_path()
_install_xfuser_stub()
if not LOCAL_WEIGHTS_DIR.exists():
pytest.skip(f"DreamX transformer weights missing: {LOCAL_WEIGHTS_DIR}")
try:
from models import Wan2_2Transformer3DModel
except Exception as exc: # noqa: BLE001 - local parity should skip missing refs.
pytest.skip(f"Cannot import official DreamX transformer: {exc}")
model = Wan2_2Transformer3DModel.from_pretrained(
str(LOCAL_WEIGHTS_DIR),
transformer_additional_kwargs=_official_transformer_kwargs(),
torch_dtype=dtype,
)
return model.to(device=device, dtype=dtype).eval()
def _load_fastvideo_transformer(device: torch.device, dtype: torch.dtype) -> torch.nn.Module:
model = _load_fastvideo_transformer_strict(torch.device("cpu"), dtype)
return model.to(device=device, dtype=dtype).eval()
def _load_fastvideo_transformer_strict(device: torch.device, dtype: torch.dtype) -> torch.nn.Module:
transformer_dir = CONVERTED_WEIGHTS_DIR / "transformer"
if not transformer_dir.exists():
pytest.skip(f"Converted DreamX transformer missing: {transformer_dir}")
safetensors_files = resolve_safetensors_files(str(transformer_dir))
config = make_dreamx_world_5b_cam_dit_config()
import fastvideo.models.dits.dreamx_world as fastvideo_dreamx
original_get_sp_world_size = fastvideo_dreamx.get_sp_world_size
fastvideo_dreamx.get_sp_world_size = lambda: 1
try:
with torch.device("meta"):
model = DreamXWorldTransformer3DModel(config=config, hf_config={})
finally:
fastvideo_dreamx.get_sp_world_size = original_get_sp_world_size
incompatible = load_model_from_full_model_state_dict(
model,
safetensors_weights_iterator(safetensors_files, to_cpu=True),
device=device,
param_dtype=dtype,
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())
return model.to(device=device, dtype=dtype).eval()
def _make_inputs(device: torch.device, dtype: torch.dtype):
torch.manual_seed(1234)
num_frames = 5
height = 64
width = 64
latent_frames = (num_frames - 1) // 4 + 1
latent_h = height // 16
latent_w = width // 16
x = torch.randn(1, 48, latent_frames, latent_h, latent_w, device=device, dtype=dtype)
context = [torch.randn(16, 4096, device=device, dtype=dtype)]
seq_len = math.ceil((latent_h * latent_w) / 4 * latent_frames)
timestep = torch.full((1, seq_len), 250, device=device, dtype=torch.long)
camera = build_dreamx_camera_condition(
["w"], [4], num_frames=num_frames, height=height, width=width, dtype=dtype, device=device
)
camera = {key: value.unsqueeze(0) for key, value in camera.items()}
return {"x": [x[0]], "context": context, "t": timestep, "seq_len": seq_len, "y_camera": camera}
def _run_official(model: torch.nn.Module, inputs: dict) -> torch.Tensor:
with torch.inference_mode():
output = model(**inputs)
if isinstance(output, list):
output = torch.stack(output, dim=0)
assert torch.is_tensor(output), f"official output is not a tensor: {type(output)}"
return output.detach().float().cpu()
def _run_fastvideo(model: torch.nn.Module, inputs: dict) -> torch.Tensor:
hidden_states = torch.stack(inputs["x"], dim=0)
encoder_hidden_states = torch.stack([
torch.cat([inputs["context"][0], inputs["context"][0].new_zeros(512 - inputs["context"][0].shape[0], 4096)])
])
with torch.inference_mode(), set_forward_context(current_timestep=0, attn_metadata=None):
output = model(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
timestep=inputs["t"],
y_camera=inputs["y_camera"],
)
assert torch.is_tensor(output), f"FastVideo output is not a tensor: {type(output)}"
return output.detach().float().cpu()
def test_dreamx_world_conversion_mapping_strict_load_smoke(monkeypatch):
_add_official_to_path()
_install_xfuser_stub()
try:
from models import Wan2_2Transformer3DModel
except Exception as exc: # noqa: BLE001
pytest.skip(f"Cannot import official DreamX transformer: {exc}")
official = Wan2_2Transformer3DModel(
dim=8,
ffn_dim=32,
num_heads=1,
num_layers=1,
add_control_adapter=True,
cam_method="prope",
)
official_state = official.state_dict()
diffusers_like_state = {
map_transformer_key(key): value.detach().clone()
for key, value in official_state.items()
}
import fastvideo.models.dits.dreamx_world as fastvideo_dreamx
monkeypatch.setattr(fastvideo_dreamx, "get_sp_world_size", lambda: 1)
with torch.device("meta"):
fastvideo = DreamXWorldTransformer3DModel(config=_make_tiny_dreamx_config(), hf_config={})
incompatible = load_model_from_full_model_state_dict(
fastvideo,
iter(diffusers_like_state.items()),
device=torch.device("cpu"),
param_dtype=torch.float32,
strict=True,
param_names_mapping=get_param_names_mapping(fastvideo.param_names_mapping),
training_mode=False,
)
assert incompatible.missing_keys == []
assert incompatible.unexpected_keys == []
assert not any(param.is_meta for param in fastvideo.parameters())
def test_dreamx_world_5b_cam_dit_config_matches_official_shape():
config = make_dreamx_world_5b_cam_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.add_control_adapter is True
assert config.cam_method == "prope"
assert config.attn_compress == 1
def test_dreamx_world_converted_5b_transformer_strict_loads():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
_load_fastvideo_transformer_strict(device, torch.bfloat16)
def test_dreamx_world_fastvideo_prope_branch_smoke():
block = DreamXWorldTransformerBlock(
8,
32,
1,
cross_attn_norm=True,
add_control_adapter=True,
cam_method="prope",
attn_compress=1,
layer_idx=0,
supported_attention_backends=(AttentionBackendEnum.TORCH_SDPA,),
)
assert block.cam_self_attn is not None
assert [
name for name, _ in block.named_parameters()
if name.startswith("cam_self_attn.")
][:8] == [
"cam_self_attn.q_proj.weight",
"cam_self_attn.q_proj.bias",
"cam_self_attn.k_proj.weight",
"cam_self_attn.k_proj.bias",
"cam_self_attn.v_proj.weight",
"cam_self_attn.v_proj.bias",
"cam_self_attn.out_proj.weight",
"cam_self_attn.out_proj.bias",
]
module = DreamXPropeSelfAttention(
dim=8,
attn_dim=8,
num_heads=1,
qk_norm="rms_norm_across_heads",
).eval()
assert module.num_heads == 1
assert module.head_dim == 8
assert tuple(module.out_proj.weight.shape) == (8, 8)
assert torch.count_nonzero(module.out_proj.weight) == 0
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for transformer parity.")
def test_dreamx_world_transformer_parity_scaffold(monkeypatch):
device = torch.device("cuda:0")
dtype = torch.float32
inputs = _make_inputs(device, dtype)
official = _load_official_transformer(device, dtype)
official_out = _run_official(official, inputs)
del official
torch.cuda.empty_cache()
import fastvideo.attention.layer as attention_layer
import fastvideo.distributed.communication_op as communication_op
import fastvideo.models.dits.dreamx_world as fastvideo_dreamx
monkeypatch.setattr(attention_layer, "get_sp_parallel_rank", lambda: 0)
monkeypatch.setattr(attention_layer, "get_sp_world_size", lambda: 1)
monkeypatch.setattr(attention_layer, "sequence_model_parallel_all_to_all_4D", lambda tensor, scatter_dim=2, gather_dim=1: tensor)
monkeypatch.setattr(attention_layer, "sequence_model_parallel_all_gather", lambda tensor, dim=-1: tensor)
monkeypatch.setattr(fastvideo_dreamx, "get_sp_world_size", lambda: 1)
monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 1)
monkeypatch.setattr(fastvideo_dreamx, "sequence_model_parallel_shard", lambda tensor, dim=1: (tensor, tensor.shape[dim]))
monkeypatch.setattr(
fastvideo_dreamx,
"sequence_model_parallel_all_gather_with_unpad",
lambda tensor, original_seq_len, dim=1: tensor.narrow(dim, 0, original_seq_len),
)
fastvideo = _load_fastvideo_transformer(device, dtype)
fastvideo_out = _run_fastvideo(fastvideo, inputs)
assert official_out.shape == fastvideo_out.shape
diff = (official_out - fastvideo_out).abs()
print(f"diff_max={diff.max().item():.6f} diff_mean={diff.mean().item():.6f}")
assert_close(fastvideo_out, official_out, atol=1e-1, rtol=1e-1)
@@ -0,0 +1,220 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World Wan2.2 VAE reuse parity scaffold.
Coverage scope: implementation_subcomponent. The official side uses
AutoencoderKLWan3_8 from DreamX, while the FastVideo side targets the native
Wan VAE. This remains a scaffold until Wan2.2 base VAE weights are staged.
"""
from __future__ import annotations
import os
from pathlib import Path
import sys
import re
import types
import pytest
import torch
from omegaconf import OmegaConf
from torch.testing import assert_close
from fastvideo.configs.pipelines.dreamx_world import make_dreamx_world_5b_cam_vae_config
from fastvideo.models.vaes.wanvae import AutoencoderKLWan
REPO_ROOT = Path(__file__).resolve().parents[3]
OFFICIAL_REF_DIR = Path(os.getenv("DREAMX_WORLD_OFFICIAL_REF_DIR", REPO_ROOT / "DreamX-World"))
WAN_BASE_DIR = Path(os.getenv("DREAMX_WORLD_WAN_BASE_DIR", REPO_ROOT / "official_weights" / "Wan2.2-TI2V-5B"))
PARITY_SCOPE = "implementation_subcomponent"
def _add_official_to_path():
if not OFFICIAL_REF_DIR.exists():
pytest.skip(f"Official reference missing: {OFFICIAL_REF_DIR}")
if str(OFFICIAL_REF_DIR) not in sys.path:
sys.path.insert(0, str(OFFICIAL_REF_DIR))
def _install_xfuser_stub() -> None:
if "xfuser" in sys.modules:
return
xfuser = types.ModuleType("xfuser")
core = types.ModuleType("xfuser.core")
distributed = types.ModuleType("xfuser.core.distributed")
long_ctx = types.ModuleType("xfuser.core.long_ctx_attention")
distributed.get_sequence_parallel_rank = lambda: 0
distributed.get_sequence_parallel_world_size = lambda: 1
distributed.get_sp_group = lambda: None
distributed.get_world_group = lambda: types.SimpleNamespace(local_rank=0, rank=0)
distributed.init_distributed_environment = lambda *args, **kwargs: None
distributed.initialize_model_parallel = lambda *args, **kwargs: None
distributed.model_parallel_is_initialized = lambda: False
class XFuserLongContextAttention:
def __call__(self, *args, **kwargs):
raise RuntimeError("xfuser stub cannot execute attention")
long_ctx.xFuserLongContextAttention = XFuserLongContextAttention
sys.modules.update({
"xfuser": xfuser,
"xfuser.core": core,
"xfuser.core.distributed": distributed,
"xfuser.core.long_ctx_attention": long_ctx,
})
def _map_residual_subkey(prefix: str, sub: str) -> str | None:
if sub == "residual.0.gamma":
return f"{prefix}.norm1.gamma"
match = re.match(r"^residual\.2\.(weight|bias)$", sub)
if match:
return f"{prefix}.conv1.{match.group(1)}"
if sub == "residual.3.gamma":
return f"{prefix}.norm2.gamma"
match = re.match(r"^residual\.6\.(weight|bias)$", sub)
if match:
return f"{prefix}.conv2.{match.group(1)}"
match = re.match(r"^shortcut\.(weight|bias)$", sub)
if match:
return f"{prefix}.conv_shortcut.{match.group(1)}"
return None
def _map_attention_subkey(prefix: str, sub: str) -> str | None:
if sub == "norm.gamma":
return f"{prefix}.norm.gamma"
match = re.match(r"^(to_qkv|proj)\.(weight|bias)$", sub)
if match:
return f"{prefix}.{match.group(1)}.{match.group(2)}"
return None
def _map_resample_subkey(prefix: str, sub: str) -> str | None:
match = re.match(r"^resample\.1\.(weight|bias)$", sub)
if match:
return f"{prefix}.resample.1.{match.group(1)}"
match = re.match(r"^time_conv\.(weight|bias)$", sub)
if match:
return f"{prefix}.time_conv.{match.group(1)}"
return None
def _map_dreamx_raw_vae_key(key: str) -> str | None:
match = re.match(r"^(conv1|conv2)\.(weight|bias)$", key)
if match:
prefix = "quant_conv" if match.group(1) == "conv1" else "post_quant_conv"
return f"{prefix}.{match.group(2)}"
match = re.match(r"^(encoder|decoder)\.conv1\.(weight|bias)$", key)
if match:
return f"{match.group(1)}.conv_in.{match.group(2)}"
match = re.match(r"^(encoder|decoder)\.head\.0\.gamma$", key)
if match:
return f"{match.group(1)}.norm_out.gamma"
match = re.match(r"^(encoder|decoder)\.head\.2\.(weight|bias)$", key)
if match:
return f"{match.group(1)}.conv_out.{match.group(2)}"
match = re.match(r"^(encoder|decoder)\.middle\.0\.(.*)$", key)
if match:
return _map_residual_subkey(f"{match.group(1)}.mid_block.resnets.0", match.group(2))
match = re.match(r"^(encoder|decoder)\.middle\.1\.(.*)$", key)
if match:
return _map_attention_subkey(f"{match.group(1)}.mid_block.attentions.0", match.group(2))
match = re.match(r"^(encoder|decoder)\.middle\.2\.(.*)$", key)
if match:
return _map_residual_subkey(f"{match.group(1)}.mid_block.resnets.1", match.group(2))
match = re.match(r"^encoder\.downsamples\.(\d+)\.downsamples\.(\d+)\.(.*)$", key)
if match:
stage = int(match.group(1))
block = int(match.group(2))
sub = match.group(3)
if block in (0, 1):
return _map_residual_subkey(f"encoder.down_blocks.{stage}.resnets.{block}", sub)
if block == 2:
return _map_resample_subkey(f"encoder.down_blocks.{stage}.downsampler", sub)
return None
match = re.match(r"^decoder\.upsamples\.(\d+)\.upsamples\.(\d+)\.(.*)$", key)
if match:
stage = int(match.group(1))
block = int(match.group(2))
sub = match.group(3)
if block in (0, 1, 2):
return _map_residual_subkey(f"decoder.up_blocks.{stage}.resnets.{block}", sub)
if block == 3:
return _map_resample_subkey(f"decoder.up_blocks.{stage}.upsampler", sub)
return None
return None
def _vae_kwargs():
config = OmegaConf.load(OFFICIAL_REF_DIR / "configs" / "wan2.2" / "wan_ti2v_5b.yaml")
return OmegaConf.to_container(config["vae_kwargs"])
def _load_official_vae(device, dtype):
_add_official_to_path()
_install_xfuser_stub()
vae_path = WAN_BASE_DIR / "Wan2.2_VAE.pth"
if not vae_path.exists():
pytest.skip(f"Wan2.2 base VAE weights missing: {vae_path}")
try:
from models import AutoencoderKLWan3_8
except Exception as exc: # noqa: BLE001
pytest.skip(f"Cannot import official DreamX VAE: {exc}")
model = AutoencoderKLWan3_8.from_pretrained(str(vae_path), additional_kwargs=_vae_kwargs())
return model.to(device=device, dtype=dtype).eval()
def _load_fastvideo_vae(device, dtype):
vae_path = WAN_BASE_DIR / "Wan2.2_VAE.pth"
if not vae_path.exists():
pytest.skip(f"Wan2.2 raw VAE weights missing: {vae_path}")
config = make_dreamx_world_5b_cam_vae_config()
config.load_encoder = True
config.load_decoder = True
model = AutoencoderKLWan(config).to(device=device, dtype=dtype)
raw_state = torch.load(str(vae_path), map_location="cpu", weights_only=True)
mapped_state = {}
for key, value in raw_state.items():
mapped_key = _map_dreamx_raw_vae_key(key)
if mapped_key is None:
raise AssertionError(f"Unmapped DreamX raw VAE key: {key}")
mapped_state[mapped_key] = value
model.load_state_dict(mapped_state, strict=True)
return model.eval()
def _normalize_fastvideo_vae_latent(latent: torch.Tensor) -> torch.Tensor:
config = make_dreamx_world_5b_cam_vae_config()
mean = torch.tensor(config.latents_mean, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1)
std = torch.tensor(config.latents_std, device=latent.device, dtype=latent.dtype).view(1, -1, 1, 1, 1)
return (latent - mean) / std
def test_dreamx_world_vae_config_matches_wan22_shape():
config = make_dreamx_world_5b_cam_vae_config()
assert config.z_dim == 48
assert config.in_channels == 12
assert config.out_channels == 12
assert config.base_dim == 160
assert config.decoder_base_dim == 256
assert config.scale_factor_temporal == 4
assert config.scale_factor_spatial == 16
assert config.patch_size == 2
assert config.is_residual is True
assert config.clip_output is False
assert len(config.latents_mean) == 48
assert len(config.latents_std) == 48
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for VAE parity.")
def test_dreamx_world_vae_encode_parity_scaffold():
device = torch.device("cuda:0")
dtype = torch.bfloat16
official = _load_official_vae(device, dtype)
fastvideo = _load_fastvideo_vae(device, dtype)
torch.manual_seed(123)
video = torch.randn(1, 3, 5, 64, 64, device=device, dtype=dtype).clamp(-1, 1)
with torch.inference_mode():
official_latent = official.encode(video).latent_dist.mean.float().cpu()
fastvideo_latent = _normalize_fastvideo_vae_latent(fastvideo.encode(video).mean).float().cpu()
assert official_latent.shape == fastvideo_latent.shape
assert_close(fastvideo_latent, official_latent, atol=5e-2, rtol=5e-2)
@@ -0,0 +1,112 @@
# SPDX-License-Identifier: Apache-2.0
"""DreamX-World pipeline parity checks.
This test compares the FastVideo pipeline's DreamX-specific conditioning and
single-step scheduler path against an explicit hand-rolled pass using the same
loaded modules. It is intentionally local and deterministic: component parity
against the official DreamX repository lives in ``tests/local_tests/dreamx_world``.
"""
from __future__ import annotations
import gc
import os
from pathlib import Path
from typing import Any, cast
import pytest
import torch
from torch.testing import assert_close
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
MODEL_DIR = Path(os.getenv("DREAMX_WORLD_MODEL_DIR", "converted_weights/dreamx_world"))
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available(),
reason="DreamX-World pipeline parity requires CUDA",
)
def _run_worker_forward_batch(worker_wrapper: Any, request_kwargs: dict[str, Any]) -> torch.Tensor:
from fastvideo.api.sampling_param import SamplingParam
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.utils import shallow_asdict
fastvideo_args = worker_wrapper.worker.fastvideo_args
sampling_param = SamplingParam.from_pretrained(fastvideo_args.model_path)
sampling_param.update({
key: value
for key, value in request_kwargs.items()
if key not in {"prompt", "output_path"}
})
sampling_param.prompt = request_kwargs["prompt"]
latents_size = [
(sampling_param.num_frames - 1) // 4 + 1,
sampling_param.height // 8,
sampling_param.width // 8,
]
n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
batch = ForwardBatch(
**shallow_asdict(sampling_param),
eta=0.0,
n_tokens=n_tokens,
VSA_sparsity=fastvideo_args.VSA_sparsity,
)
output_batch = worker_wrapper.worker.pipeline.forward(batch, fastvideo_args)
assert output_batch.output is not None
return output_batch.output.detach().cpu()
def _close_generator(generator: Any) -> None:
generator.shutdown()
gc.collect()
torch.cuda.empty_cache()
def test_dreamx_world_one_step_pipeline_latent_matches_manual_pass() -> None:
if not MODEL_DIR.exists():
pytest.fail(f"DreamX-World converted model directory is missing: {MODEL_DIR}")
from fastvideo import VideoGenerator
common_kwargs = dict(
prompt="a quiet road through a futuristic city at sunrise",
output_path="outputs_video/dreamx_world_parity",
save_video=False,
return_frames=True,
height=64,
width=64,
num_frames=9,
num_inference_steps=1,
guidance_scale=1.0,
action_list=["w"],
action_speed_list=[2.0],
seed=123,
)
generator = VideoGenerator.from_pretrained(
str(MODEL_DIR),
num_gpus=1,
use_fsdp_inference=False,
dit_cpu_offload=False,
vae_cpu_offload=True,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
output_type="latent",
override_pipeline_cls_name="DreamXWorldPipeline",
)
try:
result = cast(dict[str, Any], generator.generate_video(**common_kwargs))
pipeline_latents = cast(torch.Tensor, result["samples"]).detach().cpu()
manual_latents = generator.executor.collective_rpc(
_run_worker_forward_batch,
kwargs={"request_kwargs": common_kwargs},
)[0]
finally:
_close_generator(generator)
assert_close(pipeline_latents, manual_latents, atol=0.0, rtol=0.0)
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
"""Smoke tests for the DreamX-World-5B-Cam pipeline."""
from __future__ import annotations
import os
from pathlib import Path
from typing import Any, cast
import pytest
import torch
from PIL import Image, ImageDraw
os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA")
MODEL_DIR = Path(os.getenv("DREAMX_WORLD_MODEL_DIR", "converted_weights/dreamx_world"))
def _write_smoke_image(path: Path) -> None:
image = Image.new("RGB", (96, 96), color=(42, 76, 112))
draw = ImageDraw.Draw(image)
draw.rectangle((0, 56, 96, 96), fill=(36, 44, 52))
draw.polygon([(0, 56), (48, 30), (96, 56)], fill=(120, 142, 158))
draw.rectangle((34, 42, 62, 72), fill=(178, 198, 212))
image.save(path)
pytestmark = pytest.mark.skipif(
not torch.cuda.is_available(),
reason="DreamX-World pipeline smoke requires CUDA",
)
def test_dreamx_world_typed_surface_preflight() -> None:
import fastvideo.registry as registry
from fastvideo.api.presets import get_preset, get_presets_for_family
from fastvideo.configs.pipelines.dreamx_world import DreamXWorld5BCamPipelineConfig
from fastvideo.fastvideo_args import WorkloadType
from fastvideo.pipelines.basic.dreamx_world.dreamx_world_pipeline import (
DreamXWorldPipeline,
EntryClass,
)
assert DreamXWorldPipeline.__name__ == "DreamXWorldPipeline"
assert EntryClass is DreamXWorldPipeline
assert DreamXWorldPipeline._required_config_modules == [
"text_encoder",
"tokenizer",
"vae",
"transformer",
"scheduler",
]
default_preset, model_family = registry.get_preset_selection(
"GD-ML/DreamX-World-5B-Cam"
)
assert model_family == "dreamx_world"
assert default_preset == "dreamx_world_5b_cam"
info = registry.get_model_info(
"GD-ML/DreamX-World-5B-Cam",
workload_type=WorkloadType.I2V,
override_pipeline_cls_name="DreamXWorldPipeline",
)
assert info.pipeline_cls is DreamXWorldPipeline
assert info.pipeline_config_cls is DreamXWorld5BCamPipelineConfig
names = {p.name for p in get_presets_for_family("dreamx_world")}
assert "dreamx_world_5b_cam" in names
preset = get_preset("dreamx_world_5b_cam", "dreamx_world")
assert preset.defaults["num_inference_steps"] == 30
assert preset.defaults["height"] == 480
assert preset.defaults["width"] == 832
assert preset.defaults["num_frames"] == 161
assert preset.defaults["guidance_scale"] == 5.0
cfg = DreamXWorld5BCamPipelineConfig()
assert cfg.flow_shift == 3.0
assert cfg.ti2v_task is True
assert cfg.expand_timesteps is True
assert cfg.dit_config.arch_config.add_control_adapter is True
assert cfg.dit_config.arch_config.cam_method == "prope"
def test_dreamx_world_camera_stage_writes_y_camera() -> None:
from types import SimpleNamespace
from fastvideo.pipelines.basic.dreamx_world.stages import (
DREAMX_Y_CAMERA_KEY,
DreamXWorldCameraConditioningStage,
)
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
batch = ForwardBatch(
data_type="video",
prompt="camera smoke",
latents=torch.zeros(1, 48, 3, 8, 8, dtype=torch.bfloat16, device="cuda"),
num_frames=9,
height=64,
width=64,
action_list=["w", "d"],
action_speed_list=[2.0, 1.0],
)
out = DreamXWorldCameraConditioningStage().forward(batch, cast(Any, SimpleNamespace()))
y_camera = out.extra[DREAMX_Y_CAMERA_KEY]
assert set(y_camera) == {"viewmats", "K"}
assert y_camera["viewmats"].shape == (1, 3, 4, 4)
assert y_camera["K"].shape == (1, 3, 3, 3)
assert y_camera["viewmats"].device.type == "cuda"
assert y_camera["viewmats"].dtype == torch.bfloat16
def test_dreamx_world_pipeline_load_generate_latent_smoke(tmp_path: Path) -> None:
if not MODEL_DIR.exists():
pytest.fail(f"DreamX-World converted model directory is missing: {MODEL_DIR}")
from fastvideo import VideoGenerator
image_path = tmp_path / "dreamx_world_smoke_input.png"
_write_smoke_image(image_path)
generator = VideoGenerator.from_pretrained(
str(MODEL_DIR),
num_gpus=1,
use_fsdp_inference=False,
dit_cpu_offload=False,
vae_cpu_offload=True,
text_encoder_cpu_offload=True,
pin_cpu_memory=False,
output_type="latent",
override_pipeline_cls_name="DreamXWorldPipeline",
)
try:
result = generator.generate_video(
prompt="a quiet road through a futuristic city at sunrise",
output_path="outputs_video/dreamx_world_smoke",
save_video=False,
return_frames=True,
height=64,
width=64,
num_frames=9,
num_inference_steps=1,
guidance_scale=1.0,
image_path=str(image_path),
action_list=["w"],
action_speed_list=[2.0],
seed=0,
)
finally:
generator.shutdown()
assert isinstance(result, dict)
samples = cast(dict[str, Any], result)["samples"]
assert torch.is_tensor(samples)
assert samples.ndim == 5
assert samples.shape[1] == 48
assert torch.isfinite(samples).all()