790 lines
32 KiB
Python
790 lines
32 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
from PIL import Image
|
|
from torch import nn
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
class StubScheduler:
|
|
def __init__(self, timesteps: list[int] | None = None, *, flow_shift: float = 1.0) -> None:
|
|
self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
|
|
self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift)
|
|
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
|
|
self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
|
|
|
def set_timesteps(self, num_steps: int, device: torch.device) -> None:
|
|
self.set_timesteps_calls.append((num_steps, device))
|
|
self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device)
|
|
|
|
def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
|
|
del kwargs
|
|
self.step_calls.append((noise_pred.clone(), timestep.clone(), latents.clone()))
|
|
return (latents + noise_pred,)
|
|
|
|
|
|
class _ModeLatentDist:
|
|
def __init__(self, latents: torch.Tensor) -> None:
|
|
self._latents = latents
|
|
|
|
def mode(self) -> torch.Tensor:
|
|
return self._latents
|
|
|
|
|
|
class StubCosmos3VAE:
|
|
dtype = torch.float32
|
|
|
|
def __init__(self, z_dim: int = 2, *, temporal: int = 4, spatial: int = 8) -> None:
|
|
self.config = SimpleNamespace(
|
|
z_dim=z_dim,
|
|
scale_factor_temporal=temporal,
|
|
scale_factor_spatial=spatial,
|
|
latents_mean=[0.0] * z_dim,
|
|
latents_std=[1.0] * z_dim,
|
|
)
|
|
|
|
def encode(self, video: torch.Tensor):
|
|
latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1
|
|
latent_height = video.shape[-2] // self.config.scale_factor_spatial
|
|
latent_width = video.shape[-1] // self.config.scale_factor_spatial
|
|
latents = torch.ones(
|
|
video.shape[0],
|
|
self.config.z_dim,
|
|
latent_frames,
|
|
latent_height,
|
|
latent_width,
|
|
dtype=video.dtype,
|
|
device=video.device,
|
|
)
|
|
return SimpleNamespace(latent_dist=_ModeLatentDist(latents))
|
|
|
|
def decode(self, latents: torch.Tensor, return_dict: bool = False):
|
|
del return_dict
|
|
return (latents,)
|
|
|
|
|
|
class StubCosmos3AVAE:
|
|
def __init__(self, **kwargs: Any) -> None:
|
|
self.kwargs = kwargs
|
|
self.sample_rate = int(kwargs["sample_rate"])
|
|
self.audio_channels = int(kwargs["audio_channels"])
|
|
self.latent_ch = int(kwargs["io_channels"])
|
|
self.temporal_compression_factor = int(kwargs["hop_size"])
|
|
|
|
def get_latent_num_samples(self, num_audio_samples: int) -> int:
|
|
return int(num_audio_samples) // self.temporal_compression_factor
|
|
|
|
def get_audio_num_samples(self, num_latent_samples: int) -> int:
|
|
return int(num_latent_samples) * self.temporal_compression_factor
|
|
|
|
def decode(self, latents: torch.Tensor) -> torch.Tensor:
|
|
return torch.zeros(latents.shape[0], self.audio_channels, 8)
|
|
|
|
|
|
class StubCosmos3Transformer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
latent_channel_size: int = 2,
|
|
sound_gen: bool = False,
|
|
sound_dim: int = 3,
|
|
sound_latent_fps: float = 25.0,
|
|
action_gen: bool = False,
|
|
action_dim: int = 4,
|
|
) -> None:
|
|
super().__init__()
|
|
self.latent_channel_size = latent_channel_size
|
|
self.sound_gen = sound_gen
|
|
self.sound_dim = sound_dim
|
|
self.sound_latent_fps = sound_latent_fps
|
|
self.action_gen = action_gen
|
|
self.action_dim = action_dim
|
|
self.cached_kv: Any | None = None
|
|
self.cached_freqs_gen: Any | None = None
|
|
self.calls: list[dict[str, Any]] = []
|
|
self.reset_calls = 0
|
|
|
|
def reset_cache(self) -> None:
|
|
self.reset_calls += 1
|
|
self.cached_kv = None
|
|
self.cached_freqs_gen = None
|
|
|
|
def forward(
|
|
self,
|
|
*,
|
|
hidden_states: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
text_ids: torch.Tensor,
|
|
text_mask: torch.Tensor,
|
|
**kwargs: Any,
|
|
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
|
token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0
|
|
sound_latents = kwargs.get("sound_latents")
|
|
self.calls.append(
|
|
{
|
|
"token": token,
|
|
"timestep": timestep.clone(),
|
|
"text_mask": text_mask.clone(),
|
|
"cache_before": self.cached_kv,
|
|
"kwargs": dict(kwargs),
|
|
}
|
|
)
|
|
if self.cached_kv is None:
|
|
marker = torch.tensor([token], dtype=torch.float32)
|
|
self.cached_kv = [(marker, marker + 100)]
|
|
self.cached_freqs_gen = (marker + 200, marker + 300)
|
|
action_latents = kwargs.get("action_latents")
|
|
outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))]
|
|
if action_latents is not None:
|
|
outputs.append(torch.full_like(action_latents, float(token + 20)))
|
|
if sound_latents is not None:
|
|
outputs.append(torch.full_like(sound_latents, float(token + 10)))
|
|
return outputs[0] if len(outputs) == 1 else tuple(outputs)
|
|
|
|
|
|
def passthrough_progress_bar(iterable):
|
|
return iterable
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
|
|
module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
|
|
module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
|
|
module.ensure_initialized = lambda od_config: None
|
|
module.check_text_safety = lambda text: None
|
|
module.check_video_safety = lambda video: video
|
|
monkeypatch.setitem(sys.modules, module.__name__, module)
|
|
return module
|
|
|
|
|
|
@pytest.fixture
|
|
def make_cosmos3_pipeline():
|
|
def _make():
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import (
|
|
Cosmos3OmniDiffusersPipeline,
|
|
)
|
|
|
|
pipeline = object.__new__(Cosmos3OmniDiffusersPipeline)
|
|
nn.Module.__init__(pipeline)
|
|
pipeline.od_config = SimpleNamespace()
|
|
pipeline.device = torch.device("cpu")
|
|
pipeline.dtype = torch.float32
|
|
pipeline.transformer = StubCosmos3Transformer(latent_channel_size=2)
|
|
pipeline.vae = StubCosmos3VAE(z_dim=2)
|
|
pipeline.vae_scale_factor_temporal = 4
|
|
pipeline.vae_scale_factor_spatial = 8
|
|
pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0)
|
|
pipeline._base_scheduler_config = pipeline.scheduler.config
|
|
pipeline._engine_init_flow_shift = 1.0
|
|
pipeline._current_flow_shift = 1.0
|
|
pipeline._guidance_scale = None
|
|
pipeline._num_timesteps = None
|
|
pipeline._cache_dit_requires_paired_cfg = False
|
|
pipeline._sound_tokenizer = None
|
|
pipeline.progress_bar = passthrough_progress_bar
|
|
return pipeline
|
|
|
|
return _make
|
|
|
|
|
|
def make_sampling_params(**overrides: Any) -> SimpleNamespace:
|
|
values = {
|
|
"height": None,
|
|
"width": None,
|
|
"num_frames": None,
|
|
"num_inference_steps": None,
|
|
"guidance_scale": None,
|
|
"generator": None,
|
|
"seed": 123,
|
|
"num_outputs_per_prompt": 1,
|
|
"frame_rate": None,
|
|
"resolved_frame_rate": None,
|
|
"max_sequence_length": None,
|
|
"extra_args": {},
|
|
}
|
|
values.update(overrides)
|
|
return SimpleNamespace(**values)
|
|
|
|
|
|
def _ids(value: int) -> torch.Tensor:
|
|
return torch.tensor([[value]], dtype=torch.long)
|
|
|
|
|
|
def _mask() -> torch.Tensor:
|
|
return torch.ones(1, 1, dtype=torch.long)
|
|
|
|
|
|
def test_pipeline_registered_and_exported() -> None:
|
|
from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS
|
|
from vllm_omni.diffusion.models import cosmos3
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
|
|
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
|
from vllm_omni.diffusion.registry import (
|
|
_DIFFUSION_MODELS,
|
|
_DIFFUSION_POST_PROCESS_FUNCS,
|
|
_DIFFUSION_PRE_PROCESS_FUNCS,
|
|
)
|
|
|
|
assert issubclass(Cosmos3OmniDiffusersPipeline, nn.Module)
|
|
assert issubclass(Cosmos3OmniDiffusersPipeline, ProgressBarMixin)
|
|
assert Cosmos3OmniDiffusersPipeline.support_image_input is True
|
|
assert _DIFFUSION_MODELS["Cosmos3OmniDiffusersPipeline"] == (
|
|
"cosmos3",
|
|
"pipeline_cosmos3",
|
|
"Cosmos3OmniDiffusersPipeline",
|
|
)
|
|
assert _DIFFUSION_PRE_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_pre_process_func"
|
|
assert _DIFFUSION_POST_PROCESS_FUNCS["Cosmos3OmniDiffusersPipeline"] == "get_cosmos3_post_process_func"
|
|
assert "Cosmos3OmniDiffusersPipeline" in CUSTOM_DIT_ENABLERS
|
|
assert "Cosmos3OmniDiffusersPipeline" in cosmos3.__all__
|
|
|
|
|
|
@pytest.fixture
|
|
def stub_real_pipeline_init(monkeypatch: pytest.MonkeyPatch):
|
|
from vllm_omni.diffusion.models.cosmos3 import pipeline_cosmos3
|
|
|
|
class _StubAutoTokenizer:
|
|
@classmethod
|
|
def from_pretrained(cls, *args, **kwargs):
|
|
return SimpleNamespace()
|
|
|
|
class _StubDiffusersVAE:
|
|
config = SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, *args, **kwargs):
|
|
return cls()
|
|
|
|
def to(self, _device):
|
|
return self
|
|
|
|
class _StubDiffusersScheduler:
|
|
config = SimpleNamespace(flow_shift=1.0)
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, *args, **kwargs):
|
|
return cls()
|
|
|
|
class _StubVideoProcessor:
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
pass
|
|
|
|
monkeypatch.setattr(pipeline_cosmos3, "AutoTokenizer", _StubAutoTokenizer)
|
|
monkeypatch.setattr(pipeline_cosmos3, "DistributedAutoencoderKLWan", _StubDiffusersVAE)
|
|
monkeypatch.setattr(pipeline_cosmos3, "UniPCMultistepScheduler", _StubDiffusersScheduler)
|
|
monkeypatch.setattr(pipeline_cosmos3, "VideoProcessor", _StubVideoProcessor)
|
|
monkeypatch.setattr(pipeline_cosmos3, "get_local_device", lambda: torch.device("cpu"))
|
|
|
|
|
|
def _make_od_config(*, sound_gen: bool) -> SimpleNamespace:
|
|
tf_model_config = {
|
|
"hidden_size": 8,
|
|
"num_hidden_layers": 0,
|
|
"num_attention_heads": 2,
|
|
"num_key_value_heads": 2,
|
|
"head_dim": 4,
|
|
"intermediate_size": 16,
|
|
"vocab_size": 32,
|
|
"latent_patch_size": 1,
|
|
"latent_channel": 2,
|
|
"rope_scaling": {"mrope_section": [1, 1, 0]},
|
|
}
|
|
if sound_gen:
|
|
tf_model_config["sound_gen"] = True
|
|
return SimpleNamespace(
|
|
enable_cpu_offload=False,
|
|
enable_diffusion_pipeline_profiler=False,
|
|
model="/nonexistent/model/path",
|
|
dtype=torch.float32,
|
|
flow_shift=None,
|
|
quantization_config=None,
|
|
custom_pipeline_args={},
|
|
model_config={},
|
|
tf_model_config=tf_model_config,
|
|
)
|
|
|
|
|
|
def test_pipeline_init_skips_tokenizer_when_sound_disabled(stub_real_pipeline_init) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
|
|
|
|
pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=False))
|
|
|
|
assert pipeline._sound_tokenizer is None
|
|
assert pipeline.transformer.sound_gen is False
|
|
assert not hasattr(pipeline.transformer, "audio_proj_in")
|
|
assert not hasattr(pipeline.transformer, "audio_proj_out")
|
|
|
|
|
|
def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
|
|
stub_real_pipeline_init,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
|
|
|
|
stub_tokenizer = sound_tokenizer.Cosmos3SoundTokenizer(
|
|
StubCosmos3AVAE(sample_rate=32000, audio_channels=2, io_channels=5, hop_size=800)
|
|
)
|
|
monkeypatch.setattr(
|
|
sound_tokenizer.Cosmos3SoundTokenizer,
|
|
"from_config",
|
|
classmethod(lambda cls, od_config: stub_tokenizer),
|
|
)
|
|
|
|
pipeline = Cosmos3OmniDiffusersPipeline(od_config=_make_od_config(sound_gen=True))
|
|
|
|
assert pipeline._sound_tokenizer is stub_tokenizer
|
|
assert pipeline.transformer.sound_gen is True
|
|
assert pipeline.transformer.sound_dim == pipeline._sound_tokenizer.latent_ch == 5
|
|
assert pipeline.transformer.sound_latent_fps == pipeline._sound_tokenizer.latent_fps == 40.0
|
|
assert pipeline.transformer.audio_proj_in.in_features == 5
|
|
assert pipeline.transformer.audio_proj_out.out_features == 5
|
|
|
|
|
|
def test_preprocess_i2v_image_and_action_video_inputs() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
|
|
|
|
preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
|
|
i2v = SimpleNamespace(
|
|
prompts=[{"prompt": "A slow camera push.", "multi_modal_data": {"image": Image.new("RGB", (320, 160))}}],
|
|
sampling_params=SimpleNamespace(height=None, width=None, extra_args={}),
|
|
)
|
|
|
|
result = preprocess(i2v)
|
|
assert (result.sampling_params.height, result.sampling_params.width) == (672, 1344)
|
|
assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344)
|
|
|
|
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue")]
|
|
action = SimpleNamespace(
|
|
prompts=[{"prompt": "Move.", "multi_modal_data": {"video": frames}}],
|
|
sampling_params=SimpleNamespace(height=16, width=32, extra_args={"action_mode": "forward_dynamics"}),
|
|
)
|
|
|
|
additional = preprocess(action).prompts[0]["additional_information"]
|
|
assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
|
|
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
|
|
|
|
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
|
|
v2v = SimpleNamespace(
|
|
prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
|
|
sampling_params=SimpleNamespace(
|
|
height=16,
|
|
width=32,
|
|
extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
|
|
),
|
|
)
|
|
additional = preprocess(v2v).prompts[0]["additional_information"]
|
|
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
|
|
assert additional["condition_frame_indexes_vision"] == [0, 1]
|
|
|
|
|
|
def test_postprocess_handles_image_video_audio_and_validation() -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
|
|
|
|
func = get_cosmos3_post_process_func(SimpleNamespace())
|
|
video = torch.zeros(1, 3, 1, 4, 4)
|
|
|
|
assert func(video, output_type="latent") is video
|
|
assert func({"image": video})[0].size == (4, 4)
|
|
# Video-only postprocess returns the bare processed video (not a dict),
|
|
# matching the image/latent branches and peer audio-capable pipelines.
|
|
assert not isinstance(func({"video": video}), dict)
|
|
assert (
|
|
func(
|
|
{"video": video, "audio": torch.ones(1, 2, 16), "audio_sample_rate": 48000},
|
|
sampling_params=SimpleNamespace(extra_args={"resolved_frame_rate": 12}),
|
|
)["audio_sample_rate"]
|
|
== 48000
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="text-to-image postprocess expects"):
|
|
func({"image": torch.zeros(1, 3, 2, 4, 4)})
|
|
with pytest.raises(ValueError, match="both image and video"):
|
|
func({"image": video, "video": video})
|
|
|
|
|
|
def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> None:
|
|
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
|
|
|
|
pipeline = make_cosmos3_pipeline()
|
|
captured: list[str] = []
|
|
pipeline._tokenize_prompt = lambda text, *args, **kwargs: (captured.append(text) or _ids(len(captured)), _mask())
|
|
|
|
pipeline._format_and_tokenize_prompts(
|
|
"A robot",
|
|
"bad",
|
|
num_frames=48,
|
|
frame_rate=24,
|
|
height=720,
|
|
width=1280,
|
|
max_sequence_length=32,
|
|
sp=SimpleNamespace(
|
|
extra_args={
|
|
"negative_metadata_mode": "inverse",
|
|
# Duration/resolution metadata templates are off by default
|
|
# (commit "Change resolution and duration templates to off by
|
|
# default"); enable them explicitly to exercise the formatting.
|
|
"use_duration_template": True,
|
|
"use_resolution_template": True,
|
|
}
|
|
),
|
|
use_system_prompt=True,
|
|
is_t2i=False,
|
|
)
|
|
assert "The video is 2.0 seconds long" in captured[0]
|
|
assert "The video is not 2.0 seconds long" in captured[1]
|
|
|
|
remaps = {
|
|
"embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
|
|
"model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
|
|
"norm.weight": "transformer.language_model.norm.weight",
|
|
"norm_moe_gen.weight": "transformer.norm_moe_gen.weight",
|
|
"proj_in.weight": "transformer.proj_in.weight",
|
|
"proj_out.bias": "transformer.proj_out.bias",
|
|
"layers.3.self_attn.to_q.weight": "transformer.language_model.layers.3.self_attn.to_q.weight",
|
|
"layers.3.self_attn.to_out.weight": "transformer.language_model.layers.3.self_attn.to_out.weight",
|
|
"layers.3.self_attn.norm_q.weight": "transformer.language_model.layers.3.self_attn.norm_q.weight",
|
|
"layers.3.self_attn.add_q_proj.weight": "transformer.gen_layers.3.cross_attention.to_q.weight",
|
|
"layers.3.self_attn.to_add_out.weight": "transformer.gen_layers.3.cross_attention.to_out.weight",
|
|
"layers.3.self_attn.norm_added_q.weight": "transformer.gen_layers.3.cross_attention.norm_q.weight",
|
|
"transformer.model.layers.3.self_attn.add_k_proj.weight": (
|
|
"transformer.gen_layers.3.cross_attention.to_k.weight"
|
|
),
|
|
}
|
|
assert {key: Cosmos3OmniDiffusersPipeline._remap_ckpt_key(key) for key in remaps} == remaps
|
|
|
|
|
|
def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline) -> None:
|
|
pipeline = make_cosmos3_pipeline()
|
|
latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0))
|
|
assert latents.shape == (1, 2, 2, 2, 3)
|
|
|
|
pipeline._encode_conditioning_video = lambda *args, **kwargs: torch.full((1, 2, 2, 2, 3), 5.0)
|
|
i2v_latents, velocity_mask, image_latent = pipeline._prepare_latents_i2v(
|
|
torch.zeros(1, 3, 16, 24), 16, 24, 5, torch.Generator(device="cpu").manual_seed(0)
|
|
)
|
|
torch.testing.assert_close(i2v_latents[:, :, 0], torch.full((1, 2, 2, 3), 5.0))
|
|
assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
|
|
assert image_latent.shape == (1, 2, 1, 2, 3)
|
|
|
|
pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
|
|
v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
|
|
torch.zeros(1, 3, 5, 16, 24),
|
|
16,
|
|
24,
|
|
9,
|
|
torch.Generator(device="cpu").manual_seed(0),
|
|
[0, 1],
|
|
)
|
|
torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
|
|
assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
|
|
assert v2v_condition.shape == (1, 2, 3, 2, 3)
|
|
|
|
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
|
pipeline._sound_tokenizer = SimpleNamespace(
|
|
sample_rate=10,
|
|
latent_ch=3,
|
|
hop_size=4,
|
|
decode=lambda x: torch.ones(x.shape[0], 2, 24),
|
|
)
|
|
assert pipeline._resolve_sound_target_samples(SimpleNamespace(extra_args={"sound_duration": 2.0}), 9, 3.0) == (
|
|
20,
|
|
2.0,
|
|
10,
|
|
)
|
|
sound_latents, latent_frames = pipeline._prepare_sound_latents(21, torch.Generator(device="cpu").manual_seed(0))
|
|
assert (sound_latents.shape, latent_frames) == (torch.Size([1, 3, 6]), 6)
|
|
assert pipeline._decode_sound_latents(torch.zeros(1, 3, 6), target_audio_samples=21).shape == (1, 2, 21)
|
|
|
|
pipeline.transformer = pipeline.transformer.__class__(action_gen=True, action_dim=4)
|
|
action, action_mask, clean, raw_dim = pipeline._prepare_action_latents(
|
|
mode="forward_dynamics",
|
|
action_chunk_size=2,
|
|
raw_action_dim=None,
|
|
generator=torch.Generator(device="cpu").manual_seed(0),
|
|
sp=SimpleNamespace(extra_args={"action": [[1.0, 2.0], [3.0, 4.0]]}),
|
|
)
|
|
assert raw_dim == 2
|
|
assert action_mask.tolist() == [[[0.0], [0.0]]]
|
|
torch.testing.assert_close(action, clean)
|
|
|
|
|
|
def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
|
|
pipeline = make_cosmos3_pipeline()
|
|
latents = torch.zeros(1, 2, 1, 1, 1)
|
|
|
|
result = pipeline.diffuse(
|
|
latents=latents,
|
|
timesteps=torch.tensor([900, 100]),
|
|
cond_ids=_ids(2),
|
|
cond_mask=_mask(),
|
|
uncond_ids=_ids(1),
|
|
uncond_mask=_mask(),
|
|
guidance_scale=3.0,
|
|
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
|
|
guidance_interval=(500.0, 1000.0),
|
|
)
|
|
assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2]
|
|
torch.testing.assert_close(result, torch.full_like(latents, 6.0))
|
|
|
|
i2v = pipeline.diffuse(
|
|
latents=torch.zeros(1, 2, 2, 1, 1),
|
|
timesteps=torch.tensor([7]),
|
|
cond_ids=_ids(2),
|
|
cond_mask=_mask(),
|
|
uncond_ids=_ids(1),
|
|
uncond_mask=_mask(),
|
|
guidance_scale=1.0,
|
|
shared_kwargs={"video_shape": (2, 1, 1), "fps": 24.0},
|
|
velocity_mask=torch.tensor([[[[[0.0]], [[1.0]]]]]),
|
|
image_latent=torch.full((1, 2, 1, 1, 1), 7.0),
|
|
)
|
|
torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0))
|
|
|
|
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
|
|
video_result, action_result = pipeline.diffuse(
|
|
latents=latents,
|
|
action_latents=torch.zeros(1, 3, 4),
|
|
action_velocity_mask=torch.ones(1, 3, 1),
|
|
action_condition_latents=torch.zeros(1, 3, 4),
|
|
timesteps=torch.tensor([7, 3]),
|
|
cond_ids=_ids(2),
|
|
cond_mask=_mask(),
|
|
uncond_ids=_ids(1),
|
|
uncond_mask=_mask(),
|
|
guidance_scale=1.0,
|
|
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "action_domain_ids": torch.tensor([0])},
|
|
)
|
|
torch.testing.assert_close(video_result, torch.full_like(latents, 4.0))
|
|
torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
|
|
|
|
|
|
def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
|
|
"""With cache-dit active the uncond pass is kept even outside the guidance
|
|
interval (so cache-dit's has_separate_cfg parity stays in phase), and the
|
|
output is numerically identical to the skip path.
|
|
|
|
Contrast with ``test_diffuse_covers_cfg_and_i2v_steps`` (no marker), where
|
|
the same inputs skip the out-of-interval uncond pass: calls == [2, 1, 2].
|
|
"""
|
|
pipeline = make_cosmos3_pipeline()
|
|
# Marker normally set by ``enable_cache_for_cosmos3`` when cache-dit is on.
|
|
pipeline._cache_dit_requires_paired_cfg = True
|
|
latents = torch.zeros(1, 2, 1, 1, 1)
|
|
|
|
result = pipeline.diffuse(
|
|
latents=latents,
|
|
timesteps=torch.tensor([900, 100]),
|
|
cond_ids=_ids(2),
|
|
cond_mask=_mask(),
|
|
uncond_ids=_ids(1),
|
|
uncond_mask=_mask(),
|
|
guidance_scale=3.0,
|
|
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
|
|
guidance_interval=(500.0, 1000.0),
|
|
)
|
|
|
|
# t=900 is inside the interval (cond+uncond); t=100 is outside but the
|
|
# uncond pass is still issued -> paired cond/uncond at every step.
|
|
assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2, 1]
|
|
# Identical result to the skip path: out-of-interval combine uses scale=1.0,
|
|
# so combine_cfg_noise(cond=2, uncond=1, 1.0) == 2 == the skipped cond value.
|
|
torch.testing.assert_close(result, torch.full_like(latents, 6.0))
|
|
|
|
|
|
class TestForwardRouting:
|
|
def _install_forward_stubs(self, pipeline):
|
|
captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
|
|
|
|
def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs):
|
|
captured["format"] = {
|
|
"prompt": prompt,
|
|
"negative_prompt": negative_prompt,
|
|
"num_frames": num_frames,
|
|
"frame_rate": frame_rate,
|
|
"height": height,
|
|
"width": width,
|
|
"is_t2i": kwargs["is_t2i"],
|
|
}
|
|
return _ids(2), _mask(), _ids(1), _mask()
|
|
|
|
def fake_prepare(height, width, num_frames, generator):
|
|
captured["prepare_calls"].append((height, width, num_frames, generator.initial_seed()))
|
|
return torch.zeros(1, 2, 1, 1, 1)
|
|
|
|
def fake_diffuse(**kwargs):
|
|
captured["diffuse_calls"].append(kwargs)
|
|
outputs = [kwargs["latents"] + len(captured["diffuse_calls"])]
|
|
if kwargs.get("action_latents") is not None:
|
|
outputs.append(kwargs["action_latents"] + 3.0)
|
|
if kwargs.get("sound_latents") is not None:
|
|
outputs.append(kwargs["sound_latents"] + 2.0)
|
|
return outputs[0] if len(outputs) == 1 else tuple(outputs)
|
|
|
|
pipeline._format_and_tokenize_prompts = fake_format
|
|
pipeline._prepare_latents = fake_prepare
|
|
pipeline._set_flow_shift = lambda target: captured.setdefault("flow_shifts", []).append(target)
|
|
pipeline.diffuse = fake_diffuse
|
|
pipeline._decode_latents = lambda latents: latents
|
|
return captured
|
|
|
|
@pytest.mark.parametrize(
|
|
("prompt", "sampling_params", "expected"),
|
|
[
|
|
(
|
|
{"prompt": "A painted robot", "modalities": ["image"]},
|
|
make_sampling_params(num_outputs_per_prompt=2),
|
|
{"key": "image", "is_t2i": True, "flow": [3.0], "steps": [50, 50], "frames": 1},
|
|
),
|
|
(
|
|
"A warehouse robot",
|
|
make_sampling_params(),
|
|
{"key": "video", "is_t2i": False, "flow": [1.0], "steps": [35], "frames": 189},
|
|
),
|
|
],
|
|
)
|
|
def test_forward_defaults_and_mode_selection(
|
|
self,
|
|
make_cosmos3_pipeline,
|
|
prompt,
|
|
sampling_params,
|
|
expected,
|
|
) -> None:
|
|
pipeline = make_cosmos3_pipeline()
|
|
captured = self._install_forward_stubs(pipeline)
|
|
|
|
output = pipeline.forward(SimpleNamespace(prompts=[prompt], sampling_params=sampling_params))
|
|
|
|
assert expected["key"] in output.output
|
|
assert captured["format"]["is_t2i"] is expected["is_t2i"]
|
|
assert captured["format"]["num_frames"] == expected["frames"]
|
|
assert captured["flow_shifts"] == expected["flow"]
|
|
assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
|
|
|
|
def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
|
|
pipeline = make_cosmos3_pipeline()
|
|
captured = self._install_forward_stubs(pipeline)
|
|
image_tensor = torch.zeros(1, 3, 16, 16)
|
|
velocity_mask = torch.ones(1, 1, 1, 1, 1)
|
|
|
|
pipeline._prepare_latents_i2v = lambda *args, **kwargs: (
|
|
torch.zeros(1, 2, 1, 1, 1),
|
|
velocity_mask,
|
|
torch.zeros(1, 2, 1, 1, 1),
|
|
)
|
|
pipeline.forward(
|
|
SimpleNamespace(
|
|
prompts=[
|
|
{
|
|
"prompt": "move",
|
|
"modalities": ["video"],
|
|
"additional_information": {"preprocessed_image": image_tensor},
|
|
}
|
|
],
|
|
sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
|
|
)
|
|
)
|
|
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
|
|
|
|
video_tensor = torch.zeros(1, 3, 5, 16, 16)
|
|
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
|
|
v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
|
|
pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
|
|
torch.zeros(1, 2, 2, 1, 1),
|
|
v2v_mask,
|
|
v2v_condition,
|
|
)
|
|
pipeline.forward(
|
|
SimpleNamespace(
|
|
prompts=[
|
|
{
|
|
"prompt": "continue",
|
|
"modalities": ["video"],
|
|
"additional_information": {
|
|
"preprocessed_video": video_tensor,
|
|
"condition_frame_indexes_vision": [0],
|
|
},
|
|
}
|
|
],
|
|
sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
|
|
)
|
|
)
|
|
assert captured["flow_shifts"][-1] == 10.0
|
|
assert captured["format"]["negative_prompt"] == ""
|
|
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
|
|
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
|
|
|
|
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
|
sound_latents = torch.zeros(1, 3, 4)
|
|
pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
|
|
pipeline._prepare_sound_latents = lambda *args: (sound_latents, 4)
|
|
pipeline._decode_sound_latents = lambda *args: torch.ones(1, 2, 20)
|
|
output = pipeline.forward(
|
|
SimpleNamespace(
|
|
prompts=[{"prompt": "A robot", "modalities": ["video"], "generate_sound": True}],
|
|
sampling_params=make_sampling_params(num_frames=9, frame_rate=3.0),
|
|
)
|
|
)
|
|
assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
|
|
assert output.output["audio_sample_rate"] == 10
|
|
|
|
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
|
|
output = pipeline.forward(
|
|
SimpleNamespace(
|
|
prompts=[
|
|
{
|
|
"prompt": "Pick the block.",
|
|
"modalities": ["video"],
|
|
"additional_information": {"preprocessed_image": image_tensor},
|
|
}
|
|
],
|
|
sampling_params=make_sampling_params(
|
|
height=16,
|
|
width=16,
|
|
extra_args={
|
|
"action_mode": "policy",
|
|
"action_chunk_size": 2,
|
|
"raw_action_dim": 2,
|
|
"domain_name": "bridge_orig_lerobot",
|
|
},
|
|
),
|
|
)
|
|
)
|
|
assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
|
|
assert output.custom_output["action"].shape == (1, 2, 2)
|
|
|
|
@pytest.mark.parametrize(
|
|
("prompt", "sampling_params", "message"),
|
|
[
|
|
(["one", "two"], make_sampling_params(), "single prompt"),
|
|
([{"prompt": "one", "modalities": ["image", "video"]}], make_sampling_params(), "both image and video"),
|
|
(
|
|
[{"prompt": "x", "modalities": ["image"], "generate_sound": True}],
|
|
make_sampling_params(),
|
|
"only for video",
|
|
),
|
|
],
|
|
)
|
|
def test_forward_rejects_invalid_public_requests(
|
|
self,
|
|
make_cosmos3_pipeline,
|
|
prompt,
|
|
sampling_params,
|
|
message,
|
|
) -> None:
|
|
pipeline = make_cosmos3_pipeline()
|
|
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
|
|
|
with pytest.raises(ValueError, match=message):
|
|
pipeline.forward(SimpleNamespace(prompts=prompt, sampling_params=sampling_params))
|