Add pipeline tag, library name, and paper link

#13
by nielsr HF Staff - opened
README.md CHANGED
@@ -1,463 +1,52 @@
1
  ---
 
2
  license: other
3
  license_name: openmdw1.1-license
4
- license_link: >-
5
- https://openmdw.ai/license/1-1/
6
- library_name: cosmos
7
  tags:
8
- - nvidia
9
- - cosmos
10
- - cosmos3
11
- - vllm-omni
12
- - sglang
13
- - sglang-diffusion
14
- - diffusers
15
- - text-to-image
16
- - image-generation
17
  ---
18
 
19
  # **Cosmos 3: Omnimodal World Models for Physical AI**
20
- **[Model Collection](https://huggingface.co/collections/nvidia/cosmos3)** | **[Code](https://github.com/nvidia/cosmos)** | **[White Paper](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf)** | **[Website](https://research.nvidia.com/labs/cosmos-lab/cosmos3/)**
21
 
22
- [NVIDIA Cosmos™](https://github.com/nvidia/cosmos) is a world foundation model platform designed to accelerate the development of Physical AI by enabling machines to understand, simulate, and interact with the physical world across robotics, autonomous driving, and smart space environments, including industrial and factory-scale applications.
 
 
23
 
24
  # Model Overview: Cosmos3-Super-Text2Image
25
 
26
  ## Description
27
 
28
- Cosmos3 is a collection of Omnimodal world models capable of generating dynamic, high-quality video, image, audio, and action commands from combinations of text, image, video, and action trajectory inputs. It serves as a foundational building block for a broad range of Physical AI applications and research spanning world understanding, world generation, simulation, and embodied policy learning.
29
-
30
- This model is ready for commercial and non-commercial use.
31
-
32
- **Model Developer:** NVIDIA
33
-
34
- ### Model Versions
35
- - Cosmos3-Nano:
36
- - Given multimodal inputs including text, images, video, audio, and action trajectories, generate coherent text, images, video, audio, and action outputs for multimodal understanding, world simulation, future prediction, action reasoning, and Physical AI applications.
37
-
38
- - Cosmos3-Super:
39
- - Given multimodal inputs including text, images, video, audio, and action trajectories, generate coherent text, images, video, audio, and action outputs for multimodal understanding, world simulation, future prediction, action reasoning, and Physical AI applications.
40
-
41
- - Cosmos3-Nano-Policy-DROID:
42
- - Given language instructions and visual observations from the DROID robot platform, generate robot action trajectories for manipulation and control tasks.
43
-
44
- - Cosmos3-Super-Image2Video:
45
- - Given one input image and text instructions, generate temporally coherent video sequences that are consistent with the provided visual content.
46
-
47
- - Cosmos3-Super-Text2Image:
48
- - Given text input, generate high-fidelity images that are consistent with the provided description.
49
-
50
- ### License
51
-
52
- This model is released under the [OpenMDW1.1](https://openmdw.ai/license/1-1/)
53
-
54
- ### Deployment Geography
55
-
56
- Global
57
-
58
- ### Use Case
59
-
60
- Physical AI: Encompassing robotics, autonomous vehicles (AV), and smart space environments, including industrial and factory-scale applications.
61
-
62
- ### Release Date
63
-
64
- Hugging Face 05/31/2026 via [https://huggingface.co/collections/nvidia/cosmos3](https://huggingface.co/collections/nvidia/cosmos3)
65
- GitHub 05/31/2026 via [https://github.com/nvidia/cosmos](https://github.com/nvidia/cosmos)
66
-
67
- ## Model Architecture
68
-
69
- **Architecture Type:** Transformer
70
-
71
- **Network Architecture:** Mixture-of-Transformers (MoT)
72
-
73
- Cosmos3 is an Omni-modal foundation model built on a Mixture-of-Transformers (MoT) architecture consisting of two complementary transformer towers: an autoregressive transformer for discrete token generation and a diffusion transformer for continuous multimodal generation. During inference, text is generated through standard next-token autoregressive decoding, while non-text modalities, such as images, video, audio, and actions, are synthesized through iterative denoising. This unified architecture enables Cosmos3 to model heterogeneous modalities within a single framework while preserving generation mechanisms best suited to each modality.
74
-
75
- **This model was developed based on:** [Cosmos Framework](https://github.com/nvidia/cosmos-framework)
76
-
77
- **Number of trainable model parameters:**
78
-
79
- - Cosmos3-Nano: 16B
80
- - Cosmos3-Super: 64B
81
- - Cosmos3-Nano-Policy-DROID: 16B
82
- - Cosmos3-Super-Image2Video: 64B
83
- - Cosmos3-Super-Text2Image: 64B
84
-
85
- ## Input/Output Specifications
86
-
87
- - **Generator Input**
88
- - **Input Type(s)**: Text, Image, Video (with audio or without audio), Action Trajectory
89
- - **Input Format(s)**:
90
- - Text: String
91
- - Image: jpg, png, jpeg, webp
92
- - Video (with or without audio): mp4
93
- - Action: json (1D list)
94
- - **Input Parameters**:
95
- - Text: One-dimensional (1D)
96
- - Image: Two-dimensional (2D)
97
- - Video: Three-dimensional (3D)
98
- - Audio: One-dimensional (1D)
99
- - Action trajectory: One-dimensional (1D)
100
- - **Other Properties Related to Input**:
101
- - For video inputs, we accept various resolutions, including 720p, 480p, and 256p.
102
- - When using input video with audio muxed into the video MP4 file, the audio should have 2 channels (stereo) and a 48 kHz sample rate.
103
- - Image and video inputs are RGB color (8 bits per channel, sRGB color space); grayscale inputs are not supported.
104
- - Action input is a per-frame sequence of robot/agent state or control values (e.g., joint positions, gripper state, camera pose). The full input is a 2D array shaped (T, D), where T is the number of frames and D is the embodiment-specific dimensionality listed below.
105
- - Input action is only supported for compatible embodiments, including general camera motion (9D), autonomous vehicle (9D), egocentric motion (57D), single Franka Panda arm with RobotiQ gripper (10D), dual Franka Panda arm with RobotiQ gripper (20D), Agibot (29D), UR (10D), Google robot (10D), WidowX 250 (10D), UMI (9D).
106
- - **Input Size and Length limits:**
107
- - **Text:** 4096 tokens
108
- - **Image:** 256p, 480p, and 720p resolution at one of these aspect ratios (16:9, 4:3, 1:1, 3:4, 9:16)
109
- - **Video:** 256p, 480p, and 720p resolution at one of these aspect ratios (16:9, 4:3, 1:1, 3:4, 9:16). Max number of frames = 5.
110
- - **Audio:** Max 0.5 second
111
- - **Action:** 16 – 400 video frames
112
- - **Generator Output**
113
- - **Output Type(s)**: Image, video, audio, action, text
114
- - **Output Format(s)**:
115
- - Image: JPG
116
- - Video: MP4
117
- - Audio: Advanced Audio Coding (AAC) stream (muxed within the MP4)
118
- - Action: 1D list (.json)
119
- - Text: string
120
- - **Output Parameters**:
121
- - Image: Two-dimensional (2D)
122
- - Video: Three-dimensional (3D)
123
- - Audio: One-dimensional (1D)
124
- - Action: One-dimensional (1D)
125
- - Text: One-dimensional (1D)
126
- - **Other Properties Related to Output**:
127
- - The generated video is an MP4 file, with the resolution, frame rate, and duration specified in the input. The generated audio is encoded in AAC format, muxed into the video MP4 file with 2 channels (stereo) and a 48 kHz sample rate.
128
- - Video generation supports durations from 5 to 400 frames, with 189 frames as the default generation duration.
129
- - The generated action is only supported for compatible embodiments, including general camera motion (9D), autonomous vehicle (9D), egocentric motion (57D), single Franka Panda arm with RobotiQ gripper (10D), dual Franka Panda arm with RobotiQ gripper (20D), Agibot (29D), UR (10D), Google robot (10D), WidowX 250 (10D), UMI (9D).
130
- - Audio: 48 kHz stereo AAC stream muxed into video mp4
131
- - Video: mp4 at the FPS specified in input
132
- - Image: JPEG
133
- - **Reasoner Input**
134
- - **Input Type(s)**: Text, Text+Image, Text+Video
135
- - **Input Format(s)**:
136
- - Text: String
137
- - Image: jpg, png, jpeg, webp
138
- - Video: mp4
139
- - **Input Parameters**:
140
- - Text: One-dimensional (1D)
141
- - Image: Two-dimensional (2D)
142
- - Video: Three-dimensional (3D)
143
- - **Other Properties Related to Input**:
144
- - Video inputs are recommended at a frame rate of 4 fps.
145
- - Long-context inputs supported up to 256K tokens.
146
- - **Input Size and Length limits:**
147
- - **Text:** Up to 256K tokens (context window).
148
- - **Image:** Standard input image formats; passed as file or URL.
149
- - **Video:** mp4 at the recommended 4 fps.
150
- - **Reasoner Output**
151
- - **Output Type(s)**: Text
152
- - **Output Format(s)**:
153
- - Text: string
154
- - **Output Parameters**:
155
- - Text: One-dimensional (1D)
156
- - **Other Properties Related to Output**:
157
- - Default `max_tokens=4096+` is recommended for reasoning outputs; longer outputs may be requested.
158
- - Reasoning outputs may include structured chain-of-thought, 2D/3D point localization, and bounding-box coordinates for vision-based tasks.
159
-
160
- The video content visualizes the input text description as a short animated scene, capturing key elements within the specified time constraints.
161
-
162
- Our AI models are designed and/or optimized to run on NVIDIA GPU-accelerated systems. By leveraging NVIDIA's hardware (e.g., GPU cores) and software frameworks (e.g., CUDA libraries), the model achieves faster training and inference times compared to CPU-only solutions.
163
-
164
- ## Software Integration
165
-
166
- **Runtime Engine(s):**
167
-
168
- - [PyTorch](https://github.com/nvidia/cosmos3)
169
- - [vLLM-Omni](https://github.com/vllm-project/vllm-omni)
170
- - [Hugging Face Diffusers](https://huggingface.co/docs/diffusers/en/index)
171
- - [SGLang](https://github.com/sgl-project/sglang)
172
-
173
- **Supported Hardware Microarchitecture Compatibility:**
174
-
175
- - NVIDIA Ampere
176
- - NVIDIA Blackwell
177
- - NVIDIA Hopper
178
-
179
- **Operating System(s):**
180
-
181
- - Linux (We have not tested on other operating systems.)
182
-
183
- **Note:** Only BF16 precision is tested. Other precisions like FP4, FP8, and FP16 are not officially supported.
184
-
185
- The integration of foundation and fine-tuned models into AI systems requires additional testing using use-case-specific data to ensure safe and effective deployment. Following the V-model methodology, iterative testing and validation at both unit and system levels are essential to mitigate risks, meet technical and functional requirements, and ensure compliance with safety and ethical standards before deployment.
186
-
187
- ## Training, Testing, and Evaluation Datasets
188
-
189
- ### Dataset Overview
190
 
191
- - **Total Size:** 1.3B data points
192
- - **Total Number of Datasets:** 393 dataset entries
193
- - **Dataset partition:** Training [100%], Testing [N/A — evaluation benchmarks used separately], Validation [N/A — evaluation benchmarks used separately]
194
- - **Time period for training data collection:** 2024–2026
195
- - **Time period for testing data collection:** N/A (standard public benchmarks)
196
- - **Time period for validation data collection:** N/A (standard public benchmarks)
197
-
198
- Raw data from internal and external sources is transformed into training-ready data through multiple stages of curation, filtering, and quality review. Data acquisition spans diverse multimodal sources — robotics, autonomous driving, industrial environments, indoor and outdoor scenes, varied lighting and weather conditions, camera viewpoints, object categories, and human activities — to broaden coverage across Physical AI operating environments. Automated filtering pipelines remove corrupted, duplicate, low-quality, and restricted content. Metadata analysis, heuristic rules, and model-assisted classifiers are applied during preprocessing to flag anomalous distributions and low-diversity subsets. Human review supplements automated filtering for selected datasets, benchmark construction, and targeted quality analysis. Datasets are balanced across modalities and task categories — visual reasoning, text-to-image, text-to-video, image-to-video, audio generation, video transfer, action-conditioned generation, and action command generation — to reduce overrepresentation of narrow domains. Synthetic and simulation-based augmentation supplements coverage of rare physical interactions and edge-case scenarios. Deduplication and provenance tracking are applied across the corpus. The resulting processed data is converted into model-ready tokenized or encoded representations through modality-specific preprocessors before training begins.
199
-
200
- Training datasets passed through multiple layers of automated and manual safeguards designed to reduce the presence of harmful or policy-violating content across categories including weapons and weapons-related instructional content, criminal planning, child sexual abuse material (CSAM), non-consensual intimate imagery (NCII), sexual content involving minors, harassment, hate speech, profanity, threats and incitement to violence, self-harm or suicide-related content, and graphic violence. Data sources are reviewed for licensing compatibility, provenance, and alignment with internal data governance and safety policies before admission into training corpora. Automated filtering pipelines combine multiple detection strategies: hash-matching against known CSAM and NCII reference databases; classifier-based moderation models trained for explicit sexual content, hate speech, violence, weapons imagery, and other restricted categories; keyword and regex-based screening for criminal-planning, threats, and self-harm phrases in text data; metadata and provenance heuristics for source-level risk signals; and embedding-based anomaly detection to surface samples that fall outside expected distributions. Human review and targeted audits supplement automated filtering for selected datasets, benchmark construction, and safety-sensitive evaluation. For multimodal Physical AI data (robotics, autonomous driving, industrial scenes), additional filtering targets invalid action trajectories, physically implausible interactions, and unsafe control sequences. Synthetic and simulation-generated data are evaluated through internal validation before inclusion. Benchmark evaluations and red-team testing are applied post-training to surface remaining safety gaps across world generation, reasoning, audio, and action tasks. No large-scale data-filtering process can guarantee complete removal of all harmful content; residual risks may remain, particularly in rare edge cases or open-world deployment settings. Ongoing monitoring and dataset review continue post-release.
201
-
202
- **Data Modality and Training Data Size**
203
-
204
- | Modality | Reasoning Data Sample Count | Generation Data Sample Count |
205
- | -------- | ------------------- | -------------------- |
206
- | Text | 22M | Not Applicable |
207
- | Image | 19M | 767M |
208
- | Video | 1M | 348M |
209
- | Audio | Not Applicable | 139M |
210
- | Action | Not Applicable | 8M |
211
-
212
- **Data Collection Method by dataset**
213
-
214
- - Hybrid: Automatic/Sensors, Synthetic, Automated
215
-
216
- **Labeling Method by dataset**
217
-
218
- - Hybrid: Human, Automated
219
-
220
- **Properties:** The training, testing, and evaluation datasets consist of diverse multimodal video, image, audio, action, synthetic, and sensor-conditioned data sourced from NVIDIA-owned data and publicly available, commercially permissive datasets. These datasets are curated to exclude known restricted content and to support building an Omni model that learns to generate and reason about dynamic physical environments across world reasoning and generation tasks.
221
-
222
- ### Public Datasets
223
-
224
- | Dataset                                                             | Samples           |
225
- |---|---|
226
- | OpenImage | 1.2M |
227
- | Coyo700M | 100M |
228
- | YouTube Video | 340M |
229
- | UMI | 4.5M |
230
-
231
- ### Private Datasets
232
-
233
- | Dataset                                                             | Samples           |
234
- |---|---|
235
- | Egocentric | 7M |
236
- | Nexar | 0.6M |
237
- | AgiBot | 0.2M |
238
- | HOI | 0.3M |
239
-
240
- ### Synthetic Datasets
241
-
242
- | Dataset | Samples |
243
- |---|---|
244
- | synthetic images generated using HiDream-I1 | 15M |
245
- | synthetic images generated using Qwen-Image-2512 | 14M |
246
- | synthetic captions generated using Qwen3-VL | 1115M |
247
-
248
- ## Evaluation Datasets
249
-
250
- **Data Collection Method by dataset**
251
-
252
- - Hybrid: Automatic/Sensors, Synthetic, Automated
253
-
254
- **Labeling Method by dataset**
255
-
256
- - Hybrid: Human, Automated
257
-
258
- **Properties:** The training, testing, and evaluation datasets consist of diverse multimodal video, image, audio, action, synthetic, and sensor-conditioned data sourced from NVIDIA-owned data and publicly available, commercially permissive datasets. These datasets are curated to exclude known restricted content and to support building an Omni model that learns to generate and reason about dynamic physical environments across world reasoning and generation tasks.
259
-
260
- ## Benchmarks
261
-
262
- Please see our [technical paper](https://research.nvidia.com/labs/cosmos-lab/cosmos3/technical-report.pdf) for detailed evaluations of the base model.
263
-
264
- ### Text-to-image benchmark results
265
-
266
- ![benchmark results](assets/benchmark-text2image.png)
267
-
268
- ### Artificial Analysis Leaderboard
269
-
270
- #### Open-Source Models [2026/05/28/]
271
-
272
- ![Artificial Analysis Text-to-Image leaderboard — open-source models](assets/benchmark-text2image-leaderboard.png)
273
-
274
- #### All Models [2026/05/28/] (Including Closed-Source)
275
-
276
- ![Artificial Analysis Text-to-Image leaderboard — all models including closed-source](assets/benchmark-text2image-leaderboard-all-models.jpg)
277
-
278
- ## Qualitative examples
279
-
280
- ![Qualitative examples](assets/more_images.jpg)
281
 
282
  ## Usage
283
 
284
- - See [Cosmos](https://github.com/nvidia/cosmos) for details.
285
-
286
- ### Prompt upsampling
287
-
288
- For optimal quality, text prompts should be upsampled into a specific JSON structure. Description and code can be found [here](https://github.com/nvidia/cosmos-framework/blob/main/docs/prompt_upsampling.md).
289
-
290
- For example, for text-to-image upsampling using Opus-4.7:
291
-
292
- ```bash
293
- git clone https://github.com/NVIDIA/cosmos-framework.git packages/cosmos-framework
294
- pip install -e packages/cosmos-framework
295
-
296
- export PROMPT_UPSAMPLER_ENDPOINT_URL="https://api.anthropic.com/v1/"
297
- export PROMPT_UPSAMPLER_MODEL_NAME="claude-opus-4-7"
298
- export PROMPT_UPSAMPLER_API_TOKEN="<your_token>"
299
-
300
- python -m cosmos_framework.inference.prompt_upsampling \
301
- --input assets/original_prompt.txt \
302
- --output /tmp/upsampled_t2i_opus/ \
303
- --mode text2image \
304
- --endpoint-url "${PROMPT_UPSAMPLER_ENDPOINT_URL}" \
305
- --model "${PROMPT_UPSAMPLER_MODEL_NAME}" \
306
- --api-token "${PROMPT_UPSAMPLER_API_TOKEN}" \
307
- --resolution 768 \
308
- --aspect-ratio "1,1"
309
- ```
310
-
311
- The JSON-upsampled version of `assets/original_prompt.txt` is saved in `assets/example_caption.json` for convenience, and is used for the image generation examples below.
312
-
313
- ### vLLM-Omni
314
-
315
- #### Container
316
-
317
- ```
318
- docker pull vllm/vllm-omni:cosmos3
319
- ```
320
-
321
- #### General Invocation
322
-
323
- You can use the release-tested `vllm-omni` package for deploying an OpenAI-compatible API inference endpoint.
324
- The recommended vLLM-Omni serving configuration for `nvidia/Cosmos3-Super-Text2Image` on a 8xH100 node is:
325
-
326
- ```bash
327
- vllm serve nvidia/Cosmos3-Super-Text2Image \
328
- --omni \
329
- --host 0.0.0.0 \
330
- --port 8000 \
331
- --cfg-parallel-size 2 \
332
- --ulysses-degree 4 \
333
- --tensor-parallel-size 1 \
334
- --use-hsdp \
335
- --hsdp-shard-size 8 \
336
- --init-timeout 1800
337
- ```
338
-
339
- Setting `--enable-layerwise-offload` can help with memory usage on GPUs with less available memory; however, please note that for text2image generation, this may incur a significant performance penalty. For 4xH200 or 4xGB200 one can simply use `--cfg-parallel-size 2 --ulysses-degree 2 --tensor-parallel-size 1`.
340
-
341
- #### Examples
342
-
343
- ##### Text to image generation
344
- ```python
345
- import base64
346
- import json
347
- import requests
348
-
349
- # 1. Read JSON-upsampled prompt
350
- json_prompt = json.load(open("assets/example_caption.json"))
351
-
352
- # 2. Build your API payload
353
- payload = {
354
- "prompt": json.dumps(json_prompt),
355
- "size": "1024x1024",
356
- "n": 1, # single frame generation
357
- "num_inference_steps": 50,
358
- "guidance_scale": 4.0,
359
- "flow_shift": 3.0,
360
- "negative_prompt": "",
361
- "seed": 1143,
362
- "extra_args": {
363
- "use_resolution_template": False,
364
- "guardrails": True,
365
- },
366
- }
367
-
368
- # 3. Send the POST request
369
- url = "http://localhost:8000/v1/images/generations"
370
- print("Sending request to server...")
371
- response = requests.post(url, json=payload, headers={"Content-Type": "application/json"})
372
- response.raise_for_status()
373
-
374
- # 4. Extract the base64 data and decode it into an image
375
- response_json = response.json()
376
- b64_data = response_json["data"][0]["b64_json"]
377
- image_bytes = base64.b64decode(b64_data)
378
-
379
- # 5. Save the final PNG file
380
- with open("/tmp/cosmos3_t2i.png", "wb") as image_file:
381
- image_file.write(image_bytes)
382
- print("Saved image to /tmp/cosmos3_t2i.png")
383
- ```
384
-
385
- ![example_image](assets/example_image.png)
386
-
387
- ### SGLang
388
-
389
- [SGLang Diffusion](https://docs.sglang.io/docs/sglang-diffusion/index) can serve `nvidia/Cosmos3-Super-Text2Image` through OpenAI-compatible image generation endpoints. Install SGLang from the main branch with diffusion dependencies, then start the server:
390
-
391
- ```bash
392
- git clone --branch main https://github.com/sgl-project/sglang.git
393
- cd sglang
394
- pip install -e "python[diffusion]"
395
- pip install "cosmos-guardrail==0.3.1"
396
-
397
- sglang serve \
398
- --model-path nvidia/Cosmos3-Super-Text2Image \
399
- --num-gpus 4
400
- ```
401
-
402
- Cosmos 3 support in SGLang Diffusion currently requires the SGLang main branch. Switch to a stable SGLang release once Cosmos 3 support is included there.
403
-
404
- Example text-to-image request:
405
-
406
- ```bash
407
- curl -sS -X POST http://localhost:30000/v1/images/generations \
408
- -H "Content-Type: application/json" \
409
- -d '{
410
- "prompt": "A warehouse robot folds a blue cloth on a clean workbench.",
411
- "size": "1280x720",
412
- "n": 1,
413
- "num_inference_steps": 35,
414
- "guidance_scale": 6.0,
415
- "flow_shift": 10.0,
416
- "seed": 0,
417
- "extra_args": {
418
- "use_resolution_template": false,
419
- "guardrails": true
420
- }
421
- }'
422
- ```
423
-
424
- For complete serving instructions and request examples, see the [Cosmos3 SGLang cookbook](https://docs.sglang.io/cookbook/diffusion/Cosmos/Cosmos3).
425
-
426
  ### Diffusers
427
 
428
- Cosmos3 is fully supported within the popular HuggingFace Diffusers package. This integration makes it a supported inference backend, allowing developers to easily incorporate Cosmos3's capabilities - such as text-to-image generation - into their pipelines using the Cosmos3OmniPipeline class, as demonstrated by the provided code examples (see examples for other modalities on the HuggingFace Cosmos3 page).
429
-
430
- **Note:** This example is tested on GB200. For H100, use the [vLLM-Omni serving recipe](#vllm-omni) above, which supports multi-GPU deployment via HSDP.
431
-
432
- #### Installation
433
 
434
- To install diffusers with Cosmos3OmniPipeline:
435
- ```
436
- uv venv --python 3.13 --seed --managed-python
437
- source .venv/bin/activate
438
- uv pip install \
439
- "diffusers @ git+https://github.com/huggingface/diffusers.git" \
440
- accelerate \
441
- av \
442
- cosmos_guardrail \
443
- huggingface_hub \
444
- imageio \
445
- imageio-ffmpeg \
446
- torch \
447
- torchvision \
448
- transformers
449
- ```
450
-
451
- #### Examples
452
-
453
- ##### Text to image generation
454
  ```python
455
  import json
456
  import torch
457
  from diffusers import Cosmos3OmniPipeline
458
  from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
459
 
460
- json_prompt = json.load(open("assets/example_caption.json"))
 
 
 
 
461
 
462
  pipe = Cosmos3OmniPipeline.from_pretrained(
463
  "nvidia/Cosmos3-Super-Text2Image",
@@ -478,28 +67,37 @@ result = pipe(
478
  generator=torch.Generator(device="cuda").manual_seed(1143),
479
  )
480
 
481
- result.video[0].save("/tmp/cosmos3_t2i.png")
482
- print("Saved image to /tmp/cosmos3_t2i.png")
483
  ```
484
 
485
- ## Limitations
486
-
487
- Cosmos3 may produce imperfect outputs in challenging scenarios. Generation artifacts include temporal inconsistency, unstable camera or object motion, imprecise physical interactions, inaccurate audio-video synchronization, and action-state drift — especially in long-horizon or high-resolution outputs. Reasoning may also be incorrect: object states, causal relationships, spatial geometry, temporal ordering, agent intent, and future outcomes can be misinferred, and complex or long-context inputs may yield hallucinated entities, inconsistent interpretations, or implausible predictions. Because the model lacks an explicit physics simulator, 3D geometry, 4D space-time evolution, object permanence, contact dynamics, and physical laws are only approximated — producing artifacts such as disappearing or morphing objects, unrealistic collisions, and physically implausible motions. Quality further degrades in out-of-distribution environments, safety-critical edge cases, and domains underrepresented in training.
488
-
489
- Cosmos3 outputs should not be treated as physically accurate simulation, reliable ground-truth reasoning, or safety-certified decision making. Applications involving robotics control, autonomous systems, scientific simulation, or safety-critical planning require additional validation, external constraints, system-level safety analysis, and domain-specific guardrails before deployment.
490
 
491
- ## Inference
492
 
493
- **Acceleration Engine:** [PyTorch](https://pytorch.org/), [vLLM](https://github.com/vllm-project/vllm), [vLLM-Omni](https://github.com/vllm-project/vllm-omni), [Hugging Face Diffusers](https://github.com/huggingface/diffusers), [SGLang](https://github.com/sgl-project/sglang), [SGLang Diffusion](https://docs.sglang.io/docs/sglang-diffusion/index)
 
 
 
 
 
 
 
 
 
 
 
494
 
495
- **Test Hardware:** GB200 and H100
496
 
497
- ## Ethical Considerations
 
 
 
498
 
499
- NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. Developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse.
500
 
501
- Please make sure you have proper rights and permissions for all input image and video content; if image or video includes people, personal health information, or intellectual property, the image or video generated will not blur or maintain proportions of image subjects included.
502
 
503
- Users are responsible for model inputs and outputs. Users are responsible for ensuring safe integration of this model, including implementing guardrails as well as other safety mechanisms, prior to deployment.
504
 
505
- For more detailed information on ethical considerations for this model, please see the Model Card++ [Explainability](EXPLAINABILITY.md), [Bias](BIAS.md), [Safety & Security](SAFETY.md), and [Privacy](PRIVACY.md) subcards. Please report model quality, risk, security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
 
1
  ---
2
+ library_name: diffusers
3
  license: other
4
  license_name: openmdw1.1-license
5
+ license_link: https://openmdw.ai/license/1-1/
6
+ pipeline_tag: text-to-image
 
7
  tags:
8
+ - nvidia
9
+ - cosmos
10
+ - cosmos3
11
+ - vllm-omni
12
+ - image-generation
 
 
 
 
13
  ---
14
 
15
  # **Cosmos 3: Omnimodal World Models for Physical AI**
16
+ **[Model Collection](https://huggingface.co/collections/nvidia/cosmos3)** | **[Code](https://github.com/nvidia/cosmos)** | **[Paper](https://huggingface.co/papers/2606.02800)** | **[Website](https://research.nvidia.com/labs/cosmos-lab/cosmos3/)**
17
 
18
+ [NVIDIA Cosmos™](https://github.com/nvidia/cosmos) is a world foundation model platform designed to accelerate the development of Physical AI by enabling machines to understand, simulate, and interact with the physical world across robotics, autonomous driving, and smart space environments.
19
+
20
+ Cosmos 3 is a family of omnimodal world models designed to jointly process and generate language, image, video, audio, and action sequences within a unified mixture-of-transformers architecture.
21
 
22
  # Model Overview: Cosmos3-Super-Text2Image
23
 
24
  ## Description
25
 
26
+ Cosmos3-Super-Text2Image is a 64B parameter model designed to generate high-fidelity images from text input. It serves as a foundational building block for Physical AI applications spanning world understanding, simulation, and embodied policy learning.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27
 
28
+ - **Model Developer:** NVIDIA
29
+ - **Architecture:** Mixture-of-Transformers (MoT)
30
+ - **Trainable Parameters:** 64B
31
+ - **License:** [OpenMDW1.1](https://openmdw.ai/license/1-1/)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
32
 
33
  ## Usage
34
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
35
  ### Diffusers
36
 
37
+ Cosmos3 is fully supported within the Hugging Face [Diffusers](https://huggingface.co/docs/diffusers/en/index) library.
 
 
 
 
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  ```python
40
  import json
41
  import torch
42
  from diffusers import Cosmos3OmniPipeline
43
  from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
44
 
45
+ # Cosmos 3 works best with structured, upsampled prompts.
46
+ # See https://github.com/nvidia/cosmos for prompt upsampling details.
47
+ json_prompt = {
48
+ "comprehensive_t2i_caption": "A photorealistic studio photograph of two hands working a spinning cylinder of wet gray clay on a pottery wheel, fingers gently pinching and pulling the walls upward to form a narrow neck and rounded belly."
49
+ }
50
 
51
  pipe = Cosmos3OmniPipeline.from_pretrained(
52
  "nvidia/Cosmos3-Super-Text2Image",
 
67
  generator=torch.Generator(device="cuda").manual_seed(1143),
68
  )
69
 
70
+ result.video[0].save("cosmos3_t2i.png")
 
71
  ```
72
 
73
+ ### vLLM-Omni
 
 
 
 
74
 
75
+ You can use the `vllm-omni` package for deploying an OpenAI-compatible API inference endpoint.
76
 
77
+ ```bash
78
+ vllm serve nvidia/Cosmos3-Super-Text2Image \
79
+ --omni \
80
+ --host 0.0.0.0 \
81
+ --port 8000 \
82
+ --cfg-parallel-size 2 \
83
+ --ulysses-degree 4 \
84
+ --tensor-parallel-size 1 \
85
+ --use-hsdp \
86
+ --hsdp-shard-size 8 \
87
+ --init-timeout 1800
88
+ ```
89
 
90
+ ## Input/Output Specifications
91
 
92
+ - **Input Type**: Text
93
+ - **Output Type**: Image (JPEG)
94
+ - **Input Size**: Up to 4096 tokens
95
+ - **Supported Resolutions**: 256p, 480p, and 720p resolution at aspect ratios 16:9, 4:3, 1:1, 3:4, 9:16.
96
 
97
+ ## Limitations
98
 
99
+ Cosmos3 may produce imperfect outputs in challenging scenarios, such as imprecise physical interactions or object morphing. Because the model lacks an explicit physics simulator, physical laws are only approximated. Quality further degrades in out-of-distribution environments or safety-critical edge cases. Outputs should not be treated as physically accurate simulations or safety-certified decision-making tools.
100
 
101
+ ## Ethical Considerations
102
 
103
+ NVIDIA believes Trustworthy AI is a shared responsibility. Users are responsible for ensuring safe integration of this model, including implementing guardrails and other safety mechanisms prior to deployment. Please ensure you have proper rights and permissions for all input content.
agentic_upsampling/clients.py CHANGED
@@ -477,9 +477,8 @@ class VLMQualityJudge:
477
  api_token: str,
478
  endpoint_url: str = DEFAULT_CRITIC_ENDPOINT_URL,
479
  model: str = DEFAULT_CRITIC_MODEL,
480
- max_tokens: int = 16384,
481
  image_jpeg_quality: int | None = DEFAULT_JPEG_QUALITY,
482
- extra_body: dict[str, Any] | None = None,
483
  ) -> None:
484
  self.chat_client = OpenAIChatClient(
485
  ChatClientConfig(
@@ -488,7 +487,6 @@ class VLMQualityJudge:
488
  api_token=api_token,
489
  max_tokens=max_tokens,
490
  max_retries=3,
491
- extra_body=extra_body,
492
  )
493
  )
494
  self.image_jpeg_quality = image_jpeg_quality
 
477
  api_token: str,
478
  endpoint_url: str = DEFAULT_CRITIC_ENDPOINT_URL,
479
  model: str = DEFAULT_CRITIC_MODEL,
480
+ max_tokens: int = 8192,
481
  image_jpeg_quality: int | None = DEFAULT_JPEG_QUALITY,
 
482
  ) -> None:
483
  self.chat_client = OpenAIChatClient(
484
  ChatClientConfig(
 
487
  api_token=api_token,
488
  max_tokens=max_tokens,
489
  max_retries=3,
 
490
  )
491
  )
492
  self.image_jpeg_quality = image_jpeg_quality
agentic_upsampling/runner.py CHANGED
@@ -3,7 +3,6 @@
3
  from __future__ import annotations
4
 
5
  import json
6
- import random
7
  import traceback
8
  from concurrent.futures import ThreadPoolExecutor, as_completed
9
  from dataclasses import dataclass
@@ -113,12 +112,6 @@ class AgenticUpsamplerRunner:
113
  (item_dir / "failure.json").unlink(missing_ok=True)
114
  (item_dir / "incomplete.json").unlink(missing_ok=True)
115
  self._log(f"[prompt {item.prompt_id}] start")
116
- # One base per item, reused by every iteration (sample slot j gets base + j), so a
117
- # rewrite is compared against the same noise. A resumed run draws a new base.
118
- seed_base = self.config.seed_base
119
- if seed_base is None:
120
- seed_base = random.randint(0, 2**31 - 1 - self.config.samples_per_iteration)
121
- self._log(f"[prompt {item.prompt_id}] seed_base drawn: {seed_base}")
122
  candidates: list[dict[str, Any]] = []
123
  previous_prompt: dict[str, Any] | None = None
124
  previous_analysis: dict[str, Any] | None = None
@@ -138,7 +131,6 @@ class AgenticUpsamplerRunner:
138
  previous_analysis,
139
  previous_negative_prompt,
140
  candidates,
141
- seed_base,
142
  )
143
  except Exception as exc:
144
  if not candidates:
@@ -160,9 +152,7 @@ class AgenticUpsamplerRunner:
160
  self._log(f"[prompt {item.prompt_id}] early stop at iter={iteration}")
161
  break
162
 
163
- result = self.finalize_item(item, candidates, incomplete_error=incomplete_error)
164
- result["seed_base"] = seed_base
165
- return result
166
 
167
  def run_item_safely(self, item: PromptItem) -> dict[str, Any]:
168
  """Run one item and convert failures into structured records."""
@@ -189,7 +179,6 @@ class AgenticUpsamplerRunner:
189
  previous_analysis: dict[str, Any] | None,
190
  previous_negative_prompt: str,
191
  candidates: list[dict[str, Any]],
192
- seed_base: int,
193
  ) -> dict[str, Any]:
194
  prepared = self.prepare_iteration_prompt(
195
  item,
@@ -206,7 +195,6 @@ class AgenticUpsamplerRunner:
206
  iteration,
207
  prepared.prompt_json,
208
  prepared.negative_prompt,
209
- seed_base,
210
  )
211
  return self.finalize_iteration(item, iteration_dir, iteration, sample_candidates, sample_errors)
212
 
@@ -217,7 +205,6 @@ class AgenticUpsamplerRunner:
217
  iteration: int,
218
  prompt_json: dict[str, Any],
219
  negative_prompt: str,
220
- seed_base: int,
221
  ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
222
  """Generate seed samples concurrently, then judge successful images in sample order."""
223
  generation_outputs: dict[int, GenerationOutput] = {}
@@ -231,7 +218,6 @@ class AgenticUpsamplerRunner:
231
  sample_index,
232
  prompt_json,
233
  negative_prompt,
234
- seed_base,
235
  ): sample_index
236
  for sample_index in range(self.config.samples_per_iteration)
237
  }
@@ -314,11 +300,8 @@ class AgenticUpsamplerRunner:
314
  sample_index: int,
315
  prompt_json: dict[str, Any],
316
  negative_prompt: str,
317
- seed_base: int,
318
  ) -> dict[str, Any]:
319
- generation = self.run_generation_sample(
320
- item, iteration_dir, sample_index, prompt_json, negative_prompt, seed_base
321
- )
322
  return self.judge_iteration_sample(
323
  item,
324
  iteration_dir,
@@ -336,7 +319,6 @@ class AgenticUpsamplerRunner:
336
  sample_index: int,
337
  prompt_json: dict[str, Any],
338
  negative_prompt: str,
339
- seed_base: int,
340
  ) -> GenerationOutput:
341
  """Generate one sample image for an iteration."""
342
  sample_dir = self._sample_dir(iteration_dir, sample_index)
@@ -346,7 +328,7 @@ class AgenticUpsamplerRunner:
346
  prompt_json=prompt_json,
347
  prompt_id=item.prompt_id,
348
  output_dir=sample_dir,
349
- seed=seed_base + sample_index,
350
  negative_prompt=negative_prompt,
351
  jpeg_quality=self.config.jpeg_quality,
352
  )
@@ -440,6 +422,11 @@ class AgenticUpsamplerRunner:
440
  if self.config.verbose:
441
  print(message, flush=True)
442
 
 
 
 
 
 
443
  def _sample_dir(self, iteration_dir: Path, sample_index: int) -> Path:
444
  if self.config.samples_per_iteration == 1:
445
  return iteration_dir
 
3
  from __future__ import annotations
4
 
5
  import json
 
6
  import traceback
7
  from concurrent.futures import ThreadPoolExecutor, as_completed
8
  from dataclasses import dataclass
 
112
  (item_dir / "failure.json").unlink(missing_ok=True)
113
  (item_dir / "incomplete.json").unlink(missing_ok=True)
114
  self._log(f"[prompt {item.prompt_id}] start")
 
 
 
 
 
 
115
  candidates: list[dict[str, Any]] = []
116
  previous_prompt: dict[str, Any] | None = None
117
  previous_analysis: dict[str, Any] | None = None
 
131
  previous_analysis,
132
  previous_negative_prompt,
133
  candidates,
 
134
  )
135
  except Exception as exc:
136
  if not candidates:
 
152
  self._log(f"[prompt {item.prompt_id}] early stop at iter={iteration}")
153
  break
154
 
155
+ return self.finalize_item(item, candidates, incomplete_error=incomplete_error)
 
 
156
 
157
  def run_item_safely(self, item: PromptItem) -> dict[str, Any]:
158
  """Run one item and convert failures into structured records."""
 
179
  previous_analysis: dict[str, Any] | None,
180
  previous_negative_prompt: str,
181
  candidates: list[dict[str, Any]],
 
182
  ) -> dict[str, Any]:
183
  prepared = self.prepare_iteration_prompt(
184
  item,
 
195
  iteration,
196
  prepared.prompt_json,
197
  prepared.negative_prompt,
 
198
  )
199
  return self.finalize_iteration(item, iteration_dir, iteration, sample_candidates, sample_errors)
200
 
 
205
  iteration: int,
206
  prompt_json: dict[str, Any],
207
  negative_prompt: str,
 
208
  ) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
209
  """Generate seed samples concurrently, then judge successful images in sample order."""
210
  generation_outputs: dict[int, GenerationOutput] = {}
 
218
  sample_index,
219
  prompt_json,
220
  negative_prompt,
 
221
  ): sample_index
222
  for sample_index in range(self.config.samples_per_iteration)
223
  }
 
300
  sample_index: int,
301
  prompt_json: dict[str, Any],
302
  negative_prompt: str,
 
303
  ) -> dict[str, Any]:
304
+ generation = self.run_generation_sample(item, iteration_dir, sample_index, prompt_json, negative_prompt)
 
 
305
  return self.judge_iteration_sample(
306
  item,
307
  iteration_dir,
 
319
  sample_index: int,
320
  prompt_json: dict[str, Any],
321
  negative_prompt: str,
 
322
  ) -> GenerationOutput:
323
  """Generate one sample image for an iteration."""
324
  sample_dir = self._sample_dir(iteration_dir, sample_index)
 
328
  prompt_json=prompt_json,
329
  prompt_id=item.prompt_id,
330
  output_dir=sample_dir,
331
+ seed=self._sample_seed(sample_index),
332
  negative_prompt=negative_prompt,
333
  jpeg_quality=self.config.jpeg_quality,
334
  )
 
422
  if self.config.verbose:
423
  print(message, flush=True)
424
 
425
+ def _sample_seed(self, sample_index: int) -> int | None:
426
+ if self.config.seed_base is None:
427
+ return None
428
+ return self.config.seed_base + sample_index
429
+
430
  def _sample_dir(self, iteration_dir: Path, sample_index: int) -> Path:
431
  if self.config.samples_per_iteration == 1:
432
  return iteration_dir
modular_model_index.json DELETED
@@ -1,75 +0,0 @@
1
- {
2
- "_blocks_class_name": "Cosmos3OmniBlocks",
3
- "_class_name": "Cosmos3OmniModularPipeline",
4
- "_diffusers_version": "0.39.0.dev0",
5
- "text_tokenizer": [
6
- "transformers",
7
- "Qwen2TokenizerFast",
8
- {
9
- "pretrained_model_name_or_path": "nvidia/Cosmos3-Super-Text2Image",
10
- "revision": null,
11
- "subfolder": "text_tokenizer",
12
- "type_hint": [
13
- "transformers",
14
- "Qwen2TokenizerFast"
15
- ],
16
- "variant": null
17
- }
18
- ],
19
- "vae": [
20
- "diffusers",
21
- "AutoencoderKLWan",
22
- {
23
- "pretrained_model_name_or_path": "nvidia/Cosmos3-Super-Text2Image",
24
- "revision": null,
25
- "subfolder": "vae",
26
- "type_hint": [
27
- "diffusers",
28
- "AutoencoderKLWan"
29
- ],
30
- "variant": null
31
- }
32
- ],
33
- "transformer": [
34
- "diffusers",
35
- "Cosmos3OmniTransformer",
36
- {
37
- "pretrained_model_name_or_path": "nvidia/Cosmos3-Super-Text2Image",
38
- "revision": null,
39
- "subfolder": "transformer",
40
- "type_hint": [
41
- "diffusers",
42
- "Cosmos3OmniTransformer"
43
- ],
44
- "variant": null
45
- }
46
- ],
47
- "scheduler": [
48
- "diffusers",
49
- "UniPCMultistepScheduler",
50
- {
51
- "pretrained_model_name_or_path": "nvidia/Cosmos3-Super-Text2Image",
52
- "revision": null,
53
- "subfolder": "scheduler",
54
- "type_hint": [
55
- "diffusers",
56
- "UniPCMultistepScheduler"
57
- ],
58
- "variant": null
59
- }
60
- ],
61
- "sound_tokenizer": [
62
- "diffusers",
63
- "Cosmos3AVAEAudioTokenizer",
64
- {
65
- "pretrained_model_name_or_path": "nvidia/Cosmos3-Super-Text2Image",
66
- "revision": null,
67
- "subfolder": "sound_tokenizer",
68
- "type_hint": [
69
- "diffusers",
70
- "Cosmos3AVAEAudioTokenizer"
71
- ],
72
- "variant": null
73
- }
74
- ]
75
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
tests/test_agentic_upsampling.py CHANGED
@@ -10,7 +10,7 @@ from typing import Any
10
 
11
  from PIL import Image
12
 
13
- from agentic_upsampling.clients import ImageGenerationClient, PromptRewriterClient, VLMQualityJudge
14
  from agentic_upsampling.constants import (
15
  DEFAULT_CRITIC_ENDPOINT_URL,
16
  DEFAULT_CRITIC_MODEL,
@@ -413,32 +413,7 @@ def test_runner_early_stops_by_default(tmp_path: Path) -> None:
413
  assert result["best_iteration"] == 0
414
  assert rewriter.initial_calls == 1
415
  assert rewriter.joint_rewrite_calls == 0
416
- assert generator.seeds == [result["seed_base"]]
417
-
418
-
419
- def test_runner_draws_one_seed_base_when_unset(tmp_path: Path) -> None:
420
- """An unset seed_base is drawn once and reused, so slot j gets base + j every iteration."""
421
- rewriter = FakeRewriter()
422
- generator = FakeGenerator()
423
- runner = AgenticUpsamplerRunner(
424
- rewriter=rewriter,
425
- generator=generator, # type: ignore[arg-type]
426
- judge=FakeJudge([5.0, 6.0, 5.5, 6.5]),
427
- config=RunnerConfig(
428
- output_dir=tmp_path,
429
- max_iterations=2,
430
- samples_per_iteration=2,
431
- early_stop=False,
432
- ),
433
- )
434
-
435
- result = runner.run_item(_item())
436
-
437
- seed_base = result["seed_base"]
438
- assert seed_base is not None
439
- # Samples are generated concurrently, so compare each iteration's seeds as a set.
440
- assert sorted(generator.seeds[:2]) == [seed_base, seed_base + 1]
441
- assert sorted(generator.seeds[2:]) == [seed_base, seed_base + 1]
442
 
443
 
444
  def test_runner_can_disable_early_stop_and_select_best_sample(tmp_path: Path) -> None:
@@ -459,9 +434,7 @@ def test_runner_can_disable_early_stop_and_select_best_sample(tmp_path: Path) ->
459
 
460
  result = runner.run_item(_item("8", "exactly 12 balloons with exact color counts"))
461
 
462
- # Samples are generated concurrently, so compare each iteration's seeds as a set.
463
- assert sorted(generator.seeds[:3]) == [1000, 1001, 1002]
464
- assert sorted(generator.seeds[3:]) == [1000, 1001, 1002]
465
  assert rewriter.previous_scores == [9.0]
466
  assert result["best_iteration"] == 1
467
  assert result["best"]["selected_sample_index"] == 1
@@ -521,14 +494,3 @@ def test_extract_best_images_copies_images_and_writes_manifests(tmp_path: Path)
521
  assert (tmp_path / "export" / "best_generations.jsonl").exists()
522
  assert (tmp_path / "export" / "best_generations.csv").exists()
523
 
524
-
525
- def test_vlm_quality_judge_sends_no_extra_body_by_default() -> None:
526
- judge = VLMQualityJudge(api_token="token")
527
-
528
- assert judge.chat_client.config.extra_body is None
529
-
530
-
531
- def test_vlm_quality_judge_forwards_extra_body() -> None:
532
- judge = VLMQualityJudge(api_token="token", extra_body={"reasoning_effort": "high"})
533
-
534
- assert judge.chat_client.config.extra_body == {"reasoning_effort": "high"}
 
10
 
11
  from PIL import Image
12
 
13
+ from agentic_upsampling.clients import ImageGenerationClient, PromptRewriterClient
14
  from agentic_upsampling.constants import (
15
  DEFAULT_CRITIC_ENDPOINT_URL,
16
  DEFAULT_CRITIC_MODEL,
 
413
  assert result["best_iteration"] == 0
414
  assert rewriter.initial_calls == 1
415
  assert rewriter.joint_rewrite_calls == 0
416
+ assert generator.seeds == [None]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
417
 
418
 
419
  def test_runner_can_disable_early_stop_and_select_best_sample(tmp_path: Path) -> None:
 
434
 
435
  result = runner.run_item(_item("8", "exactly 12 balloons with exact color counts"))
436
 
437
+ assert generator.seeds == [1000, 1001, 1002, 1000, 1001, 1002]
 
 
438
  assert rewriter.previous_scores == [9.0]
439
  assert result["best_iteration"] == 1
440
  assert result["best"]["selected_sample_index"] == 1
 
494
  assert (tmp_path / "export" / "best_generations.jsonl").exists()
495
  assert (tmp_path / "export" / "best_generations.csv").exists()
496