Diffusers documentation

Cosmos 3

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Cosmos 3

NVIDIA Cosmos 3는 Physical AI를 위한 통합 월드 파운데이션 모델(WFM)입니다. 월드 생성, 물리적 추론, 액션 생성을 결합한 단일 옴니 모델로, 이전 Cosmos 릴리스의 개별 Predict, Reason, Transfer 모델을 대체합니다. 로봇, 자율주행차, 스마트 스페이스를 위한 개발을 진행하든, Cosmos 3는 물리적 세계를 시뮬레이션하고 이해하는 데 필요한 하나의 기반을 제공합니다.

이번 릴리스에 포함된 내용:

  • 모델 카드와 라이선스가 포함된 Hugging Face Hub의 모델
  • 생성 파이프라인을 위한 Cosmos 3 Diffusers 통합(이 페이지)
  • Post-training 스크립트를 사용하여 자체 데이터로 Cosmos 3를 미세 조정
  • Physical AI를 위한 오픈 합성 데이터 생성(SDG) 데이터셋

Cosmos 3의 새로운 기능

이전 Cosmos 릴리스와 가장 큰 차이점은 Cosmos 3가 Mixture-of-Transformers(MoT) 아키텍처를 기반으로 구축된 omni-model이라는 점입니다. 이전에는 월드 생성(Predict), 제어 생성(Transfer), 장면 이해(Reason), 액션-정책 생성을 위해 각각 별도의 모델을 사용했습니다. Cosmos 3는 이러한 모든 기능을 하나의 모델로 통합하여, 단일 forward pass 내에서 크로스 모달리티로 추론하고 생성합니다.

하나의 모델에서 다음을 수행할 수 있습니다:

  • 텍스트, 이미지 또는 동작 입력(이미지-비디오, 텍스트-비디오, 동작 조건부 비디오 생성)으로부터 물리적으로 타당한 비디오 세계를 생성합니다.
  • 운동, 인과관계, 공간적 관계와 같은 물리적 속성에 대해 추론합니다.
  • 현재 상태에서 미래의 비디오 및 동작 시퀀스를 예측합니다.
  • 구조적 제어를 통해 다양한 관점과 조건 사이에서 장면을 전송하세요 (출시 예정).

내부적으로 단일 Cosmos3OmniTransformer가 Qwen 스타일의 언어 모델과 확산(diffusion) 생성 경로를 병렬로 실행합니다. 텍스트 토큰은 인과적(causal) “이해” 스트림을 통해 흐르고, 비디오 및 사운드 잠재(latent)는 양방향으로 주의(attention)를 기울이는 “생성” 스트림을 통해 흐르며, 3D 멀티모달 RoPE로 결합됩니다. 아키텍처 배경에 대해서는 Cosmos World Foundation Model Platform 논문을 참고하세요.

사용 가능한 체크포인트

Hub에는 두 개의 체크포인트가 공개되어 있습니다 — nvidia/Cosmos3-Nano(작고 빠름)와 nvidia/Cosmos3-Super(크고 고품질). 동일한 파이프라인 클래스가 텍스트-이미지, 텍스트-비디오, 이미지-비디오, 그리고 (사운드 지원 체크포인트와 함께) 텍스트+이미지-사운드 포함 비디오 생성을 지원합니다. 아래 각 워크플로우에서 저장소를 선택하고 모델별 탭을 사용하세요.

Scheduler의 가이드를 확인하여 스케줄러의 속도와 품질 사이의 트레이드오프를 탐색하는 방법을 배우고, 파이프라인 간 구성 요소 재사용 섹션을 확인하여 동일한 구성 요소를 여러 파이프라인에 효율적으로 로드하는 방법을 배워보세요.

FP8 혼합 W8A8/W8A16 노이즈 제거

공식 ModelOpt FP8 체크포인트는 Hub의 fp8 리비전에 있습니다(예: revision="fp8"를 사용한 nvidia/Cosmos3-Nano).

이 모든 체크포인트는 동일한 방식으로 양자화됩니다. W8A8은 8비트 가중치와 8비트 활성화를 사용합니다(복원된 ModelOpt GEMM). W8A16은 동일한 8비트 가중치를 재사용하되 활성화 양자화를 생략합니다: FP8 가중치는 역양자화되고, 표준 선형 연산이 BF16/FP16/FP32 활성화에서 수행됩니다.

