Buckets:

mishig's picture
|
download
raw
75.2 kB

Apache License, Version 2.0(「ライセンス」)の下でライセンスされています。

このファイルは、ライセンスに準拠した場合にのみ使用できます。

ライセンスのコピーは以下から入手できます

http://www.apache.org/licenses/LICENSE-2.0

適用される法律で要求される場合、または書面による合意がある場合を除き、ソフトウェア

ライセンスに基づき配布される本ソフトウェアは、「現状のまま」

明示的または黙示的を問わず、いかなる種類の保証や条件も付随しません。

権限を管理する特定の言語については、ライセンスを参照してください。

ライセンスの下での制限。-->

Cosmos 3

NVIDIA Cosmos 3 は、Physical AI 向けの統合型ワールドファウンデーションモデル(WFM)です。ワールド生成、物理的推論、アクション生成を組み合わせる単一のオムニモデルとして機能します。以前の Cosmos リリースにおける個別の Predict、Reason、Transfer モデルに代わる存在です。ロボティクス、自動運転車、スマートスペースのいずれの開発においても、Cosmos 3 は物理的な世界をシミュレートし理解するための単一の基盤を提供します。

今回のリリースで提供される内容:

  • モデルカードとライセンスを備えた Hugging Face Hub 上のモデル
  • Cosmos 3 Diffusers による生成パイプライン統合(このページ)
  • Cosmos 3 を独自のデータでファインチューニングするためのポストトレーニングスクリプト
  • Physical AI向けのオープンな合成データ生成(SDG)データセット

Cosmos 3の新機能

従来の Cosmos リリースからの最大の変更点は、Cosmos 3 が Mixture-of-Transformers (MoT) アーキテクチャを基盤とするオムニモデルであることです。以前は、ワールド生成 (Predict)、制御された生成 (Transfer)、シーン理解 (Reason)、およびアクションポリシー生成のために、それぞれ別々のモデルを使用していました。Cosmos 3 では、これらすべてが単一のモデルに統合され、単一のフォワードパス内で、モダリティを横断して推論と生成を行います。

1つのモデルから、以下を行うことができます:

  • テキスト、画像、またはアクション入力から物理的に妥当な動画世界を生成します(画像から動画への変換、テキストから動画への変換、アクション条件付き動画生成)。
  • 運動、因果関係、空間的な関係などの物理的性質について推論します。
  • 現在の状態から、将来の動画とアクションのシーケンスを予測します。
  • 構造的な制御を用いて、視点や条件をまたいでシーンを転送します。*(近日公開)*

内部では、単一のCosmos3OmniTransformerが、Qwenスタイルの言語モデルと拡散生成パスを並行して実行します。テキストトークンは因果的な「理解」ストリームを流れ、動画と音声の潜在変数は双方向に注意を払う「生成」ストリームを流れ、3DマルチモーダルRoPEによって結合されます。アーキテクチャの背景については、Cosmos World Foundation Model Platform論文を参照してください。

利用可能なチェックポイント

Hubには2つのチェックポイントが公開されています — nvidia/Cosmos3-Nano(小型で高速)と nvidia/Cosmos3-Super(大型で高品質)です。同じパイプラインクラスが、テキストから画像へ、テキストから動画へ、画像から動画へ、そして(音声対応のチェックポイントを使用する場合)テキスト+画像から音声付き動画への変換をサポートしています。リポジトリを選択し、以下の各ワークフローでモデルごとのタブを使用してください。

スケジューラーのトレードオフ、すなわち速度と品質のバランスをどのように探求するかを学ぶには、ガイドを必ずご確認ください。また、同じコンポーネントを複数のパイプラインに効率的に読み込む方法については、パイプライン間でコンポーネントを再利用するセクションを参照してください。

FP8 混合 W8A8/W8A16 デノイジング

