diff --git a/.gitignore b/.gitignore index 9155a85cc..278eb50ff 100644 --- a/.gitignore +++ b/.gitignore @@ -242,6 +242,7 @@ data/ datasets/ *.csv *.json +!recipes/cosmos3/negative_prompt.json !tests/dfx/perf/tests/*.json !apps/ComfyUI-vLLM-Omni/example_workflows/*.json !tests/assets/soulxsinger/*.json diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 5995e0911..7fc7d0ba8 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -161,11 +161,21 @@ are not implemented yet. Models may expose additional V2V controls through `extra_params`. For example, Cosmos3 supports `condition_frame_indexes_vision` and `condition_video_keep` to select which decoded reference frames are used as clean conditioning. Cosmos3 transfer mode -also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus -transfer options such as `control_path`, `control_guidance`, -`control_guidance_interval`, `num_video_frames_per_chunk`, -`num_conditional_frames`, `show_control_condition`, and `show_input`; see the -Cosmos3 recipe for complete examples. +also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints. Each hint +may specify its own `control_path` and `control_weight`. Request-level transfer +options include +`control_guidance`, `control_guidance_interval`, +`emphasize_control_in_prompt`, `num_video_frames_per_chunk`, +`num_conditional_frames`, `show_control_condition`, and `show_input`. Transfer +uses its transfer-specific system prompt and, by default, appends a +control-adherence directive to the positive prompt. It adds duration/FPS and +resolution metadata to both CFG branches but does not add a negative prompt +automatically. The Cosmos3 recipe includes an optional reference negative +prompt and shows how to pass it. Set `emphasize_control_in_prompt`, +`use_duration_template`, or `use_resolution_template` to `false` to disable the +corresponding addition. `negative_metadata_mode` accepts `same`, `inverse`, or +`none` and defaults to `same` for transfer. See the Cosmos3 recipe for complete +examples. HTTP redirects for `image_reference.image_url` follow vLLM's `VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to diff --git a/recipes/cosmos3/Cosmos3-Nano.md b/recipes/cosmos3/Cosmos3-Nano.md index b27d1b473..4aaebf0be 100644 --- a/recipes/cosmos3/Cosmos3-Nano.md +++ b/recipes/cosmos3/Cosmos3-Nano.md @@ -195,11 +195,13 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \ # Transfer V2V with a precomputed depth control video. `control_path` can point # to a local image/video; edge and blur can also be computed from `input_reference` -# by passing `"edge":true` or `"blur":true`. +# by passing `"edge":true` or `"blur":true`. The reference negative prompt is +# optional; omit the `negative_prompt` form field to use an empty negative branch. curl -sS -X POST http://localhost:8000/v1/videos/sync \ -H "Accept: video/mp4" \ -F "model=nvidia/Cosmos3-Nano" \ -F "prompt=Generate a realistic scene following the provided control video." \ + --form-string "negative_prompt=$(jq -c . recipes/cosmos3/negative_prompt.json)" \ -F "size=1280x720" -F "num_frames=121" \ -F "num_inference_steps=50" -F "seed=125" \ -F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \ @@ -322,15 +324,34 @@ vllm serve nvidia/Cosmos3-Nano-Policy-DROID \ `extra_params={"guardrails":false}` (per request) toggles safety. The per-request flag only takes effect when the server was launched **with** guardrails enabled (it cannot re-enable them on a `--no-guardrails` server). - `use_resolution_template` / `use_duration_template` are off by default and only - needed when not using upsampled prompts that already encode resolution/duration. + Outside transfer mode, `use_resolution_template` / `use_duration_template` + are off by default and only needed when not using upsampled prompts that + already encode resolution/duration. For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the API decodes the first or last needed reference frames (`"first"` by default). - **Transfer controls:** `extra_params` may include `edge`, `blur`, `depth`, `seg`, or `wsm`. Each hint accepts `true`, a path string, or an object such as `{"control_path": "/path/to/control.mp4"}`; `edge` also accepts - `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. + `preset_edge_threshold` and `blur` accepts `preset_blur_strength`. Every hint + accepts a non-negative `control_weight`; weights are normalized across active + controls and therefore only set their relative influence. A single positive + weight always normalizes to `1.0`; use `control_guidance` to change the + absolute strength of a single control. With two or more active controls, the + per-control attention passes run replicated on every sequence-parallel + (Ulysses) rank, so Ulysses does not reduce per-rank memory or latency for + multi-control transfer requests. + Transfer always uses Cosmos3's transfer-specific system prompt. By default it + also appends a directive naming every active hint and asking the model to + follow its shape, position, and motion precisely; set the request-level + `emphasize_control_in_prompt` option to `false` for prompt ablations. Transfer + does not add a negative prompt automatically. An optional reference prompt is + provided in [`negative_prompt.json`](negative_prompt.json); compact the JSON + with `jq -c` and pass it through `negative_prompt` as shown in the transfer + example. Transfer enables duration/FPS and resolution metadata on both CFG + branches. Set `use_duration_template` or `use_resolution_template` to + `false` to disable either template. `negative_metadata_mode` accepts `same` + (the transfer default), `inverse`, or `none`. Transfer-level options include `control_guidance`, `control_guidance_interval`, `num_video_frames_per_chunk` (default `93`, `101` for WSM), `num_conditional_frames` (default `1`), diff --git a/recipes/cosmos3/Cosmos3-Super.md b/recipes/cosmos3/Cosmos3-Super.md index 54d4d7e32..c8bfb6bbd 100644 --- a/recipes/cosmos3/Cosmos3-Super.md +++ b/recipes/cosmos3/Cosmos3-Super.md @@ -125,6 +125,23 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \ under the top-level `action` field. Verified on the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`) and the rollout video reliably. +- **Transfer controls:** same semantics as Nano — `extra_params` may include + `edge`, `blur`, `depth`, `seg`, or `wsm` hints; see the Cosmos3-Nano recipe + for the request shapes and full option list. Every hint accepts a + non-negative `control_weight`; weights are normalized across active controls + and therefore only set their relative influence. A single positive weight + always normalizes to `1.0`; use `control_guidance` to change the absolute + strength of a single control. With two or more active controls, the + per-control attention passes run replicated on every sequence-parallel + (Ulysses) rank, so Ulysses does not reduce per-rank memory or latency for + multi-control transfer requests. Transfer always uses Cosmos3's + transfer-specific system prompt, appends a control-adherence directive by + default (`emphasize_control_in_prompt: false` disables it), and enables + duration/FPS and resolution metadata on both CFG branches + (`use_duration_template` / `use_resolution_template` to disable; + `negative_metadata_mode` accepts `same`, `inverse`, or `none`, defaulting to + `same`). Transfer does not add a negative prompt automatically; an optional + reference prompt is provided in [`negative_prompt.json`](negative_prompt.json). ## NPU diff --git a/recipes/cosmos3/negative_prompt.json b/recipes/cosmos3/negative_prompt.json new file mode 100644 index 000000000..9cefcb4d5 --- /dev/null +++ b/recipes/cosmos3/negative_prompt.json @@ -0,0 +1,108 @@ +{ + "subjects": [ + { + "description": "Blurry, poorly defined subjects with inconsistent shapes and unrealistic proportions.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "Extremely low-quality subjects with visible rendering artifacts, broken mesh geometry, and completely unrealistic proportions throughout.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + }, + { + "description": "Poorly generated subjects exhibiting all hallmarks of failed neural rendering \u2014 flickering edges, inconsistent depth, and uncanny spatial relationships.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0 + } + ], + "background_setting": "A poorly rendered, flat background with visible seams, repeated textures, and inconsistent depth cues. The environment lacks volumetric depth and appears as a painted backdrop rather than a three-dimensional space. Vegetation looks like flat cutouts with no volumetric depth. The background appears to have been composited from multiple source materials at different resolutions, creating visible seams and edge artifacts where elements meet. Textures swim and shift across surfaces in a way that breaks the illusion of solidity \u2014 patterns drift laterally rather than staying anchored to the geometry they belong to. Background elements flicker in and out of existence between frames, particularly at the edges of the field of view. The rendering resolution is visibly lower for distant elements, creating a jarring transition between near and far objects. Cloud textures repeat obviously in the sky with visible tiling. Water surfaces lack proper reflection and refraction, appearing as flat animated textures. Fog and atmospheric effects pop in and out rather than smoothly transitioning. Trees and vegetation exhibit obvious LOAD (level-of-detail) switching. Building facades have inconsistent window spacing and pattern repetition. The overall scene feels like a poorly assembled collage of individually rendered elements rather than a coherent whole.", + "lighting": { + "conditions": "Harsh, flat lighting with no natural variation. The scene appears uniformly lit as if by a single overhead fluorescent light, removing all sense of depth and atmosphere.", + "direction": "Inconsistent light sources \u2014 shadows point in multiple contradictory directions, breaking physical plausibility.", + "shadows": "Hard-edged, unrealistic shadows that pop in and out of existence between frames. Some objects cast no shadows while others have impossibly dark ones that don't animate smoothly with the object's motion. Shadow edges exhibit visible staircase aliasing artifacts. Shadow maps appear to have been rendered at extremely low resolution, creating blocky patterns. Self-shadowing on characters shows visible peter-panning artifacts where shadows detach from their source. Contact shadows between objects and the ground appear and disappear as objects move slightly. Shadow color is pure black with no ambient contribution, creating an unnaturally harsh contrast that flattens the image. Multiple shadow cascades have visible boundaries where resolution changes. The shadow rendering appears to be temporally unstable \u2014 even static objects have shadows that shimmer and crawl frame to frame, breaking the illusion of a stable light source.", + "illumination_effect": "No bounce light, no ambient occlusion, no subtle color interactions between surfaces. The scene looks like a poorly lit 3D render from the early 2000s." + }, + "aesthetics": { + "composition": "Cluttered, poorly framed composition with no clear focal point. Important elements are cut off by the frame edges. The rule of thirds is completely ignored, leading to an unbalanced and visually unpleasant arrangement.", + "color_scheme": "Oversaturated, garish colors that clash violently. Color banding is visible in gradient areas. The overall palette feels artificial and digitally processed rather than natural.", + "mood_atmosphere": "Unsettling, uncanny atmosphere that fails to evoke any intended emotional response. The scene feels lifeless and sterile despite attempting to portray dynamic action.", + "patterns": "Visible tiling artifacts in textures, moir\u00e9 patterns, and aliasing on edges." + }, + "cinematography": { + "camera_motion": "Extremely shaky, unstable camera with visible rolling shutter artifacts. The motion is jerky and discontinuous, causing motion sickness and making the scene impossible to follow.", + "framing": "Poorly framed shots that cut off important elements and include unnecessary empty space.", + "camera_angle": "Awkward, disorienting camera angles that provide no useful spatial information about the scene. The camera path exhibits visible mathematical artifacts suggesting simple interpolation between keyframes rather than natural camera operation. Camera motion is completely disconnected from the scene content \u2014 panning away from action, dollying during dialogue, and shaking during still moments. The camera appears to pass through solid objects occasionally. Zoom is applied digitally rather than optically, revealing progressively worse resolution. Camera motion exhibits non-physical acceleration profiles \u2014 instant starts and stops rather than smooth ease-in/ease-out. Rolling shutter simulation is applied inconsistently, present in some frames but not others. The camera occasionally exhibits impossible motion like teleporting between positions. Virtual camera stabilization creates an uncanny floating sensation disconnected from any physical camera rig.", + "depth_of_field": "Uniform focus throughout, creating a flat, documentary-like appearance with no cinematic depth separation.", + "focus": "Soft, out-of-focus imagery with visible chromatic aberration and lens distortion that was not corrected in post-processing.", + "lens_focal_length": "Inappropriate focal length causing barrel distortion and unnatural perspective compression." + }, + "style_medium": "Low quality compressed digital video with visible encoding artifacts", + "artistic_style": "Amateur, unpolished with inconsistent visual style", + "context": "A poorly produced video with numerous technical and artistic flaws that detract from any intended narrative or visual impact.", + "actions": [ + { + "time": "0:00-0:08", + "description": "Subjects attempt to move but their motion is jerky, temporally inconsistent, and physically implausible. Background elements flicker and shift between frames." + } + ], + "text_and_signage_elements": [], + "segments": [ + { + "segment_index": 0, + "time_range": "0:00-0:08", + "description": "A single continuous shot suffering from severe temporal inconsistencies \u2014 subjects that morph and deform between frames, backgrounds that shift and wobble, and rendering quality that fluctuates visibly over time. Motion blur is applied incorrectly, smearing in directions that don't match actual movement. Frame-to-frame coherence breaks down with individual pixels changing color randomly in flat areas. Texture detail level fluctuates between frames as if the rendering budget varied shot to shot. Color grading drifts over the duration with no creative motivation. Noise patterns change between frames in ways that draw attention rather than being invisible. Overall visual quality degrades progressively from start to finish.", + "key_changes": "No meaningful progression or narrative development. Visual quality degrades over time.", + "camera": "Unstable, poorly controlled camera work with visible mathematical interpolation artifacts." + } + ], + "transitions": [], + "temporal_caption": "The scene opens at 0.0 seconds with a poorly rendered establishing shot that immediately reveals low production quality. At 1.0 seconds, subjects begin to move but their motion is jerky and inconsistent, with limbs bending at unnatural angles and objects clipping through each other. From 2.0 to 4.0 seconds, the camera shakes violently while the scene exhibits visible compression artifacts, color banding in the sky, and flickering in the shadows. Between 4.0 and 6.0 seconds, temporal coherence breaks down as elements appear and disappear between frames, textures swim and morph unnaturally, and the lighting shifts abruptly without physical cause. In the final 2 seconds, the overall visual quality deteriorates further with increasing noise, blur, and a general loss of spatial coherence that makes the scene nearly unwatchable. Additionally, the frame rate appears inconsistent with visible judder and stuttering throughout. Color temperature shifts randomly between warm and cool tones with no motivation. The encode quality degrades in complex regions showing macro-blocking and mosquito noise around moving edges. Temporal noise patterns are spatially correlated, creating swimming artifacts on flat surfaces.", + "audio_description": "", + "physical_realism": "No adherence to physical laws. Objects defy gravity, pass through solid surfaces, and change mass and momentum without cause. Fluid dynamics, cloth simulation, and rigid body physics are all fundamentally broken. Furthermore, conservation of energy is violated as objects gain or lose kinetic energy spontaneously. Elastic collisions produce inelastic results and vice versa. Surface friction is inconsistent \u2014 objects slide on rough surfaces while sticking to smooth ones. Air resistance appears to affect only some objects while others move through the atmosphere unimpeded." + } diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py index 796de06f1..a0eab3b5e 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py @@ -1019,11 +1019,19 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": True})) assert cfg is not None assert list(cfg.hints) == ["edge"] + assert cfg.hints["edge"].control_weight == 1.0 + assert cfg.normalized_control_weights == [1.0] assert cfg.guidance_scale == 3.0 assert cfg.control_guidance == 1.5 assert cfg.flow_shift == 10.0 assert cfg.num_video_frames_per_chunk == 93 assert cfg.share_vision_temporal_positions is True + assert cfg.emphasize_control_in_prompt is True + prompt_ablation_cfg = transfer.resolve_transfer_config( + make_sampling_params(extra_args={"edge": True, "emphasize_control_in_prompt": False}) + ) + assert prompt_ablation_cfg is not None + assert prompt_ablation_cfg.emphasize_control_in_prompt is False # fps omitted (no fps/frame_rate on the sampling params) -> wsm preset default (10) applies. defaulted_fps_cfg = transfer.resolve_transfer_config(make_sampling_params(extra_args={"wsm": True})) assert defaulted_fps_cfg is not None @@ -1092,6 +1100,40 @@ def test_transfer_config_media_helpers_and_preprocess_budget(monkeypatch: pytest assert "preprocessed_video" not in additional +def test_transfer_control_weight_validation_and_normalization() -> None: + from vllm_omni.diffusion.models.cosmos3 import transfer + + single = transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": 0.4}})) + assert single is not None + assert single.hints["edge"].control_weight == 0.4 + assert single.normalized_control_weights == [1.0] + + multiple = transfer.resolve_transfer_config( + make_sampling_params( + extra_args={ + "edge": {"control_weight": 1.0}, + "depth": {"control_weight": 3.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)}, + } + ) + ) + assert multiple is not None + assert multiple.normalized_control_weights == [0.25, 0.75] + + with pytest.raises(ValueError, match="Unsupported.*weight"): + transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"weight": 0.5}})) + with pytest.raises(ValueError, match="finite and non-negative"): + transfer.resolve_transfer_config(make_sampling_params(extra_args={"edge": {"control_weight": -0.5}})) + with pytest.raises(ValueError, match="positive sum"): + transfer.resolve_transfer_config( + make_sampling_params( + extra_args={ + "edge": {"control_weight": 0.0}, + "depth": {"control_weight": 0.0, "control": torch.zeros(3, 1, 8, 8, dtype=torch.uint8)}, + } + ) + ) + + def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None: """When fps and frame_rate differ, transfer must pick frame_rate -- the same precedence as OmniDiffusionSamplingParams.resolved_frame_rate. Uses the real @@ -1352,6 +1394,63 @@ def test_format_and_tokenize_prompts_applies_video_templates_and_system_override ] +def test_format_and_tokenize_prompts_applies_transfer_prompt_contract(make_cosmos3_pipeline) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE, + COSMOS3_TRANSFER_SYSTEM_PROMPT, + ) + + pipeline = make_cosmos3_pipeline() + calls = _capture_tokenize_calls(pipeline) + directive = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names="edge") + + 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={ + "use_resolution_template": True, + "system_prompt": "request-level system prompt", + } + ), + use_system_prompt=True, + is_t2i=False, + system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT, + prompt_suffix=directive, + use_duration_template=True, + use_resolution_template=True, + negative_metadata_mode="same", + ) + + metadata = "The video is 2.0 seconds long and is of 24 FPS. This video is of 720x1280 resolution." + assert calls[0]["text"] == f"A robot. {metadata} {directive}" + assert calls[1]["text"] == f"bad. {metadata}" + assert all(call["use_system_prompt"] is True for call in calls) + assert all(call["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT for call in calls) + + +def test_format_and_tokenize_prompts_rejects_unknown_negative_metadata_mode(make_cosmos3_pipeline) -> None: + pipeline = make_cosmos3_pipeline() + + with pytest.raises(ValueError, match="negative_metadata_mode"): + 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="unexpected", + ) + + def test_format_and_tokenize_prompts_uses_image_templates_for_t2i(make_cosmos3_pipeline) -> None: from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import COSMOS3_T2I_SYSTEM_PROMPT @@ -1740,6 +1839,7 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_ control_guidance=1.5, control_guidance_interval=None, control_latents=[torch.zeros_like(latents)], + control_weights=[1.0], shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask}, velocity_mask=velocity_mask, condition_latents=torch.zeros_like(latents), @@ -1750,6 +1850,9 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_ (2, False), (1, True), ] + assert pipeline.transformer.calls[0]["kwargs"]["control_weights"] == [1.0] + assert "control_weights" not in pipeline.transformer.calls[1]["kwargs"] + assert pipeline.transformer.calls[2]["kwargs"]["control_weights"] == [1.0] torch.testing.assert_close(result, torch.full_like(latents, 254.0)) @@ -1857,14 +1960,35 @@ def test_diffuse_transfer_interval_switches_branch_counts(make_cosmos3_pipeline, torch.testing.assert_close(result, torch.full_like(latents, 508.0)) -@pytest.mark.parametrize(("hint_key", "expected_fps"), [("edge", 8.0), ("wsm", 10.0)]) -def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint_key: str, expected_fps: float) -> None: +@pytest.mark.parametrize( + ("hint_key", "emphasize_control", "expected_fps", "negative_prompt"), + [ + ("edge", None, 8.0, None), + ("wsm", None, 10.0, ""), + ("edge", False, 8.0, "custom negative"), + ], +) +def test_forward_transfer_uses_transfer_prompt_contract_and_source_fps_except_wsm( + make_cosmos3_pipeline, + hint_key: str, + emphasize_control: bool | None, + expected_fps: float, + negative_prompt: str | None, +) -> None: + from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE, + COSMOS3_TRANSFER_SYSTEM_PROMPT, + ) + pipeline = make_cosmos3_pipeline() captured: dict[str, Any] = {} def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): - del prompt, negative_prompt, num_frames, height, width, args, kwargs + del num_frames, height, width, args + captured["prompt"] = prompt + captured["negative_prompt"] = negative_prompt captured["format_frame_rate"] = frame_rate + captured["format_kwargs"] = kwargs return _ids(2), _mask(), _ids(1), _mask() def fake_encode(video: torch.Tensor) -> torch.Tensor: @@ -1901,33 +2025,48 @@ def test_forward_transfer_uses_source_fps_except_wsm(make_cosmos3_pipeline, hint pipeline._decode_latents = lambda latents: torch.zeros(1, 3, 5, 16, 16, device="meta") control = torch.zeros(3, 5, 16, 16, dtype=torch.uint8) + extra_args = { + hint_key: {"control": control}, + "max_frames": 5, + "num_video_frames_per_chunk": 5, + "show_control_condition": True, + } + if emphasize_control is not None: + extra_args["emphasize_control_in_prompt"] = emphasize_control + prompt_data = { + "prompt": "transfer", + "modalities": ["video"], + "additional_information": { + "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16), + "transfer_input_fps": 8.0, + }, + } + if negative_prompt is not None: + prompt_data["negative_prompt"] = negative_prompt request = SimpleNamespace( - prompts=[ - { - "prompt": "transfer", - "modalities": ["video"], - "additional_information": { - "preprocessed_transfer_video": torch.zeros(1, 3, 5, 16, 16), - "transfer_input_fps": 8.0, - }, - } - ], + prompts=[prompt_data], sampling_params=make_sampling_params( height=16, width=16, # fps omitted (no frame_rate) -> non-wsm uses the source video fps (8), wsm uses its preset (10). - extra_args={ - hint_key: {"control": control}, - "max_frames": 5, - "num_video_frames_per_chunk": 5, - "show_control_condition": True, - }, + extra_args=extra_args, ), ) output = pipeline.forward(request) assert captured["format_frame_rate"] == expected_fps + expected_negative_prompt = "" if negative_prompt is None else negative_prompt + assert captured["negative_prompt"] == expected_negative_prompt + assert captured["format_kwargs"]["use_system_prompt"] is True + assert captured["format_kwargs"]["system_prompt"] == COSMOS3_TRANSFER_SYSTEM_PROMPT + assert captured["format_kwargs"]["use_duration_template"] is True + assert captured["format_kwargs"]["use_resolution_template"] is True + assert captured["format_kwargs"]["negative_metadata_mode"] == "same" + expected_suffix = ( + None if emphasize_control is False else COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_key) + ) + assert captured["format_kwargs"]["prompt_suffix"] == expected_suffix assert captured["shared_kwargs"]["fps"] == expected_fps assert captured["flow_shifts"] == [10.0] # Transfer applies the V2V flow shift when building its timestep schedule. diff --git a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py index 5dd8a875a..2bfa7b558 100644 --- a/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py +++ b/tests/diffusion/models/cosmos3/test_cosmos3_transformer.py @@ -504,6 +504,35 @@ def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device( assert generator_param.device == accelerator_device +def test_multi_control_attention_weights_target_outputs() -> None: + from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention + + class ControlValueAttention: + def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + del k + # Each per-control K/V layout is [text, control_i, target]. + return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1) + + attention = SimpleNamespace(multi_control_attn=ControlValueAttention()) + q = torch.zeros(1, 3, 1, 1) + k = torch.zeros_like(q) + v = torch.tensor([[[[10.0]], [[20.0]], [[30.0]]]]) + text_kv = torch.zeros(1, 1, 1, 1) + + output = Cosmos3CrossAttention._forward_multi_control( + attention, + q, + k, + v, + text_kv, + text_kv, + control_token_sizes=(1, 1), + control_weights=(0.25, 0.75), + ) + + torch.testing.assert_close(output, torch.tensor([[[10.0], [20.0], [17.5]]])) + + def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 @@ -521,9 +550,21 @@ def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatc video_shape=(1, 2, 2), fps=24.0, control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], + control_weights=[1.0, 3.0], ) assert tuple(output.shape) == tuple(hidden_states.shape) + with pytest.raises(ValueError, match="control_weights length"): + model( + hidden_states=hidden_states, + timestep=torch.tensor([1.0]), + text_ids=torch.tensor([[1, 2]], dtype=torch.long), + text_mask=torch.ones(1, 2, dtype=torch.long), + video_shape=(1, 2, 2), + fps=24.0, + control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], + control_weights=[1.0], + ) with pytest.raises(ValueError, match="control latent shape"): model( hidden_states=hidden_states, diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 573f0bd79..b7f2decf9 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -151,6 +151,7 @@ def test_cosmos3_extra_registry_declares_request_and_response_params(pipeline_na "max_sequence_length", "use_resolution_template", "use_duration_template", + "negative_metadata_mode", "use_system_prompt", "system_prompt", "negative_prompt", diff --git a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py index 97cf92e26..2304c0212 100644 --- a/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py @@ -163,6 +163,14 @@ COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{widt # NOTE: Intentional typo in "give" instead of "given" to match training setup. COSMOS3_SYSTEM_PROMPT = "You are a helpful assistant who will generate videos from a give prompt." COSMOS3_T2I_SYSTEM_PROMPT = "You are a helpful assistant who will generate images from a give prompt." +COSMOS3_TRANSFER_SYSTEM_PROMPT = ( + "You are a helpful assistant that generates images or videos following the user's instructions" + " and control signals (edge maps, blur, depth, or segmentation)." +) +COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE = ( + "Follow the {hint_names} control video precisely: shape, contour, silhouette, position, and motion of every " + "visible structure must align with the {hint_names} signal at every frame." +) COSMOS3_T2V_DEFAULT_HEIGHT = 720 COSMOS3_T2V_DEFAULT_WIDTH = 1280 @@ -2068,6 +2076,11 @@ class Cosmos3OmniDiffusersPipeline( sp: OmniDiffusionSamplingParams, use_system_prompt: bool = False, is_t2i: bool = False, + system_prompt: str | None = None, + prompt_suffix: str | None = None, + use_duration_template: bool | None = None, + use_resolution_template: bool | None = None, + negative_metadata_mode: str | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Format prompts with metadata templates and tokenize. @@ -2075,14 +2088,30 @@ class Cosmos3OmniDiffusersPipeline( For T2I (``is_t2i=True``) the duration template is suppressed (no FPS or duration concept for a single image) and the image-flavored - resolution template is used. + resolution template is used. ``prompt_suffix`` is appended only to the + positive prompt, after metadata formatting. ``negative_metadata_mode`` + controls whether the negative branch receives no metadata, the same + metadata as the positive branch, or inverse metadata. """ # Route cosmos3-specific controls through ``_get_sp_param`` so they # are picked up from ``extra_args`` (OpenAI endpoint path) as well # as from direct attributes. - use_duration_template = bool(self._get_sp_param(sp, "use_duration_template", False)) and not is_t2i + if use_duration_template is None: + use_duration_template = bool(self._get_sp_param(sp, "use_duration_template", False)) + if use_resolution_template is None: + use_resolution_template = bool(self._get_sp_param(sp, "use_resolution_template", False)) + if negative_metadata_mode is None: + negative_metadata_mode = str(self._get_sp_param(sp, "negative_metadata_mode", "inverse")) + negative_metadata_mode = negative_metadata_mode.strip().lower() + if negative_metadata_mode not in {"none", "same", "inverse"}: + raise ValueError( + "Cosmos3 negative_metadata_mode must be one of 'none', 'same', or 'inverse'; " + f"got {negative_metadata_mode!r}." + ) + + use_duration_template = use_duration_template and not is_t2i dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None - if bool(self._get_sp_param(sp, "use_resolution_template", False)): + if use_resolution_template: res_tmpl = COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_RESOLUTION_TEMPLATE else: res_tmpl = None @@ -2106,37 +2135,44 @@ class Cosmos3OmniDiffusersPipeline( duration_template=dur_tmpl, resolution_template=res_tmpl, ) + if prompt_suffix: + prompt = f"{prompt.rstrip()} {prompt_suffix.lstrip()}".strip() if _is_rank_zero(): logger.info("Final prompt: '%s'", prompt) - # Negative prompt: inverse templates ("not {duration}...", "not {height}x{width}..."). - # Applied whenever the matching positive template is enabled; an empty - # negative_prompt yields output that starts with the template, not a dot. - inv_dur = COSMOS3_INVERSE_DURATION_TEMPLATE if dur_tmpl else None - if res_tmpl is None: - inv_res = None - elif is_t2i: - inv_res = COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE + if negative_metadata_mode == "none": + negative_dur_tmpl = None + negative_res_tmpl = None + elif negative_metadata_mode == "same": + negative_dur_tmpl = dur_tmpl + negative_res_tmpl = res_tmpl else: - inv_res = COSMOS3_INVERSE_RESOLUTION_TEMPLATE + negative_dur_tmpl = COSMOS3_INVERSE_DURATION_TEMPLATE if dur_tmpl else None + if res_tmpl is None: + negative_res_tmpl = None + elif is_t2i: + negative_res_tmpl = COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE + else: + negative_res_tmpl = COSMOS3_INVERSE_RESOLUTION_TEMPLATE negative_prompt = self._apply_metadata_templates( negative_prompt, num_frames, frame_rate, height, width, - duration_template=inv_dur, - resolution_template=inv_res, - force_duration_template=True, + duration_template=negative_dur_tmpl, + resolution_template=negative_res_tmpl, + force_duration_template=negative_metadata_mode == "inverse", ) - default_sys_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_SYSTEM_PROMPT - sys_prompt = self._get_sp_param(sp, "system_prompt", default_sys_prompt) or default_sys_prompt + if system_prompt is None: + default_sys_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_SYSTEM_PROMPT + system_prompt = self._get_sp_param(sp, "system_prompt", default_sys_prompt) or default_sys_prompt cond_ids, cond_mask = self._tokenize_prompt( - prompt, max_sequence_length, use_system_prompt, system_prompt=sys_prompt + prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt ) uncond_ids, uncond_mask = self._tokenize_prompt( - negative_prompt, max_sequence_length, use_system_prompt, system_prompt=sys_prompt + negative_prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt ) return cond_ids, cond_mask, uncond_ids, uncond_mask @@ -2890,6 +2926,7 @@ class Cosmos3OmniDiffusersPipeline( control_latents: list[torch.Tensor], shared_kwargs: dict[str, Any], *, + control_weights: list[float] | None = None, velocity_mask: torch.Tensor, condition_latents: torch.Tensor, guidance_interval: tuple[float, float] | None = None, @@ -2925,6 +2962,7 @@ class Cosmos3OmniDiffusersPipeline( text_ids=cond_ids, text_mask=cond_mask, control_latents=control_latents, + control_weights=control_weights, _cosmos3_cache_key="transfer_cond_full", **shared_kwargs, ) @@ -2946,6 +2984,7 @@ class Cosmos3OmniDiffusersPipeline( text_ids=uncond_ids, text_mask=uncond_mask, control_latents=control_latents, + control_weights=control_weights, _cosmos3_cache_key="transfer_uncond_full", **shared_kwargs, ), @@ -2993,6 +3032,7 @@ class Cosmos3OmniDiffusersPipeline( text_ids=uncond_ids, text_mask=uncond_mask, control_latents=control_latents, + control_weights=control_weights, _cosmos3_cache_key="transfer_uncond_full", **shared_kwargs, ), @@ -3113,7 +3153,10 @@ class Cosmos3OmniDiffusersPipeline( self._get_sp_param(sp, "max_sequence_length", COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH) or COSMOS3_DEFAULT_MAX_SEQUENCE_LENGTH ) - use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False)) + transfer_prompt_suffix = None + if transfer_config.emphasize_control_in_prompt: + hint_names = ", ".join(hint.key for hint in transfer_config.ordered_hints) + transfer_prompt_suffix = COSMOS3_TRANSFER_CONTROL_DIRECTIVE_TEMPLATE.format(hint_names=hint_names) self._guidance_scale = guidance_scale self._num_timesteps = num_inference_steps @@ -3124,6 +3167,13 @@ class Cosmos3OmniDiffusersPipeline( if generator is None: generator = torch.Generator(device=self.device).manual_seed(seed) + for ignored_key in ("system_prompt", "use_system_prompt"): + if self._get_sp_param(sp, ignored_key, None) is not None: + logger.warning_once( + "Cosmos3 transfer ignores the '%s' request parameter and always " + "tokenizes both CFG branches with the transfer-specific system prompt.", + ignored_key, + ) cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( prompt, negative_prompt, @@ -3133,8 +3183,13 @@ class Cosmos3OmniDiffusersPipeline( width, max_sequence_length, sp, - use_system_prompt, + use_system_prompt=True, is_t2i=False, + system_prompt=COSMOS3_TRANSFER_SYSTEM_PROMPT, + prompt_suffix=transfer_prompt_suffix, + use_duration_template=bool(self._get_sp_param(sp, "use_duration_template", True)), + use_resolution_template=bool(self._get_sp_param(sp, "use_resolution_template", True)), + negative_metadata_mode=str(self._get_sp_param(sp, "negative_metadata_mode", "same")), ) output_chunks: list[torch.Tensor] = [] @@ -3227,6 +3282,7 @@ class Cosmos3OmniDiffusersPipeline( control_guidance=transfer_config.control_guidance, control_guidance_interval=transfer_config.control_guidance_interval, control_latents=control_latents, + control_weights=transfer_config.normalized_control_weights, shared_kwargs=shared_kwargs, velocity_mask=velocity_mask, condition_latents=condition_latents, diff --git a/vllm_omni/diffusion/models/cosmos3/transfer.py b/vllm_omni/diffusion/models/cosmos3/transfer.py index 42cbea559..9c8a52f3f 100644 --- a/vllm_omni/diffusion/models/cosmos3/transfer.py +++ b/vllm_omni/diffusion/models/cosmos3/transfer.py @@ -14,6 +14,7 @@ generation, or action generation. from __future__ import annotations import importlib +import math from collections.abc import Mapping from dataclasses import dataclass, field from pathlib import Path @@ -25,6 +26,14 @@ import torch import torch.nn.functional as F TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm") +_TRANSFER_HINT_COMMON_FIELDS = frozenset({"control_path", "control", "control_weight"}) +_TRANSFER_HINT_FIELDS: dict[str, frozenset[str]] = { + "edge": _TRANSFER_HINT_COMMON_FIELDS | {"preset_edge_threshold"}, + "blur": _TRANSFER_HINT_COMMON_FIELDS | {"preset_blur_strength"}, + "depth": _TRANSFER_HINT_COMMON_FIELDS, + "seg": _TRANSFER_HINT_COMMON_FIELDS, + "wsm": _TRANSFER_HINT_COMMON_FIELDS, +} TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { "num_video_frames_per_chunk": 93, "num_conditional_frames": 1, @@ -33,6 +42,7 @@ TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { "show_input": False, "num_first_chunk_conditional_frames": 0, "share_vision_temporal_positions": True, + "emphasize_control_in_prompt": True, } TRANSFER_DEFAULTS: dict[str, dict[str, Any]] = { "edge": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, @@ -85,6 +95,7 @@ class Cosmos3TransferHint: key: str control_path: str | None = None control: Any | None = None + control_weight: float = 1.0 preset_edge_threshold: str = "medium" preset_blur_strength: str = "medium" @@ -103,6 +114,7 @@ class Cosmos3TransferConfig: show_input: bool = False num_first_chunk_conditional_frames: int = 0 share_vision_temporal_positions: bool = True + emphasize_control_in_prompt: bool = True num_frames: int | None = None fps: float | None = None @@ -110,6 +122,16 @@ class Cosmos3TransferConfig: def ordered_hints(self) -> list[Cosmos3TransferHint]: return [self.hints[key] for key in TRANSFER_HINT_KEYS if key in self.hints] + @property + def normalized_control_weights(self) -> list[float]: + weights = [hint.control_weight for hint in self.ordered_hints] + if any(not math.isfinite(weight) or weight < 0.0 for weight in weights): + raise ValueError(f"Cosmos3 transfer control_weight values must be finite and non-negative, got {weights}.") + total = sum(weights) + if total <= 0.0: + raise ValueError("Cosmos3 transfer control_weight values must have a positive sum.") + return [weight / total for weight in weights] + def _extra_args(sp: Any) -> Mapping[str, Any]: extra = getattr(sp, "extra_args", None) @@ -184,7 +206,7 @@ def _as_interval(value: Any) -> tuple[float, float] | None: return None if isinstance(value, str): value = [item.strip() for item in value.split(",") if item.strip()] - if not isinstance(value, (list, tuple)) or len(value) != 2: + if not isinstance(value, list | tuple) or len(value) != 2: raise ValueError("Cosmos3 transfer control_guidance_interval must contain exactly two values.") lo, hi = float(value[0]), float(value[1]) if lo > hi: @@ -220,10 +242,27 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer raise TypeError( f"Cosmos3 transfer hint '{key}' must be an object, path string, or true; got {type(raw)!r}." ) + unknown_fields = set(raw) - _TRANSFER_HINT_FIELDS[key] + if unknown_fields: + raise ValueError( + f"Unsupported Cosmos3 transfer hint '{key}' fields: {sorted(unknown_fields)}. " + f"Supported fields are: {sorted(_TRANSFER_HINT_FIELDS[key])}." + ) + try: + control_weight = float(raw.get("control_weight", 1.0)) + except (TypeError, ValueError) as exc: + raise TypeError( + f"Cosmos3 transfer hint '{key}' control_weight must be a finite non-negative number." + ) from exc + if not math.isfinite(control_weight) or control_weight < 0.0: + raise ValueError( + f"Cosmos3 transfer hint '{key}' control_weight must be finite and non-negative, got {control_weight!r}." + ) hints[key] = Cosmos3TransferHint( key=key, control_path=str(raw["control_path"]) if raw.get("control_path") is not None else None, control=raw.get("control"), + control_weight=control_weight, preset_edge_threshold=str(raw.get("preset_edge_threshold") or "medium").lower(), preset_blur_strength=str(raw.get("preset_blur_strength") or "medium").lower(), ) @@ -239,6 +278,7 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer "show_input", "num_first_chunk_conditional_frames", "share_vision_temporal_positions", + "emphasize_control_in_prompt", ) if any(_is_user_field(extra, sp, prompt_data, key) for key in transfer_only): raise ValueError("Cosmos3 transfer options were provided, but no transfer hint was selected.") @@ -293,6 +333,10 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer _param(extra, sp, prompt_data, "share_vision_temporal_positions", None), bool(TRANSFER_SAMPLE_DEFAULTS["share_vision_temporal_positions"]), ), + emphasize_control_in_prompt=_as_bool( + _param(extra, sp, prompt_data, "emphasize_control_in_prompt", None), + bool(TRANSFER_SAMPLE_DEFAULTS["emphasize_control_in_prompt"]), + ), num_frames=( int(_param(extra, sp, prompt_data, "num_frames")) if _param(extra, sp, prompt_data, "num_frames", None) is not None @@ -328,6 +372,7 @@ def resolve_transfer_config(sp: Any, prompt_data: Any = None) -> Cosmos3Transfer raise ValueError(f"Unsupported Cosmos3 edge preset: {hint.preset_edge_threshold!r}.") if hint.key == "blur" and hint.preset_blur_strength not in BLUR_DOWNUP_PRESETS: raise ValueError(f"Unsupported Cosmos3 blur preset: {hint.preset_blur_strength!r}.") + _ = config.normalized_control_weights return config diff --git a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py index cea5a32ff..c946eb4ca 100644 --- a/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py +++ b/vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py @@ -651,6 +651,19 @@ class Cosmos3CrossAttention(nn.Module): softmax_scale=1.0 / (self.head_dim**0.5), num_kv_heads=self.num_kv_heads, ) + # Multi-control attention operates on one full [control_i, target] + # sequence at a time. Keep those sequences replicated when Ulysses is + # active; sharding the concatenated [control_1, ..., control_N, target] + # layout would split the per-control ranges across ranks. + self.multi_control_attn = FrameworkAttention( + num_heads=self.num_heads, + head_size=self.head_dim, + causal=False, + softmax_scale=1.0 / (self.head_dim**0.5), + num_kv_heads=self.num_kv_heads, + prefix=f"{prefix}.multi_control_attn", + skip_sequence_parallel=True, + ) # TODO(follow-up): collapse _forward_local and _forward_sp into a single # joint-based path when NoParallelAttention can process joint_key/value. @@ -701,6 +714,61 @@ class Cosmos3CrossAttention(nn.Module): out = self.attn(q, k, v, attn_metadata) return out.reshape(B, S_gen, -1) + def _forward_multi_control( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + control_token_sizes: tuple[int, ...], + control_weights: tuple[float, ...], + ) -> torch.Tensor: + """Run independent attention per control and weight target outputs. + + For control ``i``, both its queries and the target queries attend to + ``[text, control_i, target]``. Each control keeps the output from its + own pass, while target outputs are combined using normalized weights. + """ + if len(control_token_sizes) != len(control_weights): + raise ValueError( + "Cosmos3 control token sizes and weights must have the same length: " + f"sizes={len(control_token_sizes)}, weights={len(control_weights)}." + ) + control_tokens = sum(control_token_sizes) + if control_tokens <= 0 or control_tokens >= q.shape[1]: + raise ValueError( + "Cosmos3 multi-control attention requires non-empty control and target token ranges, " + f"got control_tokens={control_tokens}, total_tokens={q.shape[1]}." + ) + + q_target = q[:, control_tokens:] + k_target = k[:, control_tokens:] + v_target = v[:, control_tokens:] + control_outputs: list[torch.Tensor] = [] + target_output: torch.Tensor | None = None + start = 0 + for size, weight in zip(control_token_sizes, control_weights, strict=True): + if size <= 0: + raise ValueError(f"Cosmos3 control token sizes must be positive, got {control_token_sizes}.") + end = start + size + q_control = q[:, start:end] + k_control = k[:, start:end] + v_control = v[:, start:end] + q_pair = torch.cat([q_control, q_target], dim=1) + k_pair = torch.cat([k_und, k_control, k_target], dim=1) + v_pair = torch.cat([v_und, v_control, v_target], dim=1) + pair_output = self.multi_control_attn(q_pair, k_pair, v_pair) + control_outputs.append(pair_output[:, :size]) + weighted_target = pair_output[:, size:] * weight + target_output = weighted_target if target_output is None else target_output + weighted_target + start = end + + if start != control_tokens or target_output is None: + raise RuntimeError("Cosmos3 multi-control attention failed to cover all control token ranges.") + output = torch.cat([*control_outputs, target_output], dim=1) + return output.reshape(q.shape[0], q.shape[1], -1) + # -- Public forward: routes to the appropriate path ---------------------- def forward( @@ -710,6 +778,8 @@ class Cosmos3CrossAttention(nn.Module): v_und: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, + control_token_sizes: tuple[int, ...] | None = None, + control_weights: tuple[float, ...] | None = None, ) -> torch.Tensor: """ Args: @@ -733,7 +803,19 @@ class Cosmos3CrossAttention(nn.Module): # Qwen3-style RoPE q, k = _apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) - if _is_sp_active(): + if control_token_sizes is not None or control_weights is not None: + if control_token_sizes is None or control_weights is None: + raise ValueError("Cosmos3 multi-control attention requires both control token sizes and weights.") + out = self._forward_multi_control( + q, + k, + v, + k_und, + v_und, + control_token_sizes, + control_weights, + ) + elif _is_sp_active(): out = self._forward_sp(q, k, v, k_und, v_und) else: out = self._forward_local(q, k, v, k_und, v_und) @@ -847,6 +929,8 @@ class Cosmos3GenDecoderLayer(nn.Module): freqs_sin: torch.Tensor | None = None, cached_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None, freqs_gen: tuple[torch.Tensor, torch.Tensor] | None = None, + control_token_sizes: tuple[int, ...] | None = None, + control_weights: tuple[float, ...] | None = None, ) -> torch.Tensor: if cached_kv is not None: if self.layer_idx is None: @@ -861,7 +945,13 @@ class Cosmos3GenDecoderLayer(nn.Module): hidden_states = self.input_layernorm(hidden_states) hidden_states = self.cross_attention( - hidden_states, k_und=k_und, v_und=v_und, freqs_cos=freqs_cos, freqs_sin=freqs_sin + hidden_states, + k_und=k_und, + v_und=v_und, + freqs_cos=freqs_cos, + freqs_sin=freqs_sin, + control_token_sizes=control_token_sizes, + control_weights=control_weights, ) hidden_states = residual + hidden_states @@ -1579,6 +1669,7 @@ class Cosmos3VFMTransformer(nn.Module): sound_latents: torch.Tensor | None = None, noisy_frame_mask: torch.Tensor | None = None, control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + control_weights: list[float] | tuple[float, ...] | torch.Tensor | None = None, transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> torch.Tensor | tuple[torch.Tensor, ...]: @@ -1601,6 +1692,8 @@ class Cosmos3VFMTransformer(nn.Module): frames are noisy, as in T2I/T2V. control_latents: Optional transfer-control latents. Controls are clean vision context and are packed before the noisy target. + control_weights: Optional non-negative relative weights for + transfer controls. Values are normalized to sum to one. Returns: [B, C, t, h, w] velocity prediction, or @@ -1629,6 +1722,31 @@ class Cosmos3VFMTransformer(nn.Module): else: control_latent_list = list(control_latents) has_control = len(control_latent_list) > 0 + if control_weights is None: + normalized_control_weights = ( + [1.0 / len(control_latent_list)] * len(control_latent_list) if has_control else [] + ) + else: + if isinstance(control_weights, torch.Tensor): + raw_control_weights = control_weights.detach().cpu().flatten().tolist() + else: + raw_control_weights = list(control_weights) + if len(raw_control_weights) != len(control_latent_list): + raise ValueError( + "Cosmos3 transfer control_weights length must match control_latents: " + f"weights={len(raw_control_weights)}, controls={len(control_latent_list)}." + ) + normalized_control_weights = [float(weight) for weight in raw_control_weights] + if any(not math.isfinite(weight) or weight < 0.0 for weight in normalized_control_weights): + raise ValueError( + "Cosmos3 transfer control_weights must be finite and non-negative, " + f"got {normalized_control_weights}." + ) + control_weight_sum = sum(normalized_control_weights) + if control_weight_sum <= 0.0: + raise ValueError("Cosmos3 transfer control_weights must have a positive sum.") + normalized_control_weights = [weight / control_weight_sum for weight in normalized_control_weights] + use_multi_control_attention = len(control_latent_list) > 1 if has_control and (has_action or has_sound): raise ValueError("Cosmos3 transfer control latents cannot be combined with action or sound latents.") if has_action and not self.action_gen: @@ -1646,6 +1764,13 @@ class Cosmos3VFMTransformer(nn.Module): # Query Ulysses state at runtime ulysses_size, _, _ = _get_ulysses_state() + if use_multi_control_attention and ulysses_size > 1: + logger.warning_once( + "Cosmos3 multi-control attention requires full [control_i | target] " + "sequences, so the GEN pathway runs replicated on every Ulysses rank. " + "Sequence parallelism will not reduce per-rank memory or latency for " + "multi-control transfer requests." + ) # Pack action/sound tokens (no learned weights) up front so the UND # cache sizing knows their token lengths. The modality projections are @@ -1722,6 +1847,9 @@ class Cosmos3VFMTransformer(nn.Module): ) hidden_controls.append(hidden_control) s_control += hidden_control.shape[1] + control_token_sizes = tuple(control.shape[1] for control in hidden_controls) + multi_control_weights = tuple(normalized_control_weights) if use_multi_control_attention else None + multi_control_token_sizes = control_token_sizes if use_multi_control_attention else None hidden_action = None hidden_sound = None if action_tokens is not None: @@ -1788,10 +1916,11 @@ class Cosmos3VFMTransformer(nn.Module): has_action=has_action, has_sound=has_sound, has_control=has_control, - ulysses_size=ulysses_size, + ulysses_size=1 if use_multi_control_attention else ulysses_size, ) freqs_cos, freqs_sin = self.cached_freqs_gen - hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) + if not use_multi_control_attention: + hidden_gen, freqs_cos, freqs_sin = self.gen_sp_prepare(hidden_gen, freqs_cos, freqs_sin) freqs_gen = (freqs_cos, freqs_sin) if len(self.gen_layers) == len(self.cached_kv): @@ -1802,6 +1931,8 @@ class Cosmos3VFMTransformer(nn.Module): v_und=v_und, freqs_cos=freqs_cos, freqs_sin=freqs_sin, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, ) # Cache-dit's block wrapper may return a tuple; unwrap it. if isinstance(hidden_gen, tuple): @@ -1813,11 +1944,14 @@ class Cosmos3VFMTransformer(nn.Module): hidden_gen, cached_kv=self.cached_kv, freqs_gen=freqs_gen, + control_token_sizes=multi_control_token_sizes, + control_weights=multi_control_weights, ) if isinstance(hidden_gen, tuple): hidden_gen = hidden_gen[0] - hidden_gen = self.gen_sp_gather(hidden_gen) + if not use_multi_control_attention: + hidden_gen = self.gen_sp_gather(hidden_gen) # Final norm and project back to latent space hidden_gen = self.norm_moe_gen(hidden_gen) diff --git a/vllm_omni/model_extras/cosmos3.py b/vllm_omni/model_extras/cosmos3.py index d0ffa5182..01cecf94f 100644 --- a/vllm_omni/model_extras/cosmos3.py +++ b/vllm_omni/model_extras/cosmos3.py @@ -9,6 +9,7 @@ COSMOS3_EXTRA_BODY_PARAMS = frozenset( "max_sequence_length", "use_resolution_template", "use_duration_template", + "negative_metadata_mode", "use_system_prompt", "system_prompt", "negative_prompt",