모든 단계에서 W8A8을 실행하면 다단계 비디오 생성 시 눈에 띄는 깜빡임이 발생할 수 있습니다. 따라서 비디오 Nano / Super / Super-I2V FP8 체크포인트는 transformer/config.json에 첫 3단계와 마지막 3단계에서 W8A16, 중간에서 W8A8을 사용하는 스케줄을 선언합니다. Diffusers는 이러한 개수를 하드코딩하지 않고 체크포인트에서 읽습니다. 스케일러 단계마다 정밀도가 한 번만 선택되도록 하여 classifier-free guidance의 cond/uncond 호출이 일치하도록 합니다.

이미지 생성 및 few-step distilled checkpoint에서는 이러한 flickering이 나타나지 않으므로, Super-T2I와 distilled 4-step FP8 리포지토리는 스케줄을 선언하지 않고 모든 단계에서 W8A8을 유지합니다. 스케줄은 ModelOpt FP8 전용입니다. 다른 양자화 백엔드(예: TorchAO)는 자체 네이티브 포워드를 유지합니다.

fp8 리비전을 ModelOpt 가이드와 동일한 복원 경로로 로드합니다(revision="fp8"에는 이미 양자화 구성이 포함되어 있습니다). 혼합 정밀도는 자동으로 체크포인트를 따릅니다:

import torch
from diffusers import Cosmos3OmniPipeline

pipe = Cosmos3OmniPipeline.from_pretrained(
    "nvidia/Cosmos3-Nano",
    revision="fp8",
    dtype=torch.bfloat16,
    device_map="cuda",
)
result = pipe(prompt="...", num_inference_steps=35)

생성 시 두 가지 선택지:

  • 기본값 (mixed_precision_format=None): 체크포인트가 diffusion_step_policy를 선언한 경우, 처음/마지막 N단계에서는 W8A16을, 중간 단계에서는 네이티브 W8A8을 실행합니다. 이는 다단계 video FP8에 대한 의도된 레시피입니다(전체 W8A8보다 플리커링이 적음). 디스틸드 4단계 및 Super-T2I FP8은 정책을 생략하므로 기본값이 이미 전체 W8A8입니다.
  • mixed_precision_format="none": 모든 스텝을 네이티브 W8A8로 유지합니다. 복원된 FP8 GEMM 대신 W8A16가 dequant + torch.nn.functional.linear를 사용하므로 더 빠르지만, 다중 스텝 비디오에서 깜빡임이 발생할 수 있습니다. 스케줄을 A/B 테스트하거나 완전히 양자화된 베이스라인과 일치시키기 위해 이 옵션을 사용하세요.

Blackwell 워크스테이션 1대에서, Cosmos3-Nano @fp8를 720×1280 / 35 스텝으로 실행할 때, 기본 혼합 스케줄은 "none" 대비 (T2I)에서 약 27% 느리고, (49프레임 T2V)에서 13% 느렸습니다. 이 수치는 처리량 보장이 아닙니다. 정책이 없는 ModelOpt FP8 체크포인트에서 첫/마지막 N 스텝 스케줄을 강제하려면 "fp8"를 전달하세요.

프롬프트 업샘플링

Cosmos 3은 길고 매우 상세한 캡션으로 학습되었습니다. 최적의 품질을 위해 짧은 텍스트 프롬프트는 파이프라인에 전달되기 전에 특정 JSON 구조로 업샘플링되어야 합니다. 업샘플러는 cosmos-framework 패키지에 포함되어 있습니다.

짧은 일반 텍스트 프롬프트로 시작하여 assets/prompt.txt에 저장합니다. 아래 텍스트-투-비디오 예시에서 원래 프롬프트는 “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.txt

그 다음 프레임워크를 설치하고 업샘플러를 실행합니다. 아래 예제는 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"

대상 워크플로(text2image, text2video, image2video)에 맞게 --mode를 전환합니다. 이 명령은 업샘플링된 프롬프트를 --output 파일에 JSON 배열로 기록합니다(--input의 비어 있지 않은 각 줄당 하나의 객체). 줄당 하나의 JSON 객체를 얻으려면 대신 .jsonl 경로를 전달합니다. image2video의 경우, --image-url(URL 또는 로컬 경로) 또는 --image-list(프롬프트당 하나의 이미지)을 통해 조건부 이미지를 제공해야 합니다.

편의상 사전 업샘플링된 긍정 프롬프트(assets/example_t2v_prompt.json)와 부정 프롬프트(assets/negative_prompt.json)가 제공되며, 아래 생성 예제에서 사용됩니다. 예제는 이러한 JSON 파일을 로드하여 json.dumps(...)를 통해 JSON 문자열로 파이프라인에 전달합니다.

텍스트-투-비디오

