Diffusers documentation
Cosmos 3
Cosmos 3
NVIDIA Cosmos 3 là một mô hình nền tảng thế giới (WFM) thống nhất cho Physical AI — một mô hình omni duy nhất kết hợp tạo thế giới, suy luận vật lý và tạo hành động. Nó thay thế các mô hình Predict, Reason và Transfer riêng biệt từ các bản phát hành Cosmos trước đó: dù bạn đang xây dựng cho robot, xe tự hành hay không gian thông minh, Cosmos 3 cung cấp cho bạn một nền tảng duy nhất để mô phỏng và hiểu thế giới vật lý.
Những gì được phát hành cùng bản cập nhật này:
- Các mô hình trên Hugging Face Hub có thẻ mô hình và giấy phép
- Tích hợp Cosmos 3 Diffusers cho các pipeline tạo (trang này)
- Post-training scripts để tinh chỉnh Cosmos 3 trên dữ liệu của bạn
- Các bộ dữ liệu tổng hợp (SDG) mở cho AI vật lý
Những điểm mới trong Cosmos 3
Thay đổi lớn nhất so với các phiên bản Cosmos trước đó là Cosmos 3 là một mô hình omni, được xây dựng trên kiến trúc Mixture-of-Transformers (MoT). Trước đây, các nhà phát triển làm việc với các mô hình riêng biệt cho việc tạo thế giới (Predict), tạo có kiểm soát (Transfer), hiểu cảnh (Reason) và tạo chính sách hành động. Cosmos 3 thống nhất tất cả những điều này trong một mô hình duy nhất, có khả năng suy luận và tạo ra nội dung xuyên suốt các phương thức (modalities) trong một lần truyền tiến (forward pass) duy nhất.
Từ một mô hình, bạn có thể:
- Tạo ra các thế giới video có tính khả thi về mặt vật lý từ văn bản, hình ảnh hoặc đầu vào hành động (tạo video từ hình ảnh, tạo video từ văn bản, tạo video có điều kiện theo hành động).
- Suy luận về các thuộc tính vật lý như chuyển động, quan hệ nhân quả và quan hệ không gian.
- Dự đoán các chuỗi video và hành động trong tương lai từ trạng thái hiện tại.
- Chuyển đổi cảnh giữa các góc nhìn và điều kiện khác nhau với kiểm soát cấu trúc (sắp ra mắt).
Bên trong, một Cosmos3OmniTransformer duy nhất chạy mô hình ngôn ngữ kiểu Qwen song song với đường dẫn sinh bằng diffusion: các token văn bản đi qua luồng “hiểu” có tính nhân quả trong khi các latent video và âm thanh đi qua luồng “sinh” có chú ý hai chiều, được kết nối bằng 3D multimodal RoPE. Xem bài báo về Nền tảng Mô hình Nền tảng Thế giới Cosmos để tìm hiểu bối cảnh kiến trúc.
Các checkpoint khả dụng
Hai checkpoint được phát hành trên Hub — nvidia/Cosmos3-Nano (nhỏ hơn, nhanh hơn) và nvidia/Cosmos3-Super (lớn hơn, chất lượng cao hơn). Cùng một lớp pipeline hỗ trợ text-to-image, text-to-video, image-to-video và (với checkpoint có khả năng tạo âm thanh) text+image-to-video-with-sound — hãy chọn một repo và sử dụng tab riêng cho từng model trong mỗi workflow bên dưới.
Hãy chắc chắn xem qua hướng dẫn về Schedulers để tìm hiểu cách khám phá sự đánh đổi giữa tốc độ và chất lượng của scheduler, và xem phần tái sử dụng các thành phần giữa các pipeline để tìm hiểu cách nạp các thành phần giống nhau vào nhiều pipeline một cách hiệu quả.
Giảm nhiễu FP8 hỗn hợp W8A8/W8A16
Các checkpoint FP8 chính thức của ModelOpt cho Nano, Super và Edge được phát hành dưới dạng các repo Hub riêng: nvidia/Cosmos3-Nano-FP8, nvidia/Cosmos3-Super-FP8 và nvidia/Cosmos3-Edge-FP8. Hãy nạp main của repo FP8. nvidia/Cosmos3-Super-Image2Video không có repo FP8 riêng; hãy nạp nó với revision="fp8".
Tất cả các checkpoint này đều được lượng tử hóa theo cùng một cách. W8A8 sử dụng trọng số 8-bit và kích hoạt 8-bit (GEMM ModelOpt được khôi phục). W8A16 tái sử dụng chính những trọng số 8-bit đó nhưng bỏ qua lượng tử hóa kích hoạt: trọng số FP8 được khử lượng tử hóa và một phép toán tuyến tính tiêu chuẩn được thực hiện trên các kích hoạt BF16/FP16/FP32.
Việc chạy W8A8 ở mỗi bước có thể gây ra hiện tượng nhấp nháy rõ rệt trong quá trình tạo video nhiều bước. Do đó, các checkpoint video Nano, Super, Super-I2V và Edge FP8 khai báo một lịch trình trong transformer/config.json: W8A16 ở 3 bước đầu và 3 bước cuối, W8A8 ở phần giữa. Diffusers đọc các giá trị đếm này từ checkpoint thay vì hardcode. Độ chính xác được chọn một lần cho mỗi bước của scheduler để các cuộc gọi cond/uncond của classifier-free guidance khớp nhau.
Việc tạo ảnh và các checkpoint đã chưng cất theo kiểu ít bước không cho thấy hiện tượng nhấp nháy, do đó Super-T2I và các kho lưu trữ chưng cất 4 bước FP8 không khai báo lịch trình và duy trì W8A8 ở mọi bước. Lịch trình này cũng chỉ dành cho ModelOpt FP8: các backend lượng tử hóa khác (ví dụ TorchAO) vẫn giữ nguyên các forward gốc của chúng.
Tải một kho FP8 có cùng đường dẫn khôi phục như hướng dẫn ModelOpt. Kho đã chứa cấu hình lượng tử hóa. Độ chính xác hỗn hợp sau đó sẽ tự động theo checkpoint:
import torch
from diffusers import Cosmos3OmniPipeline
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano-FP8",
dtype=torch.bfloat16,
device_map="cuda",
)
result = pipe(prompt="...", num_inference_steps=35)Hai lựa chọn khi tạo:
- Mặc định (
mixed_precision_format=None): nếu checkpoint khai báodiffusion_step_policy, chạy W8A16 ở N bước đầu/cuối và W8A8 bản địa ở phần giữa. Đó là công thức dự kiến cho video FP8 nhiều bước (ít nhấp nháy hơn so với toàn bộ W8A8). FP8 4 bước đã chưng cất và Super-T2I FP8 bỏ qua chính sách, nên mặc định đã là toàn bộ W8A8. mixed_precision_format="none": giữ mọi bước ở chế độ W8A8 gốc. Nhanh hơn, vì W8A16 là dequant +torch.nn.functional.linearthay vì GEMM FP8 được khôi phục, nhưng video nhiều bước có thể bị nhấp nháy. Dùng chế độ này để A/B lịch trình hoặc để khớp với baseline lượng tử hóa hoàn toàn.
Trên một workstation Blackwell, Cosmos3-Nano-FP8 ở 720×1280 / 35 bước chậm hơn khoảng 27% (T2I) và 13% (T2V 49 khung hình) khi dùng lịch trình hỗn hợp mặc định so với "none". Những con số này không phải là cam kết về thông lượng. Chỉ truyền "fp8" để ép lịch trình first/last-N trên checkpoint ModelOpt FP8 không có policy.
Mẫu hóa lại prompt
Cosmos 3 được huấn luyện trên các chú thích dài, mô tả chi tiết. Để đạt chất lượng tối ưu, các prompt văn bản ngắn nên được mở rộng (upsample) thành một cấu trúc JSON cụ thể trước khi được truyền vào pipeline. Công cụ mở rộng (upsampler) nằm trong gói cosmos-framework.
Bắt đầu từ một prompt văn bản thuần túy ngắn gọn và lưu nó vào assets/prompt.txt. Đối với ví dụ text-to-video bên dưới, prompt gốc là “A robotic arm is cleaning a plate in a kitchen”:
mkdir -p assets
echo "A robotic arm is cleaning a plate in a kitchen" > assets/prompt.txtSau đó, cài đặt framework và chạy upsampler. Ví dụ bên dưới thực hiện upsampling cho text-to-video bằng Opus-4.6:
git clone https://github.com/NVIDIA/cosmos-framework.git packages/cosmos-framework
pip install -e packages/cosmos-framework
export PROMPT_UPSAMPLER_ENDPOINT_URL="https://api.anthropic.com/v1/"
export PROMPT_UPSAMPLER_MODEL_NAME="claude-opus-4-6"
export PROMPT_UPSAMPLER_API_TOKEN="<your_token>"
python -m cosmos_framework.inference.prompt_upsampling \
--input assets/prompt.txt \
--output assets/example_t2v_prompt.json \
--mode text2video \
--endpoint-url "${PROMPT_UPSAMPLER_ENDPOINT_URL}" \
--model "${PROMPT_UPSAMPLER_MODEL_NAME}" \
--api-token "${PROMPT_UPSAMPLER_API_TOKEN}" \
--resolution 720 \
--aspect-ratio "16,9"Chuyển --mode để khớp với quy trình làm việc bạn đang nhắm tới (text2image, text2video, image2video). Lệnh sẽ ghi các prompt đã được upsample vào tệp --output dưới dạng mảng JSON (một đối tượng cho mỗi dòng không rỗng trong --input); thay vào đó, hãy truyền một đường dẫn .jsonl để nhận một đối tượng JSON cho mỗi dòng. Đối với image2video, bạn cũng phải cung cấp ảnh điều kiện thông qua --image-url (một URL hoặc đường dẫn cục bộ) hoặc --image-list (một ảnh cho mỗi prompt).
Một prompt tích cực đã được pre-upsample (assets/example_t2v_prompt.json) và một prompt tiêu cực (assets/negative_prompt.json) được cung cấp để tiện lợi, và được sử dụng trong các ví dụ tạo bên dưới. Các ví dụ này nạp các tệp JSON đó và truyền chúng vào pipeline dưới dạng chuỗi JSON thông qua json.dumps(...).
Văn bản thành video
Tạo nhiều khung hình dựa trên văn bản. Chọn num_frames dựa trên thời lượng mục tiêu — mặc định num_frames=189 tạo ra ≈ 7,9 giây ở 24 FPS. Prompt và negative prompt được đọc từ các tệp đã được JSON-upsampled như mô tả trong Prompt upsampling.
import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
num_frames=189,
height=720,
width=1280,
num_inference_steps=35,
guidance_scale=6.0,
fps=24.0,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "cosmos3_t2v.mp4", fps=24, macro_block_size=1)import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
num_frames=189,
height=720,
width=1280,
num_inference_steps=35,
guidance_scale=6.0,
fps=24.0,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "cosmos3_t2v.mp4", fps=24, macro_block_size=1)Văn bản sang hình ảnh
Tạo khung đơn. Mô hình chỉ được điều kiện bởi prompt văn bản; truyền num_frames=1. Upsample với --mode text2image để tạo prompt JSON.
import json
import torch
from diffusers import Cosmos3OmniPipeline
# JSON-upsampled prompt (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2i_prompt.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
result = pipe(prompt=json.dumps(json_prompt), num_frames=1, height=720, width=1280)
result.video[0].save("cosmos3_t2i.jpg", format="JPEG", quality=85)import json
import torch
from diffusers import Cosmos3OmniPipeline
# JSON-upsampled prompt (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2i_prompt.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
result = pipe(prompt=json.dumps(json_prompt), num_frames=1, height=720, width=1280)
result.video[0].save("cosmos3_t2i.jpg", format="JPEG", quality=85)Hình ảnh sang video
Truyền ảnh điều kiện qua image=. Pipeline sẽ neo khung 0 vào ảnh được cung cấp và khử nhiễu các khung còn lại. Ảnh sẽ được thay đổi kích thước trong khi giữ nguyên tỷ lệ khung hình, cắt theo tâm để khớp với kích thước đầu ra yêu cầu, và được chuẩn hóa với làm tròn tương đương uint8. Dùng --mode image2video để lấy mẫu lên và tạo ra prompt JSON.
import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.utils import export_to_video, load_image
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_i2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
image = load_image(
"https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/assets/robot_153.jpg"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
image=image,
num_frames=189,
height=720,
width=1280,
fps=24.0,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "cosmos3_i2v.mp4", fps=24, macro_block_size=1)import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.utils import export_to_video, load_image
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_i2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
image = load_image(
"https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/assets/robot_153.jpg"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
image=image,
num_frames=189,
height=720,
width=1280,
fps=24.0,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "cosmos3_i2v.mp4", fps=24, macro_block_size=1)Video-to-video
Truyền một đoạn video điều kiện thông qua video= (ví dụ từ load_video). Pipeline sẽ neo các khung latent đầu tiên được chỉ định bởi condition_frame_indexes_vision (mặc định là [0, 1]) vào đoạn video và khử nhiễu phần còn lại. Sử dụng condition_video_keep ("first" hoặc "last") để chọn đầu nào của đoạn video nguồn dài hơn mà các khung điều kiện sẽ được lấy từ đó. Tương tự như các chế độ khác, prompt nên tuân theo cấu trúc JSON mô tả được nêu trong Prompt upsampling.
import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video, load_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_v2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
video=video,
condition_frame_indexes_vision=[0, 1],
condition_video_keep="first",
num_frames=189,
height=720,
width=1280,
num_inference_steps=35,
guidance_scale=6.0,
fps=24.0,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "cosmos3_v2v.mp4", fps=24, macro_block_size=1)import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video, load_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_v2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
video=video,
condition_frame_indexes_vision=[0, 1],
condition_video_keep="first",
num_frames=189,
height=720,
width=1280,
num_inference_steps=35,
guidance_scale=6.0,
fps=24.0,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "cosmos3_v2v.mp4", fps=24, macro_block_size=1)Video-to-video có âm thanh
Khi checkpoint chứa sound_tokenizer, hãy thêm enable_sound=True vào lệnh gọi video-to-video để đồng thời tạo ra một bản nhạc nền đồng bộ. Dạng sóng được trả về cùng với video và có thể được ghép vào tệp MP4 bằng encode_video().
import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import encode_video, load_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_v2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
video=video,
condition_frame_indexes_vision=[0, 1],
condition_video_keep="first",
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_sound=True,
)
encode_video(
result.video,
fps=24,
audio=result.sound,
audio_sample_rate=pipe.sound_tokenizer.config.sampling_rate,
output_path="cosmos3_v2v_with_sound.mp4",
)import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import encode_video, load_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_v2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
video=video,
condition_frame_indexes_vision=[0, 1],
condition_video_keep="first",
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_sound=True,
)
encode_video(
result.video,
fps=24,
audio=result.sound,
audio_sample_rate=pipe.sound_tokenizer.config.sampling_rate,
output_path="cosmos3_v2v_with_sound.mp4",
)Tạo video từ văn bản có âm thanh
Khi checkpoint chứa sound_tokenizer, hãy truyền enable_sound=True để đồng thời tạo ra một bản nhạc nền đồng bộ. Dạng sóng được trả về cùng với video và có thể được ghép vào tệp MP4 bằng encode_video().
Đây là cùng một lệnh gọi như ví dụ text-to-video ở trên, với enable_sound=True được thêm vào:
import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.utils import encode_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2v_sound_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_sound=True,
)
encode_video(
result.video,
fps=24,
audio=result.sound,
audio_sample_rate=pipe.sound_tokenizer.config.sampling_rate,
output_path="cosmos3_with_sound.mp4",
)import json
import torch
from diffusers import Cosmos3OmniPipeline
from diffusers.utils import encode_video
# JSON-upsampled positive and negative prompts (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2v_sound_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
result = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_sound=True,
)
encode_video(
result.video,
fps=24,
audio=result.sound,
audio_sample_rate=pipe.sound_tokenizer.config.sampling_rate,
output_path="cosmos3_with_sound.mp4",
)Tạo có điều kiện theo hành động
Action runs nhóm mọi đầu vào cụ thể cho hành động vào một CosmosActionCondition được truyền qua đối số action thay vì các đối số cấp cao image / video / height / width. Đặt resolution_tier (256/480/704/720) gần với độ phân giải gốc của video đầu vào; nó chọn khung nền điều kiện. Cosmos 3 hỗ trợ ba chế độ hành động — policy, forward_dynamics và inverse_dynamics. policy và forward_dynamics chỉ điều kiện trên khung hình đầu tiên (vì vậy cả image và video đều hoạt động), trong khi inverse_dynamics yêu cầu một video. Video điều kiện cho một lần chạy hành động được đặt trên action.video (hoặc action.image), không phải trên đối số video cấp cao của pipeline.
Truyền mô tả tác vụ đơn giản dưới dạng prompt và chọn camera bằng action.view_point (mặc định là "ego_view"; ngoài ra còn có "third_person_view", "wrist_view", "concat_view"). Pipeline sẽ chuyển đổi các thông tin này thành chú thích JSON có cấu trúc mà mô hình đã được huấn luyện, do đó prompt cho tác vụ không nên được LLM-upsample.
Chính sách hành động
Việc tạo chính sách hành động dự đoán các token video và hành động trong tương lai từ khung quan sát đầu tiên, prompt văn bản và metadata của miền hành động. Ví dụ bên dưới sử dụng miền robot Bridge và ghi chunk hành động được dự đoán vào JSON trong không gian hành động đã chuẩn hóa theo mô hình.
import json
import torch
from diffusers import Cosmos3OmniPipeline, CosmosActionCondition
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video, load_video
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
prompt = "Put the pot to the left of the purple item."
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_20260501_0.mp4"
)
result = pipe(
prompt=prompt,
action=CosmosActionCondition(
mode="policy",
chunk_size=16,
domain_name="bridge_orig_lerobot",
resolution_tier=480,
video=video,
view_point="ego_view",
),
fps=5,
num_inference_steps=30,
guidance_scale=1.0,
use_system_prompt=False,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "sample.mp4", fps=5, macro_block_size=1)
if result.action is not None:
with open("sample_action.json", "w") as f:
json.dump(result.action[0].tolist(), f)import json
import torch
from diffusers import Cosmos3OmniPipeline, CosmosActionCondition
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video, load_video
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Super", dtype=torch.bfloat16, device_map="cuda" # or "mps", "xpu", "cpu"
)
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
prompt = "Put the pot to the left of the purple item."
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_20260501_0.mp4"
)
result = pipe(
prompt=prompt,
action=CosmosActionCondition(
mode="policy",
chunk_size=16,
domain_name="bridge_orig_lerobot",
resolution_tier=480,
video=video,
view_point="ego_view",
),
fps=5,
num_inference_steps=30,
guidance_scale=1.0,
use_system_prompt=False,
)
# macro_block_size=1 allows arbitrary frame sizes (Cosmos3 outputs are not always divisible by 16).
export_to_video(result.video, "sample.mp4", fps=5, macro_block_size=1)
if result.action is not None:
with open("sample_action.json", "w") as f:
json.dump(result.action[0].tolist(), f)SeaCache
SeaCache bị vô hiệu hóa theo mặc định. Cosmos 3 hỗ trợ bật nó một cách rõ ràng bằng SeaCacheConfig. SeaCache tái sử dụng các phần dư của transformer khi chỉ báo Spectral-Evolution-Aware thay đổi chậm, từ đó giảm số lần thực thi đầy đủ của transformer. Nó tính toán chỉ báo này từ các latent thị giác thô, bao gồm các khung điều kiện sạch cho việc tạo video từ ảnh. Bật nó trên transformer bằng các callback metadata của scheduler từ pipeline:
import torch
from diffusers import Cosmos3OmniPipeline, SeaCacheConfig
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16, device_map="cuda"
)
pipe.transformer.enable_cache(SeaCacheConfig(threshold=0.2, max_consecutive_cached=2))Cùng một API ở cấp độ model hoạt động với Cosmos3OmniPipeline, Cosmos3OmniModularPipeline và Cosmos3DistilledModularPipeline. SeaCache là xấp xỉ và có thể làm thay đổi kết quả đầu ra được tạo ra. Hãy tắt nó bằng pipe.transformer.disable_cache() khi bạn cần mỗi bước khử nhiễu thực thi đầy đủ transformer. Trạng thái cache được đặt lại sau mỗi lần gọi pipeline, và các nhánh hướng dẫn có điều kiện và không điều kiện giữ các lịch sử độc lập.
Cosmos 3 giữ cổng SeaCache nằm bên ngoài các lớp decoder lặp lại, do đó tương thích với biên dịch theo vùng. Biên dịch các lớp sau khi bật cache:
pipe.transformer.compile_repeated_blocks(fullgraph=True)SeaCache cũng hỗ trợ các tiện ích song song ngữ cảnh Ulysses và song song tensor dựa trên DTensor của Cosmos 3 được mô tả bên dưới. Các quyết định về cache được đồng bộ hóa giữa các rank. Không cam kết về việc biên dịch toàn bộ mô hình và các chiến lược chia mô hình khác.
Song song ngữ cảnh
Đối với video dài hoặc độ phân giải cao, một lần truyền tiến (forward pass) có thể vượt quá ngân sách bộ nhớ và độ trễ của một GPU. Cosmos 3 hỗ trợ song song ngữ cảnh (context parallelism - CP) để chia nhỏ chiều chuỗi (sequence dimension) trên nhiều GPU, tách rời phép tính chú ý (attention) để mỗi thiết bị chỉ giữ một phần các token.
Cosmos 3 hỗ trợ song song ngữ cảnh Ulysses (trao đổi chuỗi/đầu all-to-all). Ring attention không được hỗ trợ.
Khác với hầu hết các mô hình diffusers, Cosmos 3 không gắn CP vào transformer hay đường dẫn khai báo enable_parallelism(): cơ chế grouped-query attention, các luồng understanding/generation riêng biệt (luồng generation chú ý đến cả hai) và độ dài không đều theo từng luồng không thể được biểu diễn dưới dạng _cp_plan. Thay vào đó, mô hình phơi bày các điểm nối shard/gather dạng no-op nhỏ gọn, và phần cài đặt nằm trong examples/cosmos3/cosmos_parallel.py — một module tự chứa mà bạn có thể đọc từ đầu đến cuối và tùy biến. Nó cung cấp hai trục sharding trực giao, có thể kết hợp được:
| Trợ lý | Các phân mảnh | Dùng cho |
|---|---|---|
enable_cosmos3_context_parallel(transformer, cp_mesh) | chuỗi (CP / Ulysses) | độ trễ trên một model vừa với một GPU (Nano) |
enable_cosmos3_tensor_parallel(transformer, tp_mesh) | weights (TP) | đào tạo một mô hình không vừa với một GPU (Super) |
Sử dụng một trong hai hoặc cả hai cùng nhau trên lưới 2-D (tp, cp) (xem Chạy mô hình lớn với tensor parallelism).
Hai yêu cầu sau đây là đặc thù của Cosmos 3:
- Sử dụng backend attention
native. Cosmos 3 sử dụng grouped-query attention (GQA), và backend SDPA native là backend duy nhất chấp nhậnenable_gqa(cuDNN và flash từ chối tham số này). Các hàm trợ giúp sẽ mở rộng các đầu KV lên số lượng đầu query và gọi SDPA vớienable_gqa=Falseđể nó vẫn được chuyển đến kernel flash (phương án dự phòng về mặt toán học sẽ tạo ra toàn bộ điểm số[S, S]và gây tràn bộ nhớ trên các chuỗi dài). - Bậc CP (Ulysses) phải chia hết cho số lượng đầu query (32 với
Nano, 64 vớiSuper); với TP, bậc phải chia hết cho số đầu KV (8). Các luồng hiểu (văn bản) và sinh (video/âm thanh) được chia độc lập dọc theo chuỗi, và các độ dài không đều được đệm bằng số 0 bên trong đến một bội số của kích thước world.
Chạy
CLI đầy đủ examples/cosmos3/inference_cosmos3.py sử dụng Cosmos3OmniModularPipeline và tái sử dụng các hàm trợ giúp này, do đó mọi chế độ (text-to-image/video, image-to-video, âm thanh, chế độ hành động) đều chạy đa GPU thông qua --tp-degree / --cp-degree. Khởi chạy bằng torchrun; --tp-degree * --cp-degree phải bằng --nproc_per_node. Mọi rank đều tạo ra cùng một kết quả; rank 0 ghi kết quả đó.
# CP only — Nano (fits one GPU); CP degree must divide 32 query heads.
torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py --model nano --cp-degree 4 --prompt "..."
# TP only — Super; TP degree must divide 64 query heads and 8 KV heads.
torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py --model super --tp-degree 4 --prompt "..."
# TP + CP — Super, with sound (TP=2 x CP=2 across 4 GPUs).
torchrun --nproc_per_node=4 examples/cosmos3/inference_cosmos3.py \
--model super --tp-degree 2 --cp-degree 2 --enable-sound --prompt "..."Super có ~120 GB trọng số, không vừa trên một GPU 96 GB, nên cần TP; Nano vừa trên một GPU duy nhất, nên CP cho nó là tối ưu hóa độ trễ thuần túy. (Bỏ qua cả hai cờ để chạy trên một GPU.)
Đào tạo các mô hình lớn với song song tensor
CP chia nhỏ activations nhưng nhân bản toàn bộ trọng số trên mỗi rank, do đó không làm giảm dung lượng trọng số của mô hình — một mô hình không vừa trên một GPU vẫn sẽ không vừa chỉ với CP. Để chia nhỏ trọng số, enable_cosmos3_tensor_parallel(transformer, tp_mesh) áp dụng tensor parallelism kiểu Megatron trên một trục mesh thứ hai, trực giao:
- Các phép chiếu attention và MLP được chia cột/hàng trên nhóm TP (
to_q/to_k/to_v+add_q/k/vvàgate/upcủa các MLP là song song theo cột;to_out/to_add_outvàdowncủa các MLP là song song theo hàng kèm all-reduce). Mỗi rank cuối cùng sở hữuquery_heads / tpđầu query vàkv_heads / tpđầu KV. - TP kết hợp với CP trên lưới thiết bị 2 chiều
(tp, cp): TP tách đầu/cân nặng một cách bền vững, còn CP chia nhỏ chuỗi phía trên. Các ràng buộc làtpphải chia hết cho số đầu KV (8), vàtp * cpphải chia hết cho số đầu truy vấn (32 vớiNano, 64 vớiSuper). - Trọng số được nạp vào CPU và chia nhỏ lên các GPU theo từng lớp, do đó mô hình đầy đủ không bao giờ được tạo ra trên một thiết bị duy nhất.
TP thực hiện all-reduce trên mỗi khối attention và MLP, do đó tiêu tốn nhiều băng thông. Trên các máy chủ không có NVLink, đây là chi phí chiếm ưu thế; hãy ưu tiên mức độ TP nhỏ nhất sao cho trọng số vừa đủ và đưa các GPU còn lại vào CP.
Sử dụng trong pipeline mô-đun của riêng bạn
Các cờ CLI rất tiện lợi, nhưng bạn có thể gọi trực tiếp các hàm trợ giúp bằng Cosmos3OmniModularPipeline. Tải cấu hình pipeline và các thành phần lên CPU, áp dụng TP trước khi di chuyển pipeline sang GPU cục bộ của rank, chuyển sang backend native, sau đó bật CP. Không sử dụng device_map cho luồng này:
import os
import sys
import torch
import torch.distributed as dist
from diffusers import Cosmos3OmniModularPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from torch.distributed.device_mesh import init_device_mesh
# Make the helper module importable.
sys.path.insert(0, "examples/cosmos3")
from cosmos_parallel import (
enable_cosmos3_context_parallel,
enable_cosmos3_flash_attention,
enable_cosmos3_tensor_parallel,
)
# torchrun sets RANK / WORLD_SIZE / LOCAL_RANK. Pick tp_degree * cp_degree == world size.
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
dist.init_process_group("nccl")
mesh = init_device_mesh("cuda", (tp_degree, cp_degree), mesh_dim_names=("tp", "cp"))
# Load components on CPU first; a TP-sharded model may not fit one GPU.
pipe = Cosmos3OmniModularPipeline.from_pretrained(model_id)
pipe.load_components(dtype=torch.bfloat16)
pipe.enable_safety_checker()
if tp_degree > 1:
enable_cosmos3_tensor_parallel(pipe.transformer, mesh["tp"]) # shard weights -> GPUs
pipe.to(f"cuda:{local_rank}") # move the replicated remainder
pipe.transformer.set_attention_backend("native")
if cp_degree > 1:
enable_cosmos3_context_parallel(pipe.transformer, mesh["cp"]) # shard the sequence
elif tp_degree > 1:
enable_cosmos3_flash_attention(pipe.transformer) # GQA-safe dense attention
# Modular pipelines replace components through update_components().
scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
pipe.update_components(scheduler=scheduler)
# A single output name returns that value; a list returns a dictionary.
outputs = pipe(
prompt='{"scene":"A robot arm in a kitchen"}',
num_frames=189,
height=720,
width=1280,
output=["videos", "sound", "sampling_rate", "action"],
)
videos = outputs["videos"]
sound = outputs["sound"] # None unless sound generation was requested.
action = outputs["action"] # None unless an action workflow produced actions.Chỉ dành cho CP (không chia shard trọng số), hãy sử dụng lưới 1-D: init_device_mesh("cuda", (world_size,), mesh_dim_names=("cp",)) và chỉ cần enable_cosmos3_context_parallel.
enable_safety_checker() tải và bật trình kiểm tra mặc định; disable_safety_checker() tắt nó một cách tường minh. Hãy sử dụng các phương thức pipeline đó thay vì tham số khởi tạo enable_safety_checker= của pipeline tác vụ hoặc tham số gọi enable_safety_check=. Các pipeline dạng module cũng không trả về Cosmos3OmniPipelineOutput: hãy dùng output="videos" cho riêng các khung hình, hoặc dùng danh sách output và từ điển được trả về như minh họa ở trên thay vì result.video, result.sound hoặc result.action.
Trên một số cấu hình đa GPU, lệnh all-to-all NCCL đầu tiên có thể bị treo. Nếu một lần chạy CP bị dừng ở đầu bước khử nhiễu đầu tiên, hãy đặt
NCCL_P2P_DISABLE=1trong môi trường trước khi khởi chạytorchrun.
CP và TP kết hợp với tất cả các quy trình làm việc ở trên (text-to-video, image-to-video, text-to-video có âm thanh, và sinh có điều kiện hành động) cũng như với cả checkpoint Nano và Super — chỉ có các dòng xây dựng pipeline và thiết lập song song là thay đổi.
Các mẫu metadata
tokenize_prompt thêm các câu metadata ngắn vào bên trong thông điệp của người dùng để LLM nhìn thấy các điều kiện mà mô hình đã được huấn luyện. Prompt tích cực nhận các câu như “The video is 7.9 seconds long and is of 24 FPS.” và “This video is of 720x1280 resolution.”; prompt tiêu cực nhận phần ngược lại (”… is not …”).
Cả hai đều được bật theo mặc định. Tắt từng cặp thông qua __call__:
result = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
num_frames=189,
height=720,
width=1280,
fps=24.0,
add_duration_template=False, # skip the duration sentence on both prompts
add_resolution_template=False, # skip the resolution sentence on both prompts
)add_duration_template không có tác dụng khi num_frames == 1 (chế độ ảnh); trong trường hợp đó, chỉ có câu về độ phân giải được thêm vào.
Trình kiểm tra an toàn
Cosmos3 tích hợp cosmos_guardrail CosmosSafetyChecker và chạy nó mặc định. Guardrail văn bản từ chối các prompt không an toàn trước khi sinh nội dung (ValueError); guardrail video chạy trên các khung hình đã giải mã và hoặc làm mờ các khuôn mặt được phát hiện hoặc từ chối đầu ra. Đầu ra âm thanh không được bảo vệ bởi guardrail.
Cài đặt phụ thuộc tùy chọn để bật trình kiểm tra mặc định:
pip install cosmos_guardrailTrình kiểm tra là bắt buộc theo Thỏa thuận Giấy phép Mô hình Mở của NVIDIA. Hai cờ bên dưới tồn tại cho các quy trình kiểm thử và phát triển, nơi rào chắn sẽ là thừa (ví dụ: đầu vào đã được làm sạch, hoặc bạn đang cố ý chạy thử pipeline trên các đầu vào biên).
Tắt khi khởi tạo (không có trình kiểm tra nào được khởi tạo, do đó không có mô hình guardrail nào được tải xuống hoặc nạp vào bộ nhớ):
import torch
from diffusers import Cosmos3OmniPipeline
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano",
dtype=torch.bfloat16,
device_map="cuda", # or "mps", "xpu", "cpu"
enable_safety_checker=False,
)Tắt cho một lần gọi (trình kiểm tra vẫn được nạp — hữu ích cho việc bỏ qua một lần duy nhất trong khi giữ mặc định bật cho các lần gọi sau):
result = pipe(
prompt=prompt,
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_safety_check=False,
)Để cung cấp một trình kiểm tra tùy chỉnh (ví dụ: một lớp con không thực hiện gì cho các bài kiểm tra nhanh), hãy truyền nó dưới dạng safety_checker=:
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano",
dtype=torch.bfloat16,
device_map="cuda", # or "mps", "xpu", "cpu"
safety_checker=MyCustomSafetyChecker(),
)Cosmos3OmniPipeline
lớp diffusers.Cosmos3OmniPipeline
< source >( transformer: Cosmos3OmniTransformertext_tokenizer: AutoTokenizervae: AutoencoderKLWanscheduler: UniPCMultistepSchedulersound_tokenizer: diffusers.models.autoencoders.autoencoder_cosmos3_audio.Cosmos3AVAEAudioTokenizer | None = Nonesafety_checker: diffusers.pipelines.cosmos.pipeline_cosmos3_omni.CosmosSafetyChecker | None = Noneenable_safety_checker: bool = Truedefault_use_system_prompt: bool = Trueuse_native_flow_schedule: bool = False )
Giải mã âm thanh tiềm ẩn [C, T] thành dạng sóng [audio_ch, N].
Thêm/xóa chiều batch mà bộ giải mã của bộ phân từ âm thanh kỳ vọng.
prepare_latents
< source >( image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, NoneType] = Nonevideo: typing.Union[list[PIL.Image.Image], torch.Tensor, numpy.ndarray, NoneType] = Nonecondition_frame_indexes_vision: collections.abc.Iterable[int] = (0, 1)condition_video_keep: typing.Literal['first', 'last'] = 'first'num_frames: int | None = Noneheight: int | None = Nonewidth: int | None = Nonefps: float = 24.0latents: typing.Optional[torch.Tensor] = Nonesound_latents: typing.Optional[torch.Tensor] = Noneaction_latents: typing.Optional[torch.Tensor] = Nonegenerator: typing.Optional[torch.Generator] = Nonedevice: str = 'cuda'dtype: dtype = torch.bfloat16enable_sound: bool = Falseaction: CosmosActionCondition | None = None )
Xây dựng điều kiện + nhiễu ban đầu cho một mẫu duy nhất.
tokenize_prompt
< source >( prompt: strnegative_prompt: str | None = Nonenum_frames: int = 189height: int = 720width: int = 1280fps: float = 24.0use_system_prompt: bool | None = Noneadd_resolution_template: bool = Trueadd_duration_template: bool = Trueaction_mode: str | None = Noneaction_view_point: str | None = None )
Áp dụng các mẫu bổ trợ prompt và thực hiện tokenize các prompt cond/uncond thông qua mẫu chat được cấu hình.
Pipeline này không chạy một text encoder riêng biệt: transformer Cosmos3 kết hợp nhận các token ID thô cùng với các token thị giác (và tùy chọn là âm thanh).
Khi negative_prompt là None, một chuỗi rỗng sẽ được sử dụng; trang tài liệu của Cosmos3 liệt kê các negative prompt kiểm soát chất lượng được khuyến nghị để truyền vào một cách tường minh cho text2video / image2video. Các mẫu thời lượng và độ phân giải sẽ được nối vào prompt, và các mẫu ngược sẽ được nối vào negative prompt, khi được bật.
Khi action_mode được đặt, prompt sẽ được chuyển đổi thành caption JSON hành động có cấu trúc mà mô hình đã được huấn luyện (xem _build_action_json_prompt), sử dụng action_view_point cho trường framing; các mẫu metadata phẳng sẽ bị bỏ qua vì JSON đã chứa duration/fps/resolution/aspect_ratio.
- tất cả
- call
Cosmos3OmniModularPipeline
Cosmos 3 cũng có sẵn dưới dạng pipeline Modular Diffusers. Pipeline dựa trên tác vụ Cosmos3OmniPipeline vẫn khả dụng; pipeline modular tồn tại song song với nó và bao phủ các chế độ giống nhau (text2image, text2video, image2video, video2video, tạo có điều kiện hành động, và transfer (kiểm soát cấu trúc), kèm âm thanh tùy chọn khi checkpoint hỗ trợ).
import torch
from diffusers import Cosmos3OmniModularPipeline
pipe = Cosmos3OmniModularPipeline.from_pretrained(
"nvidia/Cosmos3-Nano", dtype=torch.bfloat16
)
pipe.load_components(dtype=torch.bfloat16)
pipe.enable_safety_checker()
videos = pipe(
prompt='{"scene":"A robot arm in a kitchen"}',
num_frames=1,
height=720,
width=1280,
output="videos",
)
# Modular pipelines expose declared outputs directly instead of using the task pipeline's
# `return_dict`/`Cosmos3OmniPipelineOutput` API.
image = videos[0]Bạn cũng có thể tải thông qua ModularPipeline và để cấu hình kho chọn lớp blocks:
import torch
from diffusers import ModularPipeline
pipe = ModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", dtype=torch.bfloat16)
pipe.load_components(dtype=torch.bfloat16)
pipe.enable_safety_checker()
videos = pipe(
prompt='{"scene":"A robot arm in a kitchen"}', num_frames=1, height=720, width=1280, output="videos"
)Để kiểm tra hoặc tùy chỉnh một quy trình làm việc Cosmos mô-đun cụ thể, hãy sử dụng available_workflows + get_workflow():
available = pipe.blocks.available_workflows
image2video_blocks = pipe.blocks.get_workflow("image2video")Các ví dụ mô-đun cho tất cả các quy trình làm việc hiện có
Pipeline dạng module hỗ trợ các chữ ký gọi giống như pipeline theo tác vụ. Các đoạn mã bên dưới phản ánh mọi ví dụ sinh được hiển thị ở trên (text2video, text2image, image2video, video2video, video2video_sound, text2video_sound, và action_policy). Transfer (kiểm soát cấu trúc) có các đầu vào riêng và được trình bày riêng trong Modular transfer bên dưới.
import json
import torch
from diffusers import Cosmos3OmniModularPipeline, CosmosActionCondition
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import encode_video, export_to_video, load_image, load_video
pipe = Cosmos3OmniModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", dtype=torch.bfloat16)
pipe.load_components(dtype=torch.bfloat16)
pipe.enable_safety_checker()
pipe.to("cuda") # or "mps", "xpu", "cpu"
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
# text2video
json_prompt = json.load(open("assets/example_t2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
videos = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
num_frames=189,
height=720,
width=1280,
num_inference_steps=35,
guidance_scale=6.0,
fps=24.0,
output="videos",
)
export_to_video(videos, "cosmos3_modular_t2v.mp4", fps=24, macro_block_size=1)
# text2image
json_prompt = json.load(open("assets/example_t2i_prompt.json"))
videos = pipe(prompt=json.dumps(json_prompt), num_frames=1, height=720, width=1280, output="videos")
videos[0].save("cosmos3_modular_t2i.jpg", format="JPEG", quality=85)
# image2video
json_prompt = json.load(open("assets/example_i2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
image = load_image("https://github.com/nvidia-cosmos/cosmos-dependencies/releases/download/assets/robot_153.jpg")
videos = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
image=image,
num_frames=189,
height=720,
width=1280,
fps=24.0,
output="videos",
)
export_to_video(videos, "cosmos3_modular_i2v.mp4", fps=24, macro_block_size=1)
# video2video
json_prompt = json.load(open("assets/example_v2v_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt_i2v.json"))
video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4"
)
videos = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
video=video,
condition_frame_indexes_vision=[0, 1],
condition_video_keep="first",
num_frames=189,
height=720,
width=1280,
num_inference_steps=35,
guidance_scale=6.0,
fps=24.0,
output="videos",
)
export_to_video(videos, "cosmos3_modular_v2v.mp4", fps=24, macro_block_size=1)
# video2video_sound
outputs = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
video=video,
condition_frame_indexes_vision=[0, 1],
condition_video_keep="first",
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_sound=True,
output=["videos", "sound", "sampling_rate"],
)
encode_video(
outputs["videos"],
fps=24,
audio=outputs["sound"],
audio_sample_rate=outputs["sampling_rate"],
output_path="cosmos3_modular_v2v_with_sound.mp4",
)
# text2video_sound
json_prompt = json.load(open("assets/example_t2v_sound_prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
outputs = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_sound=True,
output=["videos", "sound", "sampling_rate"],
)
encode_video(
outputs["videos"],
fps=24,
audio=outputs["sound"],
audio_sample_rate=outputs["sampling_rate"],
output_path="cosmos3_modular_t2v_with_sound.mp4",
)
# action_policy
prompt = "Put the pot to the left of the purple item."
action_video = load_video(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/action/bridge_20260501_0.mp4"
)
outputs = pipe(
prompt=prompt,
action=CosmosActionCondition(
mode="policy",
chunk_size=16,
domain_name="bridge_orig_lerobot",
resolution_tier=480,
video=action_video,
view_point="ego_view",
),
fps=5,
num_inference_steps=30,
guidance_scale=1.0,
use_system_prompt=False,
output=["videos", "action"],
)
export_to_video(outputs["videos"], "cosmos3_modular_action_policy.mp4", fps=5, macro_block_size=1)
if outputs["action"] is not None:
with open("cosmos3_modular_action_policy.json", "w") as f:
json.dump(outputs["action"][0].tolist(), f)Chuyển giao mô-đun (kiểm soát cấu trúc)
Transfer làm theo một video điều khiển đã tính trước (cạnh, làm mờ, độ sâu, phân đoạn hoặc bản đồ kịch bản thế giới) được truyền qua control_videos= dưới dạng ánh xạ {hint: video}. Nó chỉ dành cho video (không có image / video / action / enable_sound), prompt là chú thích JSON đã được upsampling trước (xem Prompt upsampling), và các clip dài được tạo tự hồi quy theo từng khối num_video_frames_per_chunk rồi được ghép tự động. guidance_scale là CFG văn bản thông thường; control_guidance (!= 1.0) ngoài ra còn khuếch đại tín hiệu điều khiển. Các giá trị khởi đầu được khuyến nghị cho từng hint:
| Gợi ý | guidance_scale | control_guidance | flow_shift | Hình học |
|---|---|---|---|---|
| Viền / Làm mờ / Độ sâu | 3.0 | 1.5 | 10.0 | 121 khung hình @ 30 FPS |
| Phân đoạn | 3.0 | 2.0 | 10.0 | 121 khung hình @ 30 FPS |
| Kịch bản thế giới (WSM) | 1.0 | 3.0 | 10.0 | 101 khung hình @ 10 FPS |
Diffusers không cung cấp các tài nguyên điều khiển. Các tài nguyên có sẵn (một video điều khiển + prompt.json tương ứng cho mỗi gợi ý, cùng với negative_prompt.json dùng chung) nằm trong Cosmos cookbook. Đối với ví dụ về cạnh bên dưới, hãy tải chúng vào thư mục assets/ cục bộ:
base=https://github.com/NVIDIA/cosmos/raw/refs/heads/main/cookbooks/cosmos3/generator/transfer/assets
mkdir -p assets/edge
curl -sL "$base/edge/control_edge.mp4" -o assets/edge/control_edge.mp4
curl -sL "$base/edge/prompt.json" -o assets/edge/prompt.json
curl -sL "$base/negative_prompt.json" -o assets/negative_prompt.jsonimport json
import torch
from diffusers import Cosmos3OmniModularPipeline
from diffusers.schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
from diffusers.utils import export_to_video, load_video
pipe = Cosmos3OmniModularPipeline.from_pretrained("nvidia/Cosmos3-Nano", dtype=torch.bfloat16)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config, flow_shift=10.0, use_karras_sigmas=False
)
# Downloaded into assets/ from the Cosmos cookbook (see the curl snippet above).
json_prompt = json.load(open("assets/edge/prompt.json"))
negative_prompt = json.load(open("assets/negative_prompt.json"))
control_edge = load_video("assets/edge/control_edge.mp4")
videos = pipe(
prompt=json.dumps(json_prompt),
negative_prompt=json.dumps(negative_prompt),
control_videos={"edge": control_edge},
num_frames=121,
height=720,
width=1280,
fps=30.0,
num_inference_steps=35,
guidance_scale=3.0,
control_guidance=1.5,
output="videos",
)
export_to_video(videos, "cosmos3_modular_transfer_edge.mp4", fps=30, macro_block_size=1)Chuyển văn bản thành ảnh và ảnh thành video đã được chưng cất (ít bước)
Các checkpoint distilled ít bước được phục vụ bởi Cosmos3DistilledModularPipeline (các khối: Cosmos3DistilledBlocks); Cosmos3OmniModularPipeline và Cosmos3OmniPipeline cơ sở không hỗ trợ chúng. num_inference_steps được cố định bằng độ dài của cấu hình pipeline distilled_sigmas (từ modular_model_index.json của checkpoint) và guidance_scale bị ép về 1.0 vì guidance đã được tích hợp vào trọng số — việc truyền bất kỳ giá trị nào khác cho một trong hai tham số này sẽ gây ra lỗi, và negative_prompt sẽ được cảnh báo và bỏ qua.
Các checkpoint đã chưng cất FP8 (revision="fp8") không khai báo chính sách độ chính xác hỗn hợp, do đó mọi bước đều duy trì W8A8 bản địa.
Prompts tuân theo cùng cấu trúc JSON mô tả như các model không được chưng cất, do đó văn bản ngắn phải được upsample trước — hãy sử dụng --mode text2image (T2I) hoặc --mode image2video (I2V) như đã mô tả trong Prompt upsampling, sau đó truyền JSON qua json.dumps(...).
import json
import torch
from diffusers import Cosmos3DistilledModularPipeline
from diffusers.utils import export_to_video, load_image
# JSON-upsampled prompt (see "Prompt upsampling" above).
json_prompt = json.load(open("assets/example_t2i_prompt.json"))
repo = "nvidia/Cosmos3-Super-Text2Image-4Step"
pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo, dtype=torch.bfloat16)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
# text-to-image (distilled)
videos = pipe(
prompt=json.dumps(json_prompt),
num_frames=1,
height=720,
width=1280,
output="videos",
)
videos[0].save("cosmos3_distilled_t2i.jpg", format="JPEG", quality=85)
# image-to-video (distilled) — load the I2V repo instead
# JSON-upsampled prompt (see "Prompt upsampling" above); upsampled from the source prompt
# "The right robotic hand picks up the red sphere on the shelf."
json_prompt_i2v = json.load(open("assets/example_i2v_prompt.json"))
repo_i2v = "nvidia/Cosmos3-Super-Image2Video-4Step"
pipe = Cosmos3DistilledModularPipeline.from_pretrained(repo_i2v, dtype=torch.bfloat16)
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda")
image = load_image(
"https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg"
)
videos = pipe(
prompt=json.dumps(json_prompt_i2v),
image=image,
num_frames=189,
height=720,
width=1280,
output="videos",
)
export_to_video(videos, "cosmos3_distilled_i2v.mp4", fps=24, macro_block_size=1)lớp diffusers.Cosmos3OmniModularPipeline
< source >( blocks: diffusers.modular_pipelines.modular_pipeline.ModularPipelineBlocks | None = Nonepretrained_model_name_or_path: str | os.PathLike | None = Nonecomponents_manager: diffusers.modular_pipelines.components_manager.ComponentsManager | None = Nonecollection: str | None = Noneworkflow: str | None = Nonemodular_config_dict: dict[str, typing.Any] | None = Noneconfig_dict: dict[str, typing.Any] | None = None**kwargs )
Một ModularPipeline cho việc tạo nội dung omni của Cosmos 3.
- tất cả
- call
Cosmos3DistilledModularPipeline
lớp diffusers.Cosmos3DistilledModularPipeline
< source >( blocks: diffusers.modular_pipelines.modular_pipeline.ModularPipelineBlocks | None = Nonepretrained_model_name_or_path: str | os.PathLike | None = Nonecomponents_manager: diffusers.modular_pipelines.components_manager.ComponentsManager | None = Nonecollection: str | None = Noneworkflow: str | None = Nonemodular_config_dict: dict[str, typing.Any] | None = Noneconfig_dict: dict[str, typing.Any] | None = None**kwargs )
ModularPipeline cho việc tạo nội dung Cosmos 3 omni đã được chưng cất (ít bước).
Các checkpoint đã được chưng cất tích hợp classifier-free guidance vào trọng số và lấy mẫu theo lịch cố định được đọc từ cấu hình distilled_sigmas của pipeline (được nạp từ modular_model_index.json), do đó guidance_scale và num_inference_steps là cố định và negative_prompt không được hỗ trợ.
- tất cả
- call
CosmosActionCondition
lớp diffusers.CosmosActionCondition
< source >( mode: typing.Literal['policy', 'forward_dynamics', 'inverse_dynamics']chunk_size: intdomain_name: strresolution_tier: int = 480raw_actions: typing.Optional[torch.Tensor] = Noneimage: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, NoneType] = Nonevideo: typing.Union[list, numpy.ndarray, torch.Tensor, NoneType] = Noneview_point: str = 'ego_view' )
Tham số
- mode (
str) — Nhiệm vụ hành động. Một trong số"forward_dynamics"(triển khai video tương lai từ khung hình đầu tiên và chuỗiraw_actionsđã cho),"inverse_dynamics"(suy luận các hành động kết nối các khung hình điều kiện), hoặc"policy"(triển khai đồng thời video tương lai và các hành động từ khung hình đầu tiên). - chunk_size (
int) — Số bước chuyển tiếp hành động trong chunk. Video điều kiện tương ứng trải dàichunk_size + 1khung hình. - domain_name (
str) — Lĩnh vực hiện thân (embodiment domain) dùng để chọn trọng số chiếu hành động nhận thức theo lĩnh vực. Phải là một trong các lĩnh vực hiện thân Cosmos 3 đã đăng ký. Tham số này cũng xác định chiều hành động không đệm (unpadded action width) được dùng để cắt các hành động dự đoán, được giải quyết nội bộ từ tên này (xem_EMBODIMENT_TO_RAW_ACTION_DIM). - resolution_tier (
int, mặc định là480) — Phân cấp độ phân giải tier cho điều kiện hành động (một trong số256,480,704,720). Phân cấp này chọn một khung hình có sẵn với tỷ lệ khung hình gần nhất với đầu vào; đầu vào sẽ được thu nhỏ (không bao giờ phóng to) và được đệm vào trong khung hình đó để phục vụ cho việc điều kiện. Đây không phải là kích thước khung hình đầu ra, vốn bám theo nội dung đầu vào. Hãy khớp phân cấp với độ phân giải gốc của đầu vào: một phân cấp thấp hơn sẽ làm mất chi tiết, trong khi một phân cấp cao hơn không bổ sung độ phân giải (không phóng to), lãng phí tài nguyên tính toán cho phần đệm, và gây ra sự không khớp giữa huấn luyện và suy luận, có thể làm giảm chất lượng. - raw_actions (
torch.Tensor, tùy chọn) — Các vector hành động miền thô có hình dạng[T, raw_action_dim]điều khiển"forward_dynamics". Các chuỗi ngắn hơnchunk_sizesẽ lặp lại hành động cuối cùng; các chuỗi dài hơn sẽ bị cắt bớt. Các kênh vượt quáaction_dimcủa mô hình sẽ bị từ chối, và các đầu vào hẹp hơn sẽ được đệm bằng số 0 lên đếnaction_dim. - image (
PIL.Image.Image,np.ndarrayhoặctorch.Tensor, tùy chọn) — Khung điều kiện cho"policy"/"forward_dynamics". Loại trừ lẫn nhau vớivideo. - video (
list,np.ndarray, hoặctorch.Tensor, tùy chọn) — Video điều kiện, bắt buộc cho"inverse_dynamics". Với"policy"/"forward_dynamics", chỉ khung hình đầu tiên được sử dụng. Loại trừ lẫn nhau vớiimage. - view_point (
str, mặc định là"ego_view") — Nhãn góc nhìn camera được dùng để điền vào trườngcinematography.framingcủa chú thích hành động. Một trong số"ego_view","third_person_view","wrist_view"hoặc"concat_view". Mô hình hành động được huấn luyện trên các chú thích JSON có cấu trúc chứa câu mô tả góc nhìn này; một nhãn không được công nhận sẽ làm mất trường framing (kèm theo cảnh báo).
Nhóm tất cả các đầu vào cần thiết cho một tác vụ tạo có điều kiện hành động của Cosmos 3.
Truyền điều này vào Cosmos3OmniPipeline.__call__() thông qua đối số action thay vì các đối số cấp cao image / height / width, vốn được dành riêng cho các lần chạy t2v, i2v.
Cosmos3OmniPipelineOutput
lớp diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput
< source >( video: Anysound: typing.Optional[torch.Tensor] = Noneaction: list[torch.Tensor] | None = None )
Tham số
- video — Video được tạo ra. Kiểu chính xác phụ thuộc vào
output_typeđược truyền vào pipeline: một danh sách các khung hình PIL cho"pil"(mặc định), mộtnp.ndarraycó hình dạng[T, H, W, C]cho"np", mộttorch.Tensorcó hình dạng[T, C, H, W]cho"pt", hoặc một tensor latent thô khioutput_type="latent". - sound — Dạng sóng âm thanh đã giải mã có hình dạng
[C, N].Nonekhienable_sound=False. - action — Các token hành động dự đoán.
Nonetrừ khi một chế độ hành động dự đoán các hành động.
Dataclass đầu ra cho Cosmos3OmniPipeline.