Nano、Super、および Edge 用の ModelOpt FP8 公式チェックポイントは、それぞれ独立した Hub リポジトリとして公開されています: nvidia/Cosmos3-Nano-FP8、nvidia/Cosmos3-Super-FP8、および nvidia/Cosmos3-Edge-FP8。FP8 リポジトリの main を読み込みます。nvidia/Cosmos3-Super-Image2Video には独立した FP8 リポジトリはありません。 revision="fp8" を指定して読み込みます。

これらのチェックポイントはすべて同じ方法で量子化されています。W8A8は8ビットの重みと8ビットの活性化値を使用します(復元されたModelOpt GEMM)。W8A16は同じ8ビットの重みを再利用しますが、活性化値の量子化をスキップします。FP8の重みはデクォンタイズされ、BF16/FP16/FP32の活性化値に対して標準的な線形演算が実行されます。

すべてのステップでW8A8を実行すると、マルチステップの動画生成において目に見えるちらつきが発生する可能性があります。そのため、動画のNano、Super、Super-I2V、およびEdge FP8のチェックポイントは、transformer/config.jsonにスケジュールを宣言しています。最初の3ステップと最後の3ステップではW8A16、中間ではW8A8を使用します。Diffusersはこれらのカウントをハードコードするのではなく、チェックポイントから読み取ります。精度はスケジューラーのステップごとに1回選択されるため、classifier-free guidanceのcond/uncond呼び出しが一致します。

画像生成および少数ステップ蒸留チェックポイントでは、そのようなフリッカーは観測されません。そのため、Super-T2Iおよび蒸留4ステップFP8リポジトリはスケジュールを宣言せず、すべてのステップでW8A8を維持します。また、このスケジュールはModelOpt FP8専用です。その他の量子化バックエンド(TorchAOなど)は、それぞれのネイティブなフォワード処理を維持します。

ModelOpt ガイドと同じ復元パスを持つ FP8 リポジトリを読み込みます。リポジトリにはすでに量子化設定が含まれています。混合精度はチェックポイントに従って自動的に設定されます。

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)

生成時の2つの選択肢:

  • デフォルト (mixed_precision_format=None): チェックポイントが diffusion_step_policy を宣言している場合、最初の N ステップと最後の N ステップでは W8A16 を実行し、中間ステップではネイティブの W8A8 を実行します。これはマルチステップの video FP8 に対して意図されたレシピです(全ステップ W8A8 よりもフリッカーが少ない)。蒸留された 4 ステップおよび Super-T2I FP8 はポリシーを省略しているため、デフォルトではすでに全ステップ W8A8 になります。
  • mixed_precision_format="none": すべてのステップをネイティブの W8A8 で実行します。W8A16 はデクアンタ化と torch.nn.functional.linear による処理であり、復元された FP8 GEMM とは異なるため、より高速です。ただし、マルチステップの動画ではちらつきが発生する可能性があります。スケジュールの A/B テストや、完全に量子化されたベースラインとの比較に使用してください。

Blackwellワークステーション1台で、720×1280 / 35ステップのCosmos3-Nano-FP8では、デフォルトの混合スケジュールを使用した場合、"none"を使用した場合と比較して、(T2I)で約27%遅く、(49フレームのT2V)で約13%遅くなりました。これらの数値はスループットを保証するものではありません。ポリシーが設定されていないModelOpt FP8チェックポイントに対して、最初のNステップ/最後の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"

ターゲットとするワークフローに合わせて、--mode を切り替えてください(text2image、text2video、image2video)。このコマンドは、アップサンプリングされたプロンプトを JSON 配列として --output ファイルに書き出します(--input の空でない各行につき 1 つのオブジェクト)。1 行につき 1 つの JSON オブジェクトを取得するには、代わりに .jsonl パスを指定してください。image2video の場合、--image-url(URL またはローカルパス)または --image-list(プロンプトごとに 1 枚の画像)を介して条件付け画像を指定する必要があります。

利便性のため、事前にアップサンプリングされた正のプロンプト(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 アップサンプリング済みファイルから読み込まれます。

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)

テキストから画像へ

単一フレームの生成。モデルはテキストプロンプトのみで条件付けられます。num_frames=1を渡してください。--mode text2imageでアップサンプリングして、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)

