diff --git a/.gitignore b/.gitignore index 3b3663ec7..0734cada6 100644 --- a/.gitignore +++ b/.gitignore @@ -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/ diff --git a/examples/inference/basic/basic_dreamx_world.py b/examples/inference/basic/basic_dreamx_world.py new file mode 100644 index 000000000..990d8ca97 --- /dev/null +++ b/examples/inference/basic/basic_dreamx_world.py @@ -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() diff --git a/fastvideo/configs/models/dits/__init__.py b/fastvideo/configs/models/dits/__init__.py index 094ddb474..4171fdcc2 100644 --- a/fastvideo/configs/models/dits/__init__.py +++ b/fastvideo/configs/models/dits/__init__.py @@ -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" ] diff --git a/fastvideo/configs/models/dits/dreamx_world.py b/fastvideo/configs/models/dits/dreamx_world.py new file mode 100644 index 000000000..286e2c3ef --- /dev/null +++ b/fastvideo/configs/models/dits/dreamx_world.py @@ -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" diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index 31eb042ac..68fb8f3e6 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -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" ] diff --git a/fastvideo/configs/pipelines/dreamx_world.py b/fastvideo/configs/pipelines/dreamx_world.py new file mode 100644 index 000000000..3e6b8d140 --- /dev/null +++ b/fastvideo/configs/pipelines/dreamx_world.py @@ -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 diff --git a/fastvideo/models/dits/dreamx_world.py b/fastvideo/models/dits/dreamx_world.py new file mode 100644 index 000000000..b6fcc92a9 --- /dev/null +++ b/fastvideo/models/dits/dreamx_world.py @@ -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 diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index dc0246e68..c105143dd 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -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) diff --git a/fastvideo/pipelines/basic/dreamx_world/__init__.py b/fastvideo/pipelines/basic/dreamx_world/__init__.py new file mode 100644 index 000000000..8dbc4bff1 --- /dev/null +++ b/fastvideo/pipelines/basic/dreamx_world/__init__.py @@ -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", +] diff --git a/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py b/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py new file mode 100644 index 000000000..b330c1125 --- /dev/null +++ b/fastvideo/pipelines/basic/dreamx_world/camera_conditioning.py @@ -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} diff --git a/fastvideo/pipelines/basic/dreamx_world/config.py b/fastvideo/pipelines/basic/dreamx_world/config.py new file mode 100644 index 000000000..cce9047e8 --- /dev/null +++ b/fastvideo/pipelines/basic/dreamx_world/config.py @@ -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", +] diff --git a/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py new file mode 100644 index 000000000..72656bb8b --- /dev/null +++ b/fastvideo/pipelines/basic/dreamx_world/dreamx_world_pipeline.py @@ -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 diff --git a/fastvideo/pipelines/basic/dreamx_world/presets.py b/fastvideo/pipelines/basic/dreamx_world/presets.py new file mode 100644 index 000000000..05fb28aa2 --- /dev/null +++ b/fastvideo/pipelines/basic/dreamx_world/presets.py @@ -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, ) diff --git a/fastvideo/pipelines/basic/dreamx_world/stages.py b/fastvideo/pipelines/basic/dreamx_world/stages.py new file mode 100644 index 000000000..feb0d0471 --- /dev/null +++ b/fastvideo/pipelines/basic/dreamx_world/stages.py @@ -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 diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 2a4fdb391..8c1938ac6 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -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, ) diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 647178e65..c6f1467ca 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -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, diff --git a/fastvideo/tests/ssim/test_dreamx_world_similarity.py b/fastvideo/tests/ssim/test_dreamx_world_similarity.py new file mode 100644 index 000000000..9f6861921 --- /dev/null +++ b/fastvideo/tests/ssim/test_dreamx_world_similarity.py @@ -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), + }, + ) diff --git a/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py new file mode 100644 index 000000000..6297863fd --- /dev/null +++ b/scripts/checkpoint_conversion/dreamx_world_to_diffusers.py @@ -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 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() diff --git a/tests/local_tests/dreamx_world/PORT_STATUS.md b/tests/local_tests/dreamx_world/PORT_STATUS.md new file mode 100644 index 000000000..e4e7bb20b --- /dev/null +++ b/tests/local_tests/dreamx_world/PORT_STATUS.md @@ -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 +``` diff --git a/tests/local_tests/dreamx_world/README.md b/tests/local_tests/dreamx_world/README.md new file mode 100644 index 000000000..c19051482 --- /dev/null +++ b/tests/local_tests/dreamx_world/README.md @@ -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. diff --git a/tests/local_tests/dreamx_world/__init__.py b/tests/local_tests/dreamx_world/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/tests/local_tests/dreamx_world/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py new file mode 100644 index 000000000..2448e81e2 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_camera_conditioning_parity.py @@ -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) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py b/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py new file mode 100644 index 000000000..5815478a0 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_conversion.py @@ -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 diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py new file mode 100644 index 000000000..b9993e479 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_pipeline_config.py @@ -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}) == {} + diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py new file mode 100644 index 000000000..b394bbae7 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_scheduler_parity.py @@ -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) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py new file mode 100644 index 000000000..d51d83314 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_text_encoder_parity.py @@ -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) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py new file mode 100644 index 000000000..fff655317 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_transformer_parity.py @@ -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) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py new file mode 100644 index 000000000..c994b4c12 --- /dev/null +++ b/tests/local_tests/dreamx_world/test_dreamx_world_vae_parity.py @@ -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) diff --git a/tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py b/tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py new file mode 100644 index 000000000..ce839132d --- /dev/null +++ b/tests/local_tests/pipelines/test_dreamx_world_pipeline_parity.py @@ -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) diff --git a/tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py b/tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py new file mode 100644 index 000000000..bd956bdcd --- /dev/null +++ b/tests/local_tests/pipelines/test_dreamx_world_pipeline_smoke.py @@ -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()