텍스트만으로 조건부 멀티프레임 생성. 목표 길이에 따라 num_frames를 선택하세요. 기본값인 num_frames=189는 24 FPS에서 약 7.9초를 생성합니다. 프롬프트와 네거티브 프롬프트는 Prompt upsampling에서 설명한 JSON 업샘플링 파일에서 읽힙니다.

Nano
훌륭합니다
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)

텍스트-이미지

단일 프레임 생성. 모델은 텍스트 프롬프트에만 조건이 설정되며 num_frames=1을 전달합니다. JSON 프롬프트를 생성하려면 --mode text2image로 업샘플링합니다.

Nano
훌륭합니다
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)

이미지-비디오

image=를 통해 조건화 이미지를 전달합니다. 파이프라인은 0번째 프레임을 제공된 이미지에 고정하고 나머지 프레임을 디노이즈합니다. 이미지는 종횡비를 유지하면서 리사이즈되고, 요청된 출력 크기로 중앙 크롭되며, uint8에 해당하는 반올림으로 정규화됩니다. --mode image2video로 업샘플링하여 JSON 프롬프트를 생성합니다.

Nano
훌륭합니다
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)

비디오-투-비디오

video=를 통해 조건부 클립을 전달합니다(예: load_video에서 가져온 클립). 파이프라인은 condition_frame_indexes_vision(기본값 [0, 1])로 지정된 앞쪽 잠재 프레임(leading latent frames)을 해당 클립에 고정하고, 나머지 프레임을 디노이즈합니다. condition_video_keep("first" 또는 "last")를 사용하여 더 긴 소스 클립의 어느 끝에서 조건부 프레임을 가져올지 선택합니다. 다른 모드와 마찬가지로, 프롬프트는 Prompt upsampling에 설명된 설명형 JSON 구조를 따르도록 작성해야 합니다.

Nano
훌륭합니다
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)

음성이 포함된 비디오-투-비디오

체크포인트에 sound_tokenizer가 포함된 경우, video-to-video 호출에 enable_sound=True를 추가하여 동기화된 오디오 트랙을 함께 생성할 수 있습니다. 웨이브폼은 비디오와 함께 반환되며, encode_video()를 사용하여 MP4에 mux할 수 있습니다.

Nano
훌륭합니다
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",
)

음성이 포함된 텍스트-비디오

sound_tokenizer가 포함된 체크포인트인 경우, enable_sound=True를 전달하여 동기화된 오디오 트랙을 함께 생성합니다. 웨이브폼은 비디오와 함께 반환되며, encode_video()를 사용하여 MP4에 믹스할 수 있습니다.

위 텍스트-투-비디오 예시와 동일한 호출이며 enable_sound=True가 추가되었습니다:

Nano
훌륭합니다
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",
)

액션 조건 생성

Action runs는 각 액션별 입력을 CosmosActionCondition으로 묶어 상위 수준의 image / video / height / width 인수가 아닌 action 인수를 통해 전달합니다. resolution_tier (256/480/704/720)은 입력 비디오의 네이티브 해상도에 가깝게 설정해야 하며, 이는 조건부 캔버스를 선택합니다. Cosmos 3는 policy, forward_dynamics, inverse_dynamics의 세 가지 액션 모드를 지원합니다. policy와 forward_dynamics는 첫 번째 프레임에만 조건을 부여하므로 image와 video 모두 사용할 수 있지만, inverse_dynamics는 video가 필요합니다. 액션 실행의 조건부 비디오는 파이프라인의 상위 수준 video 인수가 아닌 action.video (또는 action.image)에 설정합니다.

prompt에 일반 작업 설명을 전달하고 action.view_point로 카메라를 선택합니다(기본값 "ego_view"; "third_person_view", "wrist_view", "concat_view"도 사용 가능). 파이프라인은 이들을 모델 학습에 사용된 구조화된 JSON 캡션으로 변환하므로, 액션 프롬프트는 LLM 업샘플링을 하지 않아야 합니다.

액션 정책

Action policy generation은 첫 번째 관찰 프레임, 텍스트 프롬프트, 그리고 액션 도메인 메타데이터로부터 미래의 비디오 및 액션 토큰을 예측합니다. 아래 예제는 Bridge 로봇 도메인을 사용하며, 모델 정규화 액션 공간의 예측된 액션 청크를 JSON으로 기록합니다.

Nano
훌륭합니다
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)

SeaCache