画像から動画へ

image= を介して条件付け画像を渡します。パイプラインはフレーム 0 を指定された画像に固定し、残りのフレームをデノイズします。画像はアスペクト比を維持したままリサイズされ、要求された出力サイズにセンタークロップされ、uint8 相当の丸めによって正規化されます。--mode image2video を使用してアップサンプリングし、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

video= を介して条件付け用のクリップを渡します(load_video から取得するなど)。パイプラインは、condition_frame_indexes_vision(デフォルトは [0, 1])で指定される先頭の潜在フレームをクリップに固定し、残りのフレームをデノイズします。より長いソースクリップから条件付けフレームをどの端から取得するかは、condition_video_keep("first" または "last")で選択できます。他のモードと同様に、プロンプトは Prompt upsampling で説明されている記述的な JSON 構造に従う必要があります。

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

When the checkpoint carries a sound_tokenizer, add enable_sound=True to the video-to-video call to jointly generate a synchronized audio track. The waveform is returned alongside the video and can be muxed into the MP4 with 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",
)

音声付きのテキストから動画へ

チェックポイントにsound_tokenizerが含まれている場合、enable_sound=Trueを渡すことで、同期したオーディオトラックを同時に生成できます。波形は動画とともに返され、encode_video()を使用して MP4 にマックスできます。

これは、上記のテキストから動画への例と同じ呼び出しにenable_sound=Trueを追加したものです。

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

アクション条件付き生成

Action runs は、各アクション固有の入力を CosmosActionCondition にまとめ、トップレベルの image / video / height / width 引数の代わりに action 引数として渡します。resolution_tier (256/480/704/720) は、入力動画のネイティブ解像度に近い値に設定してください。これにより、条件付け用のキャンバスが選択されます。Cosmos 3 は 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 に書き出します。

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 はデフォルトで無効になっています。Cosmos 3 では、SeaCacheConfig を使用して明示的に有効化できます。SeaCache は、Spectral-Evolution-Aware 指標がゆっくりと変化するときに transformer の残差を再利用し、完全な transformer 実行の回数を削減します。この指標は、画像から動画への生成におけるクリーンな条件付けフレームを含む、生のビジョン潜在変数から計算されます。パイプラインからのスケジューラメタデータコールバックを使用して、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は近似であり、生成される出力が変わる可能性があります。すべてのデノイジングステップでトランスフォーマーを完全に実行する必要がある場合は、pipe.transformer.disable_cache()で無効にしてください。キャッシュの状態は各パイプライン呼び出し後にリセットされ、条件付きと非条件付きのガイダンス分岐はそれぞれ独立した履歴を保持します。

Cosmos 3 は SeaCache ゲートを反復デコーダー層の外側に配置しているため、リージョナルコンパイルと互換性があります。キャッシュを有効化した後に、各層をコンパイルします。

pipe.transformer.compile_repeated_blocks(fullgraph=True)

SeaCache は、以下に説明されている Cosmos 3 の Ulysses コンテキスト並列および DTensor ベースのテンソル並列ヘルパーもサポートしています。キャッシュの決定はランク間で同期されます。フルモデルのコンパイルやその他のモデルシャーディング戦略については、サポートを主張していません。

コンテキスト並列

長時間の動画や高解像度の場合、単一のフォワードパスで1つのGPUのメモリおよびレイテンシの予算を超える可能性があります。Cosmos 3は**コンテキスト並列性 (CP)**をサポートしており、シーケンス次元を複数のGPUにシャード化して分割し、アテンションの計算を分割することで、各デバイスがトークンのスライスのみを保持するようにします。

Cosmos 3 は Ulysses コンテキスト並列性(all-to-all によるシーケンス/ヘッドの交換)をサポートしています。リングアテンションはサポートされていません。

