Add Cosmos3 model support (#3454)

Signed-off-by: Maciej Bala <mbala@nvidia.com>
Signed-off-by: MaciejBalaNV <mbala@nvidia.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>

tests/entrypoints/openai_api/test_image_server.py

@@ -584,6 +584,7 @@ def test_generate_single_image(test_client):
     img_bytes = base64.b64decode(data["data"][0]["b64_json"])
     img = Image.open(io.BytesIO(img_bytes))
     assert img.size == (64, 64)  # Our mock returns 64x64 images
+    assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"]
 
 
 def test_generate_images_async_omni_sampling_params(async_omni_test_client):
@@ -1011,6 +1012,42 @@ def test_with_custom_parameters(test_client):
     assert response.status_code == 200
 
 
+def test_flow_shift_forwarded_to_extra_args(test_client):
+    """flow_shift must reach the diffusion sampling params via extra_args.
+
+    Regression: ``ImageGenerationRequest`` had no ``flow_shift`` field and the
+    single-stage handler never forwarded it, so a request like
+    ``{"flow_shift": 10.0}`` was silently dropped and Cosmos3 T2I always ran at
+    its hardcoded per-mode default shift. The pipeline reads
+    ``extra_args["flow_shift"]`` (via ``_get_sp_param``), so it must land there.
+    """
+    response = test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a robot in a lab",
+            "size": "960x960",
+            "num_inference_steps": 50,
+            "guidance_scale": 4.0,
+            "flow_shift": 10.0,
+        },
+    )
+    assert response.status_code == 200
+    captured = test_client.app.state.engine_client.captured_sampling_params_list[0]
+    assert captured.extra_args["flow_shift"] == 10.0
+
+
+def test_flow_shift_absent_when_not_requested(test_client):
+    """Omitting flow_shift must not inject an override, so the pipeline keeps
+    its per-mode default (e.g. Cosmos3 T2I shift=3.0)."""
+    response = test_client.post(
+        "/v1/images/generations",
+        json={"prompt": "a tree", "size": "1024x1024"},
+    )
+    assert response.status_code == 200
+    captured = test_client.app.state.engine_client.captured_sampling_params_list[0]
+    assert "flow_shift" not in (captured.extra_args or {})
+
+
 def test_invalid_size(test_client):
     """Test with invalid size parameter - rejected by Pydantic"""
     response = test_client.post(