SeaCache는 기본적으로 비활성화되어 있습니다. Cosmos 3는 SeaCacheConfig를 사용하여 이를 명시적으로 활성화할 수 있습니다. SeaCache는 Spectral-Evolution-Aware 지표가 느리게 변할 때 transformer 잔차를 재사용하여 전체 transformer 실행 횟수를 줄입니다. 이 지표는 이미지-비디오 생성의 경우 클린 조건 프레임이 포함된 원시 비전 잠재 변수(raw vision latents)에서 계산됩니다. 파이프라인의 스케줄러 메타데이터 콜백을 통해 transformer에서 활성화하십시오:

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))

동일한 모델 레벨 API가 Cosmos3OmniPipeline, Cosmos3OmniModularPipeline, Cosmos3DistilledModularPipeline에서도 동작합니다. SeaCache는 근사적이며 생성된 출력을 변경할 수 있습니다. 모든 denoising 단계에서 전체 transformer가 실행되어야 하는 경우 pipe.transformer.disable_cache()를 사용하여 비활성화할 수 있습니다. 캐시 상태는 각 pipeline 호출 후 초기화되며, 조건부 및 무조건부 guidance 분기는 독립적인 히스토리를 유지합니다.

Cosmos 3은 SeaCache 게이트를 반복되는 디코더 레이어 외부에 유지하므로 지역 컴파일과 호환됩니다. 캐시를 활성화한 후 레이어를 컴파일합니다:

pipe.transformer.compile_repeated_blocks(fullgraph=True)

SeaCache는 아래에 문서화된 Cosmos 3 Ulysses context-parallel 및 DTensor 기반 tensor-parallel 헬퍼도 지원합니다. 캐시 결정은 rank 간에 동기화됩니다. 전체 모델 컴파일 및 기타 모델 샤딩 전략은 지원하지 않습니다.

컨텍스트 병렬성

긴 영상이나 고해상도의 경우, 단일 forward pass가 GPU 한 대의 메모리 및 지연 시간 예산을 초과할 수 있습니다. Cosmos 3은 컨텍스트 병렬성(context parallelism, CP)을 지원하여 시퀀스 차원을 여러 GPU에 걸쳐 샤딩하고, 어텐션 계산을 분할하여 각 디바이스가 토큰의 일부만 보유하도록 합니다.

Cosmos 3은 Ulysses 컨텍스트 병렬성(all-to-all sequence/head 교환)을 지원합니다. Ring attention은 지원하지 않습니다.

대부분의 diffusers 모델과 달리 Cosmos 3은 CP를 transformer나 선언형 enable_parallelism() 경로에 연결하지 않습니다. 그룹 쿼리 어텐션, 독립적인 이해/생성 스트림(생성 스트림은 두 스트림 모두에 어텐션을 수행), 스트림별 불규칙한 길이 등은 _cp_plan으로 표현할 수 없기 때문입니다. 대신 모델은 작은 no-op shard/gather 시ーム을 노출하고, 실제 구현은 examples/cosmos3/cosmos_parallel.py에 있습니다. 처음부터 끝까지 읽고 수정할 수 있는 자체 완결형 모듈입니다. 두 가지 직교적이고 조합 가능한 샤딩 축을 제공합니다:

헬퍼Shards사용
enable_cosmos3_context_parallel(transformer, cp_mesh)시퀀스(CP / Ulysses)GPU 1개에 맞는 모델(Nano)의 지연 시간
enable_cosmos3_tensor_parallel(transformer, tp_mesh)가중치 (TP)하나의 GPU에 맞지 않는 모델 학습 (Super)

2-D (tp, cp) 메시(텐서 병렬을 통한 대형 모델 피팅 참조)에서 각각 단독으로 또는 함께 사용할 수 있습니다.

Cosmos 3에 특화된 두 가지 요구 사항이 있습니다:

  • native attention 백엔드를 사용하세요. Cosmos 3는 grouped-query attention(GQA)을 사용하며, enable_gqa를 허용하는 유일한 백엔드는 native SDPA입니다(cuDNN과 flash는 이를 거부합니다). 헬퍼 함수들은 KV 헤드를 쿼리 헤드 수로 확장한 뒤 enable_gqa=False로 SDPA를 호출하여 여전히 flash 커널로 디스패치되도록 합니다(수학 폴백을 사용하면 전체 [S, S] 점수가 메모리에 할당되어 긴 시퀀스에서 OOM이 발생합니다).
  • CP(Ulysses) 도는 query-head 수(32 for Nano, 64 for Super)를 나누어야 하며, TP의 경우 도는 KV heads(8)를 나누어야 합니다. 이해(text) 및 생성(video/sound) 스트림은 시퀀스 방향으로 독립적으로 샤딩되며, 불규칙한 길이는 내부적으로 world size의 배수로 zero-padding됩니다.

실행