ほとんどの diffusers モデルとは異なり、Cosmos 3 は CP を transformer や宣言的な enable_parallelism() パスに組み込んでいません。その grouped-query attention、独立した理解/生成ストリーム(生成ストリームは両方にアテンションを向ける)、およびストリームごとに異なる長さは、_cp_plan として表現できないためです。代わりに、モデルは小さな no-op の shard/gather シームを公開し、実装は examples/cosmos3/cosmos_parallel.py にあります。これは、最初から最後まで読み通して、必要に応じて調整できる自己完結型モジュールです。2つの直交し、組み合わせ可能なシャーディング軸を提供します。

ヘルパー シャード 用途
enable_cosmos3_context_parallel(transformer, cp_mesh) シーケンス(CP / Ulysses) 1つのGPUに収まるモデル(Nano)のレイテンシ
enable_cosmos3_tensor_parallel(transformer, tp_mesh) 重み (TP) 1つのGPUに収まらないモデルの学習(Super)

2次元の(tp, cp)メッシュ(テンソル並列による大規模モデルの適合を参照)で、どちらか一方を単独で使用するか、両方を組み合わせて使用できます。

Cosmos 3 に固有の要件は2つあります:

  • native アテンションバックエンドを使用します。Cosmos 3 はグループクエリアテンション (GQA) を使用しており、enable_gqa を受け入れるのはネイティブ SDPA バックエンドのみです (cuDNN と flash はこれを拒否します)。ヘルパーは KV ヘッドをクエリヘッド数に展開し、enable_gqa=False で SDPA を呼び出すことで、引き続き flash カーネルにディスパッチされます (math フォールバックでは完全な [S, S] スコアがマテリアライズされ、長いシーケンスで OOM になります)。
  • CP (Ulysses) の度合いは、クエリヘッド数(Nano では 32、Super では 64)で割り切れる必要があります。TP の場合、度合いは KV ヘッド数(8)で割り切れる必要があります。理解(テキスト)と生成(動画/音声)のストリームは、シーケンス方向に独立してシャード化され、不揃いな長さは内部でワールドサイズの倍数にゼロパディングされます。

実行する

完全な 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であり、1つの96 GB GPUには収まらないため、TPが必要です。Nanoは単一GPUに収まるため、CPは純粋なレイテンシ最適化です。(両方のフラグを省略すると、単一GPUで実行されます。)

テンソル並列性を用いた大規模モデルの学習

CP は activations をシャードしますが、すべての重みを各ランクに複製するため、モデルの重みのフットプリントを削減することはありません。1 つの GPU に収まらないモデルは、CP だけでは依然として収まりません。weights をシャードするには、enable_cosmos3_tensor_parallel(transformer, tp_mesh) が、第 2 の直交メッシュ軸上で Megatron スタイルのテンソル並列性を適用します:

  • 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 個のクエリヘッドと kv_heads / tp 個の KV ヘッドを所有することになります。
  • TP は 2 次元の (tp, cp) デバイスメッシュ上で CP と組み合わせられます。TP はヘッド/重みを永続的に分割し、CP はその上にシーケンスをシャードします。制約は、tp が KV ヘッド数(8)を割り、tp * cp がクエリヘッド数(Nano では 32、Super では 64)を割りることです。
  • 重みはCPUに読み込まれ、レイヤーごとにGPUにシャード化されるため、完全なモデルが単一のデバイス上にマテリアライズされることはありません。

TP はすべての attention ブロックと MLP ブロックで all-reduce を発行するため、帯域幅の負荷が重くなります。NVLink を備えていないホストでは、これが支配的なコストとなります。重みが収まる最小の TP 度を優先し、残りの 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 次元メッシュを使用します: 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 を返しません。フレームのみを必要とする場合は output="videos" を使用するか、上記に示したように出力リストと返された辞書を使用し、result.video、result.sound、または result.action を使用しないでください。

On some multi-GPU topologies the first NCCL all-to-all can hang. If a CP run stalls at the start of the first denoising step, set NCCL_P2P_DISABLE=1 in the environment before launching torchrun.