전체 CLI examples/cosmos3/inference_cosmos3.py는 Cosmos3OmniModularPipeline를 사용하고 이러한 헬퍼를 재사용하므로, 모든 모달리티(텍스트-이미지/비디오, 이미지-비디오, 사운드, 액션 모드)가 --tp-degree / --cp-degree를 통해 멀티 GPU로 실행됩니다. torchrun으로 시작하며, --tp-degree * --cp-degree는 --nproc_per_node와 같아야 합니다. 모든 랭크가 동일한 출력을 생성하며, 랭크 0이 이를 기록합니다.

# 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의 약 120 GB의 가중치는 96 GB GPU 하나에 맞지 않으므로 TP가 필요합니다. Nano는 단일 GPU에 맞으므로 CP는 순수한 지연 시간 최적화입니다. (두 플래그를 모두 생략하면 단일 GPU로 실행됩니다.)

텐서 병렬성을 활용한 대규모 모델 학습

CP는 activations만 샤딩하고 모든 rank에 가중치를 복제하므로 모델의 가중치 footprint를 줄이지 않습니다. 단일 GPU에 맞지 않는 모델은 CP만으로는 여전히 맞지 않습니다. weights를 샤딩하려면, enable_cosmos3_tensor_parallel(transformer, tp_mesh)는 두 번째 직교 메쉬 축에서 Megatron 스타일 tensor parallelism을 적용합니다:

  • Attention 및 MLP 투영은 TP 그룹에 걸쳐 열/행 샤딩됩니다(to_q/to_k/to_v + add_q/k/v 및 MLP의 gate/up은 열 병렬, to_out/to_add_out 및 MLP의 down은 all-reduce를 사용하는 행 병렬). 각 랭크는 query_heads / tp개의 query head와 kv_heads / tp개의 KV head를 소유하게 됩니다.
  • TP는 2차원 (tp, cp) 디바이스 메시에서 CP와 조합됩니다. TP는 헤드/가중치를 영구적으로 분할하고, CP는 그 위에 시퀀스를 샤딩합니다. 제약 조건은 tp가 KV 헤드(8)를 나누어야 하고, tp * cp가 쿼리 헤드를 나누어야 한다는 것입니다(Nano의 경우 32, Super의 경우 64).
  • Weights는 CPU로 로드된 후 레이어별로 GPU에 샤드되므로, 전체 모델이 단일 디바이스에 완전히 생성되지 않습니다.

TP은 모든 attention 및 MLP 블록에서 all-reduce를 수행하므로 대역폭 부하가 큽니다. NVLink가 없는 호스트에서는 이것이 지배적인 비용이 되므로, 가중치가 들어가는 가장 작은 TP degree를 사용하고 나머지 GPU는 CP에 할당하는 것이 좋습니다.

…

CLI 플래그는 편리하지만, Cosmos3OmniModularPipeline로 헬퍼를 직접 호출할 수도 있습니다. 파이프라인 구성과 컴포넌트를 CPU에 로드한 후, 파이프라인을 랭크 로컬 GPU로 이동하기 전에 TP를 적용하고, native 백엔드로 전환한 다음 CP를 활성화합니다. 이 플로우에는 device_map을 사용하지 마세요:

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.

CP 전용(가중치 샤딩 없음)의 경우 1-D 메시를 사용하세요: init_device_mesh("cuda", (world_size,), mesh_dim_names=("cp",))을 초기화하고 enable_cosmos3_context_parallel만 활성화하세요.

enable_safety_checker()는 기본 체크러를 로드하고 활성화하며, disable_safety_checker()는 이를 명시적으로 비활성화합니다. 태스크 파이프라인의 enable_safety_checker= 생성 인자나 enable_safety_check= 호출 인자 대신 해당 파이프라인 메서드를 사용하세요. 모듈형 파이프라인은 Cosmos3OmniPipelineOutput를 반환하지 않습니다. result.video, result.sound, result.action 대신 프레임만 필요할 경우 output="videos"를 사용하거나, 위와 같이 출력 목록과 그 반환된 딕셔너리를 사용하세요.

일부 멀티 GPU 토폴로지에서는 첫 번째 NCCL all-to-all이 멈출 수 있습니다. CP 실행이 첫 번째 denoising 단계 시작 시점에 멈추면, torchrun을 실행하기 전에 환경에서 NCCL_P2P_DISABLE=1을 설정하세요.

CP and TP compose with all the workflows above (text-to-video, image-to-video, text-to-video with sound, and action-conditioned generation) and with both the Nano and Super checkpoints — only the pipeline construction and the parallelism setup lines change.

메타데이터 템플릿

tokenize_prompt는 사용자 메시지 내부에 짧은 메타데이터 문장을 추가하여 LLM이 모델 학습 시 사용된 조건을 인식할 수 있도록 합니다. 긍정 프롬프트에는 “비디오의 길이는 7.9초이며 24 FPS입니다.”와 “이 비디오의 해상도는 720x1280입니다.”와 같은 문장이, 부정 프롬프트에는 그 반대(“… 은/는 아닙니다.”)가 포함됩니다.

둘 다 기본적으로 활성화되어 있습니다. __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은 num_frames == 1(이미지 모드)일 때 효과가 없으며, 이 경우 해상도 문장만 추가됩니다.

안전 검사기

Cosmos3은 cosmos_guardrail의 CosmosSafetyChecker를 연결하고 기본적으로 실행합니다. 텍스트 가드레일은 생성 전에 안전하지 않은 프롬프트를 거부하며(ValueError), 비디오 가드레일은 디코딩된 프레임에서 실행되어 감지된 얼굴을 픽셀화하거나 출력을 거부합니다. 오디오 출력에는 가드레일이 적용되지 않습니다.

기본 체크기를 활성화하려면 선택적 종속성을 설치하세요:

pip install cosmos_guardrail

NVIDIA Open Model License Agreement에 따라 checker는 필수입니다. 아래 두 플래그는 가드레일이 중복되는 테스트 및 개발 워크플로우(예: 입력이 이미 정리된 경우, 또는 의도적으로 엣지 입력으로 파이프라인을 실행하는 경우)를 위해 존재합니다.

생성 시 비활성화(checker가 인스턴스화되지 않으므로 guardrail 모델이 다운로드되거나 메모리에 로드되지 않음):

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,
)

단일 호출에 대해 비활성화(checker는 로드된 상태로 유지됩니다 — 기본값은 유지하면서 일시적으로 우회하는 데 유용합니다):

result = pipe(
    prompt=prompt,
    num_frames=189,
    height=720,
    width=1280,
    fps=24.0,
    enable_safety_check=False,
)

사용자 지정 검사기(예: 빠른 테스트를 위한 no-op 하위 클래스)를 제공하려면 safety_checker=로 전달하세요:

pipe = Cosmos3OmniPipeline.from_pretrained(
    "nvidia/Cosmos3-Nano",
    dtype=torch.bfloat16,
    device_map="cuda",  # or "mps", "xpu", "cpu"
    safety_checker=MyCustomSafetyChecker(),
)

Cosmos3OmniPipeline

클래스 diffusers.Cosmos3OmniPipeline

< >

( 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 )

decode_sound

< >

( latent: Tensor )

소리 잠재 변수 [C, T]를 웨이브폼 [audio_ch, N]으로 디코딩합니다.

sound 토크나이저 디코더가 기대하는 배치 차원을 추가/제거합니다.

prepare_latents

< >

( 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: Iterable = (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 )

단일 샘플에 대한 조건부 생성 및 초기 노이즈를 구성합니다.

tokenize_prompt

< >

( 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 )

prompt-augmentation 템플릿을 적용하고 구성된 채팅 템플릿을 통해 cond/uncond 프롬프트를 토큰화합니다.

이 파이프라인은 별도의 텍스트 인코더를 실행하지 않습니다. 결합된 Cosmos3 트랜스포머가 원시 토큰 ID를 비전(그리고 선택적으로 사운드) 토큰과 함께 소비합니다.

negative_prompt가 None인 경우 빈 문자열이 사용됩니다. Cosmos3 문서 페이지에는 text2video / image2video에 대해 명시적으로 전달할 것을 권장하는 품질 관리용 negative prompt가 문서화되어 있습니다. 활성화된 경우, duration 및 resolution 템플릿은 prompt에, inverse 템플릿은 negative prompt에 추가됩니다.

action_mode가 설정되면 프롬프트는 모델이 학습에 사용된 구조화된 액션 JSON 캡션으로 변환됩니다(자세한 내용은 _build_action_json_prompt 참조). 이때 action_view_point가 framing 필드에 사용됩니다. JSON에 이미 duration/fps/resolution/aspect_ratio가 포함되어 있으므로 평탄한 메타데이터 템플릿은 건너뜁니다.

  • 전체
  • 호출

Cosmos3OmniModularPipeline

Cosmos 3는 Modular Diffusers 파이프라인으로도 사용할 수 있습니다. 태스크 기반의 Cosmos3OmniPipeline은 계속 사용할 수 있으며, 모듈형 파이프라인은 이를 공존하며 동일한 모드(text2image, text2video, image2video, video2video, 액션 조건 생성, transfer(구조적 제어))를 지원하며, 체크포인트가 지원하는 경우 선택적으로 사운드를 포함할 수 있습니다.

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]