CP と TP は、上記のすべてのワークフロー(テキストから動画へ、画像から動画へ、音声付きのテキストから動画へ、アクション条件付き生成)および Nano と Super の両方のチェックポイントと組み合わせることができます。変更されるのは、パイプラインの構築と並列化の設定を行う行のみです。

メタデータテンプレート

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 では、チェッカーの使用が必須です。以下の2つのフラグは、ガードレールが冗長になるテストや開発ワークフロー(例:入力がすでにクリア済みである場合、または意図的にエッジケースの入力に対してパイプラインを実行している場合)のために用意されています。

構築時に無効化(チェッカーがインスタンス化されないため、ガードレールモデルはダウンロードされず、メモリにも読み込まれません):

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

単一の呼び出しに対して無効にする(チェッカーは読み込み状態を維持します。デフォルトを有効にしたまま、一時的にバイパスする場合に便利です):

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

diffusers.Cosmos3OmniPipeline[[diffusers.Cosmos3OmniPipeline]]

diffusers.Cosmos3OmniPipeline(transformer: Cosmos3OmniTransformer, text_tokenizer: AutoTokenizer, vae: AutoencoderKLWan, scheduler: UniPCMultistepScheduler, sound_tokenizer: diffusers.models.autoencoders.autoencoder_cosmos3_audio.Cosmos3AVAEAudioTokenizer | None = None, safety_checker: diffusers.pipelines.cosmos.pipeline_cosmos3_omni.CosmosSafetyChecker | None = None, enable_safety_checker: bool = True, default_use_system_prompt: bool = True, use_native_flow_schedule: bool = False)

ソース

decode_sound[[diffusers.Cosmos3OmniPipeline.decode_sound]]

decode_sound(latent: Tensor)

ソース

音声潜在変数[C, T]を波形[audio_ch, N]にデコードします。

サウンドトークナイザのデコーダが期待するバッチ次元を追加/削除します。

prepare_latents[[diffusers.Cosmos3OmniPipeline.prepare_latents]]

prepare_latents(image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, NoneType] = None, video: typing.Union[list[PIL.Image.Image], torch.Tensor, numpy.ndarray, NoneType] = None, condition_frame_indexes_vision: collections.abc.Iterable[int] = (0, 1), condition_video_keep: typing.Literal['first', 'last'] = 'first', num_frames: int | None = None, height: int | None = None, width: int | None = None, fps: float = 24.0, latents: typing.Optional[torch.Tensor] = None, sound_latents: typing.Optional[torch.Tensor] = None, action_latents: typing.Optional[torch.Tensor] = None, generator: typing.Optional[torch.Generator] = None, device: str = 'cuda', dtype: dtype = torch.bfloat16, enable_sound: bool = False, action: CosmosActionCondition | None = None)

ソース

戻り値:

初期ノイズテンソルと、ビジョン、サウンド、および任意のアクションモダリティ用の条件マスク/メタデータ。

単一サンプル用の条件付けと初期ノイズを構築します。

tokenize_prompt[[diffusers.Cosmos3OmniPipeline.tokenize_prompt]]

tokenize_prompt(prompt: str, negative_prompt: str | None = None, num_frames: int = 189, height: int = 720, width: int = 1280, fps: float = 24.0, use_system_prompt: bool | None = None, add_resolution_template: bool = True, add_duration_template: bool = True, action_mode: str | None = None, action_view_point: str | None = None)

ソース

戻り値:

(cond_input_ids, uncond_input_ids) — このサンプルのトークンIDのリスト。

設定済みのチャットテンプレートを使用して、プロンプト拡張テンプレートを適用し、cond/uncondプロンプトをトークン化します。

このパイプラインでは、個別のテキストエンコーダーを実行しません。ジョイントのCosmos3トランスフォーマーは、ビジョン(および任意でサウンド)トークンと並んで、生のトークンIDを直接消費します。

negative_prompt が None の場合、空の文字列が使用されます。Cosmos3 のドキュメントページには、text2video / image2video で明示的に渡すことを推奨する品質管理用のネガティブプロンプトが記載されています。有効な場合、期間と解像度のテンプレートはプロンプトに、逆テンプレートはネガティブプロンプトに追加されます。