ModularPipeline를 통해 로드하고 저장소 설정이 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"
)

특정 Cosmos 모듈형 워크플로를 검사하거나 사용자 정의하려면 available_workflows + get_workflow()를 사용하세요:

available = pipe.blocks.available_workflows
image2video_blocks = pipe.blocks.get_workflow("image2video")

기존 워크플로우 전체에 대한 모듈형 예제

모듈형 파이프라인은 태스크 파이프라인과 동일한 호출 시그니처를 지원합니다. 아래 스니펫은 위에서 보여준 모든 생성 예시(text2video, text2image, image2video, video2video, video2video_sound, text2video_sound, action_policy)를 그대로 반영합니다. Transfer(구조적 제어)는 자체 입력을 가지며, 아래 Modular transfer 섹션에서 별도로 보여줍니다.

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)

모듈식 전송(구조적 제어)

Transfer는 control_videos=를 통해 {hint: video} 매핑으로 전달되는 사전 계산된 제어 비디오(edge, blur, depth, segmentation 또는 world-scenario map)를 따릅니다. 비디오 전용이며(image / video / action / enable_sound 없음), 프롬프트는 사전 업샘플링된 JSON 캡션입니다(Prompt upsampling 참조). 긴 클립은 num_video_frames_per_chunk 단위로 자동 회귀적으로 생성된 후 자동으로 이어 붙여집니다. guidance_scale은 일반적인 텍스트 CFG이며, control_guidance(!= 1.0)은 제어 신호를 추가로 증폭합니다. 힌트별 권장 시작 값:

힌트guidance_scalecontrol_guidanceflow_shift기하
Edge / Blur / Depth3.01.510.030 FPS 기준 121프레임
분할3.02.010.030 FPS 기준 121프레임
월드 시나리오(WSM)1.03.010.0101프레임 @ 10FPS

Diffusers에는 control 자산이 포함되지 않습니다. 완성된 자산(힌트별 control 비디오와 매칭되는 prompt.json, 그리고 공유되는 negative_prompt.json)은 Cosmos cookbook에 있습니다. 아래 edge 예제의 경우, 이를 로컬 assets/ 폴더로 다운로드하세요:

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.json
import 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)

디스틸드(소수 스텝) 텍스트-이미지 및 이미지-비디오

Few-step distilled checkpoints는 Cosmos3DistilledModularPipeline(블록: Cosmos3DistilledBlocks)에 의해 제공됩니다. 기본 Cosmos3OmniModularPipeline과 Cosmos3OmniPipeline은 이를 지원하지 않습니다. num_inference_steps는 체크포인트의 modular_model_index.json에 포함된 distilled_sigmas 파이프라인 구성의 길이에 고정되며, 가이던스가 가중치에 이미 반영되어 있으므로 guidance_scale은 1.0으로 강제됩니다. 이 두 값 중 어느 하나에 다른 값을 전달하면 오류가 발생하며, negative_prompt는 경고가 표시된 후 무시됩니다.

FP8 증류 체크포인트(revision="fp8")는 혼합 정밀도 정책을 선언하지 않으므로 모든 단계가 네이티브 W8A8로 유지됩니다.

Prompts는 비 증류 모델과 동일한 설명형 JSON 구조를 따르므로 짧은 텍스트는 먼저 업샘플링해야 합니다. Prompt upsampling에 설명된 대로 --mode text2image(T2I) 또는 --mode image2video(I2V)를 사용하고, json.dumps(...)를 통해 JSON을 전달합니다.

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)

클래스 diffusers.Cosmos3OmniModularPipeline

< >

( 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 )

Cosmos 3 omni 생성을 위한 ModularPipeline입니다.

  • 전체
  • 호출

Cosmos3DistilledModularPipeline

클래스 diffusers.Cosmos3DistilledModularPipeline

< >

( 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 )

디스틸(소수 스텝) Cosmos 3 omni 생성을 위한 ModularPipeline.

Distilled checkpoints는 classifier-free guidance를 가중치에 고정하여 학습하며, 파이프라인의 distilled_sigmas 설정(modular_model_index.json에서 가져옴)에서 읽은 고정 스케줄로 샘플링합니다. 따라서 guidance_scale와 num_inference_steps는 고정되어 있으며 negative_prompt는 지원되지 않습니다.

  • 전체
  • 호출

CosmosActionCondition

클래스 diffusers.CosmosActionCondition