action_modeが設定されている場合、プロンプトは代わりに、モデルが学習に使用した構造化アクションJSONキャプションに変換されます(_build_action_json_promptを参照)。この際、フレーミングフィールドにはaction_view_pointが使用されます。JSONにはduration/fps/resolution/aspect_ratioが既に含まれているため、フラットなメタデータテンプレートはスキップされます。

  • すべて
  • call

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 を通して読み込み、リポジトリの構成でブロッククラスを選択させることもできます。

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} マッピングとして渡される、事前計算された制御動画(エッジ、ぼかし、深度、セグメンテーション、またはワールドシナリオマップ)に従って行われます。これは動画専用であり、image / video / action / enable_sound はサポートされていません。プロンプトは事前アップサンプリング済みの JSON キャプションです(Prompt upsampling を参照)。また、長いクリップは num_video_frames_per_chunk フレームごとのチャンクで自己回帰的に生成され、自動的に結合されます。guidance_scale は通常のテキスト CFG です。control_guidance(!= 1.0)は、制御シグナルをさらに増幅します。各ヒントごとの推奨初期値は以下の通りです。

ヒント guidance_scale control_guidance flow_shift 幾何学
エッジ / ブラー / 深度 3.0 1.5 10.0 30 FPSで121フレーム
セグメンテーション 3.0 2.0 10.0 30 FPSで121フレーム
ワールドシナリオ(WSM) 1.0 3.0 10.0 101フレーム @ 10 FPS

Diffusersにはコントロールアセットは同梱されていません。完成済みのアセット(各ヒントに対応するコントロール動画とprompt.json、および共有のnegative_prompt.json)は、Cosmos cookbookに格納されています。以下のエッジの例では、これらをローカルのassets/フォルダにダウンロードしてください。

base=https://github.com/NVIDIA/cosmos/raw/6e2c69afafd6444c1ca14a021aba93fcc6b7996a/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)

Distilled (few-step) テキストから画像へ、画像から動画へ[[diffusers.Cosmos3OmniModularPipeline]]

Few-step distilled checkpoints は Cosmos3DistilledModularPipeline(ブロック: Cosmos3DistilledBlocks)によって提供されます。ベースの Cosmos3OmniModularPipeline および Cosmos3OmniPipeline ではサポートされていません。num_inference_steps は、チェックポイントの modular_model_index.json にある distilled_sigmas パイプライン設定の長さに固定され、guidance が重みに組み込まれているため guidance_scale は 1.0 に強制されます。どちらかの値を他の値に設定するとエラーが発生し、negative_prompt は警告が表示された上で無視されます。

FP8蒸留チェックポイント(revision="fp8")は混合精度ポリシーを宣言しないため、すべてのステップでネイティブのW8A8が維持されます。

プロンプトは、蒸留されていないモデルと同じ記述的な 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[[diffusers.Cosmos3OmniModularPipeline]]

diffusers.Cosmos3OmniModularPipeline(blocks: diffusers.modular_pipelines.modular_pipeline.ModularPipelineBlocks | None = None, pretrained_model_name_or_path: str | os.PathLike | None = None, components_manager: diffusers.modular_pipelines.components_manager.ComponentsManager | None = None, collection: str | None = None, workflow: str | None = None, modular_config_dict: dict[str, typing.Any] | None = None, config_dict: dict[str, typing.Any] | None = None, **kwargs)

ソース

Cosmos 3 のオムニ生成用の ModularPipeline。

  • すべて
  • call

Cosmos3DistilledModularPipeline[[diffusers.Cosmos3DistilledModularPipeline]]

diffusers.Cosmos3DistilledModularPipeline[[diffusers.Cosmos3DistilledModularPipeline]]