< >

( 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' )

매개변수

  • mode (str) — 작업 유형입니다. "forward_dynamics" (첫 번째 프레임과 주어진 raw_actions 시퀀스로부터 미래 비디오를 롤아웃), "inverse_dynamics" (조건부 프레임들을 연결하는 액션을 추론), 또는 "policy" (첫 번째 프레임으로부터 미래 비디오와 액션을 함께 롤아웃) 중 하나입니다.
  • chunk_size (int) — chunk의 액션 전이 단계 수입니다. 짝을 이루는 조건부 비디오는 chunk_size + 1 프레임에 걸쳐 있습니다.
  • domain_name (str) — 도메인 인식 동작 투영 가중치를 선택하는 Embodiment 도메인입니다. 등록된 Cosmos 3 embodiment domains 중 하나여야 합니다. 또한 예측된 동작을 슬라이싱하는 데 사용되는 패딩되지 않은 동작 폭을 고정하며, 이 이름에서 내부적으로 해석됩니다( _EMBODIMENT_TO_RAW_ACTION_DIM 참조).
  • resolution_tier (int, 기본값 480) — Action conditioning tier (256, 480, 704, 720 중 하나). tier는 입력과 가장 유사한 종횡비를 가진 미리 정의된 캔버스를 선택합니다. 입력은 다운스케일링(업스케일링은 절대 수행되지 않음)된 후 패딩되어 해당 캔버스에 맞춰 conditioning에 사용됩니다. 이는 입력 콘텐츠에 따라 결정되는 출력 프레임 크기가 아닙니다. tier는 입력의 네이티브 해상도와 일치하도록 설정해야 합니다. 낮은 tier를 사용하면 디테일이 손실되고, 높은 tier를 사용하면 해상도가 증가하지 않으며(업스케일링 없음), 패딩으로 인한 계산 자원이 낭비되고, 학습/추론 불일치가 발생하여 품질 저하를 초래할 수 있습니다.
  • raw_actions (torch.Tensor, 선택 사항) — "forward_dynamics"를 구동하는 [T, raw_action_dim] 형태의 원시 도메인 액션 벡터입니다. chunk_size보다 짧은 시퀀스는 마지막 액션을 반복하고, 더 긴 시퀀스는 잘립니다. 모델의 action_dim을 초과하는 채널은 거부되며, 더 좁은 입력은 action_dim까지 영 패딩됩니다.
  • image (PIL.Image.Image, np.ndarray 또는 torch.Tensor, 선택 사항) — "policy" / "forward_dynamics"의 조건 프레임입니다. video와 상호 배타적입니다.
  • video (list, np.ndarray, 또는 torch.Tensor, 선택 사항) — "inverse_dynamics"에 필요한 조건부 비디오입니다. "policy" / "forward_dynamics"의 경우 첫 번째 프레임만 사용됩니다. image와 상호 배타적입니다.
  • view_point (str, 기본값 "ego_view") — 액션 캡션의 cinematography.framing 필드를 채우는 데 사용되는 카메라 시점 라벨입니다. "ego_view", "third_person_view", "wrist_view", "concat_view" 중 하나여야 합니다. 액션 모델은 이 시점 문장을 포함하는 구조화된 JSON 캡션으로 학습되었습니다. 인식되지 않는 라벨이 사용되면 framing 필드가 삭제되고 경고가 표시됩니다.

Cosmos 3의 동작 조건 생성 작업에 필요한 모든 입력을 그룹화합니다.

Cosmos3OmniPipeline.__call__()에 action 인수를 통해 전달하세요. 최상위 image / height / width 인수는 t2v, i2v 실행용으로 예약되어 있습니다.

Cosmos3OmniPipelineOutput

클래스 diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput

< >

( video: Anysound: typing.Optional[torch.Tensor] = Noneaction: list[torch.Tensor] | None = None )

매개변수

  • video — 생성된 비디오. 정확한 유형은 파이프라인에 전달된 output_type에 따라 달라집니다. 기본값인 "pil"의 경우 PIL 프레임 목록, "np"의 경우 [T, H, W, C] 형태의 np.ndarray, "pt"의 경우 [T, C, H, W] 형태의 torch.Tensor, 그리고 output_type="latent"인 경우 원시 latent 텐서를 반환합니다.
  • sound — [C, N] 형태의 디코딩된 오디오 웨이브폼. enable_sound=False인 경우 None.
  • action — 예측된 액션 토큰. 액션 모드로 액션을 예측하지 않는 한 None입니다.

Cosmos3OmniPipeline의 출력 dataclass입니다.

Update on GitHub