diffusers.Cosmos3DistilledModularPipeline(blocks: diffusers.modular_pipelines.modular_pipeline.ModularPipelineBlocks | None = None, pretrained_model_name_or_path: str | os.PathLike | None = None, components_manager: diffusers.modular_pipelines.components_manager.ComponentsManager | None = None, collection: str | None = None, workflow: str | None = None, modular_config_dict: dict[str, typing.Any] | None = None, config_dict: dict[str, typing.Any] | None = None, **kwargs)

ソース

蒸留(少数ステップ)Cosmos 3 omni 生成用の ModularPipeline です。

蒸留済みチェックポイントは、classifier-free guidance を重みに焼き込み、パイプラインの distilled_sigmas 設定(modular_model_index.json から設定されます)から読み取られる固定スケジュールでサンプリングを行うため、guidance_scale と num_inference_steps は固定されており、negative_prompt はサポートされていません。

  • すべて
  • call

CosmosActionCondition[[diffusers.CosmosActionCondition]]

diffusers.CosmosActionCondition[[diffusers.CosmosActionCondition]]

diffusers.CosmosActionCondition(mode: typing.Literal['policy', 'forward_dynamics', 'inverse_dynamics'], chunk_size: int, domain_name: str, resolution_tier: int = 480, raw_actions: typing.Optional[torch.Tensor] = None, image: typing.Union[PIL.Image.Image, numpy.ndarray, torch.Tensor, NoneType] = None, video: typing.Union[list, numpy.ndarray, torch.Tensor, NoneType] = None, view_point: str = 'ego_view')

ソース

パラメータ:

mode (str) : 実行するアクションタスク。"forward_dynamics"(最初のフレームと指定されたraw_actionsシーケンスから将来の動画をロールアウトする)、"inverse_dynamics"(条件付けフレームを結ぶアクションを推論する)、または"policy"(最初のフレームから将来の動画とアクションを同時にロールアウトする)のいずれかです。

chunk_size (int) : チャンク内のアクション遷移ステップ数。対応する条件付け動画はchunk_size + 1フレームにまたがります。

domain_name (str) : ドメイン認識アクション投影の重みを決定する Embodiment ドメイン。登録済みの Cosmos 3 Embodiment ドメインのいずれかである必要があります。また、予測アクションをスライスするために使用されるパディングなしのアクション幅も、この名前から内部で解決されるため固定されます(_EMBODIMENT_TO_RAW_ACTION_DIM を参照)。

resolution_tier (int、デフォルトは 480) : アクション条件付けの解像度 tier(256、480、704、720 のいずれか)。tier は、入力に最も近いアスペクト比を持つ事前定義済みのキャンバスを選択します。入力はダウンスケール(アップスケールは行われません)され、条件付けのためにそのキャンバスにパディングされます。これは出力フレームサイズではなく、出力は入力コンテンツに追従します。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 のアクション条件付き生成タスクに必要なすべての入力をグループ化します。

トップレベルの image / height / width 引数ではなく、action 引数を通じて Cosmos3OmniPipeline.__call__() に渡してください。これらの引数は t2v、i2v の実行に予約されています。

Cosmos3OmniPipelineOutput[[diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput]]

diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput[[diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput]]

diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput(video: Any, sound: typing.Optional[torch.Tensor] = None, action: list[torch.Tensor] | None = None)

ソース

パラメータ:

video : 生成された動画。正確な型は、パイプラインに渡されるoutput_typeによって異なります。"pil"(デフォルト)の場合はPILフレームのリスト、"np"の場合はnp.ndarray(形状[T, H, W, C])、"pt"の場合はtorch.Tensor(形状[T, C, H, W])、output_type="latent"の場合は生の潜在テンソルです。

sound : 形状が[C, N]のデコード済み音声波形。enable_sound=Falseの場合はNone。

action : 予測されたアクショントークン。アクションモードがアクションを予測する場合を除き、None です。

Cosmos3OmniPipelineの出力データクラスです。

Xet Storage Details

Size:
75.2 kB
·
Xet hash:
b8d4d22bd23a82720b479ea21ed885ea644ff67852baeaed026a82ba5f8265dc

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.