Diffusers documentation
Cosmos 3
Cosmos 3
NVIDIA Cosmos 3 est un modèle de fondation unifié pour le monde (WFM) destiné à l'IA physique — un modèle omni unique qui combine la génération de mondes, le raisonnement physique et la génération d'actions. Il remplace les modèles distincts Predict, Reason et Transfer des versions précédentes de Cosmos : que vous développiez pour la robotique, les véhicules autonomes ou les espaces intelligents, Cosmos 3 vous offre une base unique pour simuler et comprendre le monde physique.
Ce qui est livré avec cette version :
- Modèles sur le Hugging Face Hub avec fiches de modèle et licences
- Intégration de Cosmos 3 Diffusers pour les pipelines de génération (cette page)
- Scripts d'entraînement postérieur pour l'ajustement fin de Cosmos 3 sur vos propres données
- Jeux de données ouverts de génération de données synthétiques (SDG) pour l'IA physique
Nouveautés de Cosmos 3
La plus grande évolution par rapport aux versions précédentes de Cosmos est que Cosmos 3 est un omni-modèle, construit sur une architecture Mixture-of-Transformers (MoT). Auparavant, les développeurs travaillaient avec des modèles distincts pour la génération de mondes (Predict), la génération contrôlée (Transfer), la compréhension de scènes (Reason) et la génération de politiques d'action. Cosmos 3 unifie l'ensemble de ces fonctionnalités dans un seul modèle qui raisonne et génère à travers les modalités lors d'un unique passage avant.
À partir d'un seul modèle, vous pouvez :
- Générer des mondes vidéo physiquement plausibles à partir de texte, d'images ou d'entrées d'action (génération vidéo image-vers-vidéo, texte-vers-vidéo, génération vidéo conditionnée par l'action).
- Raisonnez sur des propriétés physiques telles que le mouvement, la causalité et les relations spatiales.
- Prédisez les séquences futures de vidéos et d’actions à partir de l’état actuel.
- Transférez des scènes entre différents points de vue et conditions avec un contrôle structurel (bientôt disponible).
Sous le capot, une seule instance de Cosmos3OmniTransformer exécute un modèle de langage de type Qwen en parallèle avec un chemin de génération par diffusion : les jetons de texte circulent dans un flux causal d’« compréhension » tandis que les latents vidéo et audio circulent dans un flux de « génération » à attention bidirectionnelle, le tout étant relié par un RoPE multimodal 3D. Consultez le paper de la plateforme Cosmos World Foundation Model pour le contexte architectural.
Points de contrôle disponibles
Deux points de contrôle sont publiés sur le Hub — nvidia/Cosmos3-Nano (plus petit, plus rapide) et nvidia/Cosmos3-Super (plus grand, qualité supérieure). La même classe de pipeline prend en charge le texte-vers-image, le texte-vers-vidéo, l'image-vers-vidéo et (avec un point de contrôle capable de générer du son) le texte+image-vers-vidéo-avec-son — choisissez un dépôt et utilisez l'onglet par modèle dans chaque workflow ci-dessous.
Pensez à consulter le guide des Schedulers pour apprendre à explorer le compromis entre la vitesse et la qualité du scheduler, et à consulter la section réutilisation des composants entre les pipelines pour apprendre à charger efficacement les mêmes composants dans plusieurs pipelines.
Dénbruitage mixte FP8 W8A8/W8A16
Les points de contrôle FP8 ModelOpt officiels pour Nano, Super et Edge sont publiés sous forme de dépôts Hub distincts : nvidia/Cosmos3-Nano-FP8, nvidia/Cosmos3-Super-FP8 et nvidia/Cosmos3-Edge-FP8. Chargez main du dépôt FP8. nvidia/Cosmos3-Super-Image2Video ne dispose pas de dépôt FP8 distinct ; chargez-le avec revision="fp8".
Tous ces points de contrôle sont quantifiés de la même manière. W8A8 utilise des poids en 8 bits et des activations en 8 bits (le GEMM ModelOpt restauré). W8A16 réutilise ces mêmes poids en 8 bits mais saute la quantification des activations : le poids FP8 est déquantifié et une opération linéaire standard est exécutée sur des activations BF16/FP16/FP32.
L’exécution de W8A8 à chaque étape peut provoquer des scintillements visibles dans la génération de vidéos multi-étapes. Les points de contrôle vidéo Nano, Super, Super-I2V et Edge FP8 déclarent donc un planning dans transformer/config.json : W8A16 pour les 3 premières et les 3 dernières étapes, W8A8 en milieu de séquence. Diffusers lit ces compteurs depuis le point de contrôle au lieu de les coder en dur. La précision est choisie une fois par étape du planificateur afin que les appels cond/uncond du guidage sans classifieur correspondent.
La génération d’images et les points de contrôle distillés en peu d’étapes ne présentent pas ce scintillement, si bien que Super-T2I et les dépôts distillés 4 étapes FP8 ne déclarent aucun calendrier et restent en W8A8 à chaque étape. Le calendrier est également exclusif à ModelOpt FP8 : les autres backends de quantification (par exemple TorchAO) conservent leurs forwards natifs.
Chargez un dépôt FP8 avec le même chemin de restauration que le guide ModelOpt. Le dépôt contient déjà la configuration de quantification. La précision mixte suit alors automatiquement le point de contrôle :
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)Deux choix au moment de la génération :
- Par défaut (
mixed_precision_format=None) : si le point de contrôle déclarediffusion_step_policy, exécutez W8A16 sur les N premières/dernières étapes et W8A8 natif en milieu de parcours. C'est la recette prévue pour la vidéo video FP8 multi-étapes (moins de scintillement que le tout-W8A8). Les modèles FP8 distillés 4 étapes et Super-T2I omettent la politique, donc le comportement par défaut est déjà tout W8A8. mixed_precision_format="none": conserve chaque étape en W8A8 natif. Plus rapide, car W8A16 correspond à une déquantification +torch.nn.functional.linearplutôt qu’au GEMM FP8 restauré, mais les vidéos multi-étapes peuvent scintiller. Utilisez ce mode pour comparer le planning ou pour correspondre à une référence entièrement quantifiée.
Sur une seule station de travail Blackwell, Cosmos3-Nano-FP8 à 720×1280 / 35 étapes était environ 27 % plus lent (T2I) et 13 % plus lent (T2V 49 images) avec l’ordonnancement mixte par défaut que avec "none". Ces chiffres ne constituent pas une garantie de débit. Passez "fp8" uniquement pour forcer l’ordonnancement first/last-N sur un point de contrôle ModelOpt FP8 qui n’a pas de politique.
Échantillonnage supérieur des prompts
Cosmos 3 a été entraîné sur des légendes longues et très descriptives. Pour une qualité optimale, les prompts texte courts doivent être échantillonnés vers le haut dans une structure JSON spécifique avant d’être transmis au pipeline. L’échantillonneur se trouve dans le package cosmos-framework.
Commencez par un court prompt en texte brut et enregistrez-le dans assets/prompt.txt. Pour l’exemple de génération de vidéo à partir de texte ci-dessous, le prompt original est « 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.txtInstallez ensuite le framework et exécutez l'upsampler. L'exemple ci-dessous effectue un upsampling pour la génération de vidéo à partir de texte en utilisant 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"Basculez --mode pour correspondre au flux de travail visé (text2image, text2video, image2video). La commande écrit les prompts suréchantillonnés dans le fichier --output sous forme de tableau JSON (un objet par ligne non vide de --input) ; passez plutôt un chemin .jsonl pour obtenir un objet JSON par ligne. Pour image2video, vous devez également fournir l’image de conditionnement via --image-url (une URL ou un chemin local) ou --image-list (une image par prompt).
Un prompt positif pré-échantillonné (assets/example_t2v_prompt.json) et un prompt négatif (assets/negative_prompt.json) sont fournis pour votre confort et sont utilisés dans les exemples de génération ci-dessous. Les exemples chargent ces fichiers JSON et les transmettent au pipeline sous forme de chaînes JSON via json.dumps(...).
Texte vers vidéo
Génération multi-images conditionnée uniquement sur le texte. Choisissez num_frames en fonction de la durée cible — la valeur par défaut num_frames=189 produit ≈ 7,9 s à 24 FPS. Le prompt et le prompt négatif sont lus depuis les fichiers JSON-échantillonnés décrits dans Échantillonnage des prompts.
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)Texte vers image
Génération d'une image unique. Le modèle est conditionné uniquement sur le prompt texte ; passez num_frames=1. Utilisez --mode text2image pour générer le 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)Image vers vidéo
Passez une image de conditionnement via image=. Le pipeline ancre le cadre 0 à l’image fournie et débruite le reste. L’image est redimensionnée en préservant son ratio d’aspect, recadrée au centre à la taille de sortie demandée, et normalisée avec un arrondi équivalent à uint8. Utilisez --mode image2video pour générer le 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)Vidéo-vers-vidéo
Passez un clip de conditionnement via video= (par exemple, issu de load_video). Le pipeline ancre les trames latentes initiales définies par condition_frame_indexes_vision (par défaut [0, 1]) au clip et débruite le reste. Utilisez condition_video_keep ("first" ou "last") pour choisir à quelle extrémité d'un clip source plus long les trames de conditionnement sont prélevées. Comme pour les autres modes, le prompt doit suivre la structure JSON descriptive décrite dans 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)Vidéo-vers-vidéo avec son
Lorsque le point de contrôle contient un sound_tokenizer, ajoutez enable_sound=True à l’appel vidéo-versus-vidéo pour générer conjointement une piste audio synchronisée. La forme d’onde est renvoyée avec la vidéo et peut être multiplexée dans le MP4 à l’aide de 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",
)Texte vers vidéo avec son
Lorsque le point de contrôle contient un sound_tokenizer, passez enable_sound=True pour générer conjointement une piste audio synchronisée. La forme d’onde est renvoyée avec la vidéo et peut être multiplexée dans le MP4 à l’aide de encode_video().
C'est le même appel que l'exemple de texte vers vidéo ci-dessus, avec l'ajout de 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",
)Génération conditionnée par l'action
Action runs group every action-specific input into a CosmosActionCondition passed via the action argument instead of the top-level image / video / height / width arguments. Set resolution_tier (256/480/704/720) close to the input video’s native resolution; it selects the conditioning canvas. Cosmos 3 supports three action modes — policy, forward_dynamics, and inverse_dynamics. policy and forward_dynamics condition only on the first frame (so an image or a video both work), while inverse_dynamics requires a video. The conditioning video for an action run is set on action.video (or action.image), not on the pipeline’s top-level video argument.
Transmettez une simple description de tâche en tant que prompt et sélectionnez la caméra avec action.view_point (par défaut "ego_view" ; également "third_person_view", "wrist_view", "concat_view"). Le pipeline transforme ces éléments en la légende JSON structurée sur laquelle le modèle a été entraîné, les prompts d’action ne doivent donc pas être enrichis par un LLM.
Politique d'action
La génération de politique d’action prédit les futurs jetons vidéo et d’action à partir de la première image d’observation, du prompt texte et des métadonnées du domaine d’action. L’exemple ci-dessous utilise le domaine du robot Bridge et écrit le bloc d’actions prédit au format JSON dans l’espace d’actions normalisé par le modèle.
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 est désactivé par défaut. Cosmos 3 prend en charge son activation explicite via SeaCacheConfig. SeaCache réutilise les résidus du transformer lorsque l'indicateur Spectral-Evolution-Aware évolue lentement, réduisant ainsi le nombre d'exécutions complètes du transformer. Il calcule cet indicateur à partir des latents visuels bruts, y compris les images de conditionnement propres pour la génération d'images en vidéo. Activez-le sur le transformer à l'aide des callbacks de métadonnées du planificateur fournis par le 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))La même API au niveau du modèle fonctionne avec Cosmos3OmniPipeline, Cosmos3OmniModularPipeline et Cosmos3DistilledModularPipeline. SeaCache est approximatif et peut modifier les sorties générées. Désactivez-le avec pipe.transformer.disable_cache() lorsque vous avez besoin que chaque étape de débruitage exécute le transformeur complet. L'état du cache est réinitialisé après chaque appel au pipeline, et les branches de guidage conditionnel et non conditionnel conservent des historiques indépendants.
Cosmos 3 conserve la porte SeaCache en dehors de ses couches décodeur répétées, ce qui le rend compatible avec la compilation régionale. Compilez les couches après avoir activé le cache :
pipe.transformer.compile_repeated_blocks(fullgraph=True)SeaCache prend également en charge les utilitaires de parallélisme de contexte Ulysses et de parallélisme de tenseurs basés sur DTensor pour Cosmos 3, documentés ci-dessous. Les décisions de cache sont synchronisées entre les rangs. La compilation du modèle complet et les autres stratégies de fractionnement de modèle ne sont pas revendiquées.
Parallélisme de contexte
Pour les vidéos longues ou les hautes résolutions, un seul passage avant peut dépasser le budget mémoire et latence d'un GPU. Cosmos 3 prend en charge le parallélisme de contexte (CP) pour répartir la dimension de séquence sur plusieurs GPU, en divisant le calcul d'attention afin que chaque appareil ne conserve qu'une portion des jetons.
Cosmos 3 prend en charge le parallélisme de contexte Ulysses (échange de séquences/têtes all-to-all). L'attention en anneau n'est pas prise en charge.
Contrairement à la plupart des modèles diffusers, Cosmos 3 ne branche pas le CP dans le transformeur ni dans le chemin déclaratif enable_parallelism() : son attention à requêtes groupées, ses flux de compréhension/génération distincts (le flux de génération fait attention aux deux) et ses longueurs irrégulières par flux ne peuvent pas être exprimés sous forme de _cp_plan. À la place, le modèle expose de petites coutures de shard/gather sans effet, et l’implémentation se trouve dans examples/cosmos3/cosmos_parallel.py — un module autonome que vous pouvez lire de bout en bout et adapter. Il propose deux axes de sharding orthogonaux et composables :
| Assistant | Shards | Utiliser pour |
|---|---|---|
enable_cosmos3_context_parallel(transformer, cp_mesh) | séquence (CP / Ulysses) | latence sur un modèle qui tient sur une seule GPU (Nano) |
enable_cosmos3_tensor_parallel(transformer, tp_mesh) | poids (TP) | entraîner un modèle qui ne tient pas sur une seule GPU (Super) |
Utilisez l’un ou l’autre seul, ou les deux ensemble, sur un maillage 2-D (tp, cp) (voir Adapter de grands modèles avec le parallélisme de tenseurs).
Deux exigences sont spécifiques à Cosmos 3 :
- Utilisez le backend d’attention
native. Cosmos 3 utilise l’attention à requêtes groupées (GQA), et le backend SDPA natif est le seul à accepterenable_gqa(cuDNN et flash le rejettent). Les fonctions utilitaires étendent les têtes KV au nombre de têtes de requête et appellent SDPA avecenable_gqa=Falseafin qu’il se distribue toujours vers le noyau flash (le repli mathématique matérialiserait les scores complets[S, S]et provoquerait une erreur de mémoire sur les séquences longues). - Le degré de CP (Ulysses) doit diviser le nombre de têtes de requête (32 pour
Nano, 64 pourSuper) ; pour le TP, le degré doit diviser les têtes KV (8). Les flux de compréhension (texte) et de génération (vidéo/son) sont fragmentés indépendamment le long de la séquence, et les longueurs irrégulières sont complétées par des zéros en interne jusqu'à un multiple de la taille du monde.
L'exécuter
Le CLI complet examples/cosmos3/inference_cosmos3.py utilise Cosmos3OmniModularPipeline et réutilise ces utilitaires, de sorte que toute modalité (texte-vers-image/vidéo, image-vers-vidéo, son, modes d'action) s'exécute sur plusieurs GPU via --tp-degree / --cp-degree. Lancez avec torchrun ; --tp-degree * --cp-degree doit être égal à --nproc_per_node. Chaque rang produit la même sortie ; le rang 0 l'écrit.
# 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 "..."Les ~120 Go de poids de Super ne tiennent pas sur un seul GPU de 96 Go, il faut donc du TP ; Nano tient sur un GPU unique, donc le CP pour lui est une pure optimisation de latence. (Omettez les deux drapeaux pour exécuter en mono-GPU.)
Ajustement de grands modèles avec le parallélisme de tenseurs
CP fragmente les activations mais réplique chaque poids sur chaque rang, ce qui ne réduit pas l’empreinte mémoire des poids d’un modèle — un modèle qui ne tient pas sur une seule GPU ne tiendra toujours pas avec CP seul. Pour fragmenter les poids, enable_cosmos3_tensor_parallel(transformer, tp_mesh) applique le parallélisme de tenseurs de type Megatron sur un second axe de maillage orthogonal :
- Les projections d’attention et de MLP sont shardées en colonnes/lignes au sein du groupe TP (
to_q/to_k/to_v+add_q/k/vet lesgate/updes MLP sont parallèles en colonnes ;to_out/to_add_outet lesdowndes MLP sont parallèles en lignes avec un all-reduce). Chaque rank possède au finalquery_heads / tptêtes de requête etkv_heads / tptêtes KV. - TP se combine avec CP sur un maillage d’appareils 2D
(tp, cp): TP divise de manière persistante les têtes/poids, tandis que CP fragmente la séquence par-dessus. Les contraintes sont quetpdivise les têtes KV (8) et quetp * cpdivise les têtes de requête (32 pourNano, 64 pourSuper). - Les poids sont chargés sur le CPU et répartis sur les GPU couche par couche, de sorte que le modèle complet n'est jamais matérialisé sur un seul appareil.
TP émet un all-reduce sur chaque bloc d’attention et de MLP, ce qui est très gourmand en bande passante. Sur les hôtes sans NVLink, c’est le coût dominant ; préférez le degré de TP le plus petit qui permette aux poids de tenir en mémoire et placez les GPU restants dans CP.
L'utiliser dans votre propre pipeline modulaire
Les drapeaux CLI sont pratiques, mais vous pouvez appeler les helpers directement avec Cosmos3OmniModularPipeline. Chargez la configuration du pipeline et les composants sur le CPU, appliquez le TP avant de déplacer le pipeline sur le GPU local au rang, passez au backend native, puis activez le CP. N'utilisez pas device_map pour ce flux :
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.Pour CP uniquement (sans partage de poids), utilisez un maillage 1-D : init_device_mesh("cuda", (world_size,), mesh_dim_names=("cp",)) et simplement enable_cosmos3_context_parallel.
enable_safety_checker() charge et active le vérificateur par défaut ; disable_safety_checker() le désactive explicitement. Utilisez ces méthodes de pipeline plutôt que l’argument de construction enable_safety_checker= du pipeline de tâche ou l’argument d’appel enable_safety_check=. Les pipelines modulaires ne renvoient pas non plus Cosmos3OmniPipelineOutput : utilisez output="videos" pour les seules images, ou une liste de sorties et son dictionnaire renvoyé comme indiqué ci-dessus, au lieu de result.video, result.sound ou result.action.
Sur certaines topologies multi-GPU, le premier all-to-all NCCL peut bloquer. Si une exécution CP stagne au début de la première étape de débruitage, définissez
NCCL_P2P_DISABLE=1dans l’environnement avant de lancertorchrun.
CP et TP se combinent avec tous les flux de travail ci-dessus (texte-vers-vidéo, image-vers-vidéo, texte-vers-vidéo avec son, et génération conditionnée par l'action) ainsi qu'avec les points de contrôle Nano et Super — seules les lignes de construction du pipeline et de configuration du parallélisme changent.
Modèles de métadonnées
tokenize_prompt ajoute de courtes phrases de métadonnées dans le message utilisateur afin que le LLM perçoive les conditions avec lesquelles le modèle a été entraîné. Le prompt positif reçoit des phrases comme « La vidéo dure 7,9 secondes et est de 24 FPS. » et « Cette vidéo est d’une résolution de 720x1280. » ; le prompt négatif reçoit l’inverse (« … n’est pas … »).
Les deux sont activés par défaut. Désactivez l'une ou l'autre paire via __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 n'a aucun effet lorsque num_frames == 1 (mode image) ; dans ce cas, seule la phrase de résolution est ajoutée.
Vérificateur de sécurité
Cosmos3 intègre le cosmos_guardrail CosmosSafetyChecker et l'exécute par défaut. Le garde-fou texte rejette les invites non sûres avant la génération (ValueError) ; le garde-fou vidéo s'exécute sur les images décodées et applique soit un floutage des visages détectés, soit un rejet de la sortie. La sortie audio n'est pas protégée par un garde-fou.
Installez la dépendance facultative pour activer le vérificateur par défaut :
pip install cosmos_guardrailLe vérificateur est obligatoire dans le cadre de l’accord de licence des modèles ouverts NVIDIA. Les deux drapeaux ci-dessous existent pour les flux de travail de test et de développement où la garde-fou serait redondante (par exemple, si l’entrée a déjà été nettoyée, ou si vous exercez intentionnellement le pipeline sur des entrées limites).
Désactiver à la construction (aucun vérificateur n'est instancié, donc aucun modèle de garde-fou n'est téléchargé ni chargé en mémoire) :
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,
)Désactiver pour un seul appel (le vérificateur reste chargé — utile pour un contournement ponctuel tout en conservant l’activation par défaut pour les appels suivants) :
result = pipe(
prompt=prompt,
num_frames=189,
height=720,
width=1280,
fps=24.0,
enable_safety_check=False,
)Pour fournir un vérificateur personnalisé (par exemple, une sous-classe sans opération pour des tests rapides), passez-le en tant que safety_checker= :
pipe = Cosmos3OmniPipeline.from_pretrained(
"nvidia/Cosmos3-Nano",
dtype=torch.bfloat16,
device_map="cuda", # or "mps", "xpu", "cpu"
safety_checker=MyCustomSafetyChecker(),
)Cosmos3OmniPipeline
classe 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 )
Décode un latent sonore [C, T] en une forme d'onde [audio_ch, N].
Ajoute/supprime la dimension de lot attendue par le décodeur du tokenizer audio.
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 )
Construire le conditionnement et le bruit initial pour un seul échantillon.
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 )
Appliquez les modèles d'augmentation de prompt et tokenisez les prompts cond/uncond via le modèle de chat configuré.
Ce pipeline n'exécute pas d'encodeur de texte séparé : le transformeur joint Cosmos3 consomme les identifiants de jetons bruts en même temps que les jetons visuels (et, le cas échéant, sonores).
Lorsque negative_prompt est None, une chaîne vide est utilisée ; la page de documentation de Cosmos3 documente les prompts négatifs de contrôle qualité recommandés à passer explicitement pour text2video / image2video. Les modèles de durée et de résolution sont ajoutés au prompt, et les modèles inverses sont ajoutés au prompt négatif, lorsque cette option est activée.
Lorsque action_mode est défini, le prompt est plutôt converti en la légende JSON d’action structurée sur laquelle le modèle a été entraîné (voir _build_action_json_prompt), en utilisant action_view_point pour le champ de cadrage ; les modèles de métadonnées plates sont ignorés car le JSON contient déjà duration/fps/resolution/aspect_ratio.
- tous
- appel
Cosmos3OmniModularPipeline
Cosmos 3 est également disponible en tant que pipeline Modular Diffusers. Le pipeline basé sur les tâches Cosmos3OmniPipeline reste disponible ; le pipeline modulaire coexiste avec lui et couvre les mêmes modes (text2image, text2video, image2video, video2video, génération conditionnée par l'action et transfer (contrôle structurel), avec un son optionnel lorsque le checkpoint le prend en charge).
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]Vous pouvez également charger via ModularPipeline et laisser la configuration du dépôt sélectionner la classe de blocs :
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"
)Pour inspecter ou personnaliser un flux de travail modulaire Cosmos spécifique, utilisez available_workflows + get_workflow() :
available = pipe.blocks.available_workflows
image2video_blocks = pipe.blocks.get_workflow("image2video")Exemples modulaires pour tous les flux de travail existants
Le pipeline modulaire prend en charge les mêmes signatures d’appel que le pipeline de tâche. Les extraits ci-dessous reproduisent chaque exemple de génération présenté plus haut (text2video, text2image, image2video, video2video, video2video_sound, text2video_sound et action_policy). Le transfert (contrôle structurel) possède ses propres entrées et est présenté séparément dans Transfert modulaire ci-dessous.
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)Transfert modulaire (contrôle structurel)
Le transfert suit une vidéo de contrôle précalculée (contour, flou, profondeur, segmentation ou carte de scénario de monde) transmise via control_videos= sous forme de mapping {hint: video}. Il est réservé aux vidéos (pas de image / video / action / enable_sound), le prompt est une légende JSON pré-échantillonnée (voir Échantillonnage du prompt), et les clips longs sont générés de manière autorégressive par blocs de num_video_frames_per_chunk puis assemblés automatiquement. guidance_scale est le CFG texte habituel ; control_guidance (!= 1.0) amplifie en outre le signal de contrôle. Valeurs de départ recommandées par indice :
| Astuce | guidance_scale | control_guidance | flow_shift | Géométrie |
|---|---|---|---|---|
| Bord / Flou / Profondeur | 3.0 | 1.5 | 10.0 | 121 images @ 30 FPS |
| Segmentation | 3.0 | 2.0 | 10.0 | 121 images @ 30 FPS |
| Scénario mondial (WSM) | 1.0 | 3.0 | 10.0 | 101 images @ 10 i/s |
Diffusers ne fournit pas les assets de contrôle. Des assets prêts à l’emploi (une vidéo de contrôle + un prompt.json correspondant par indice, ainsi qu’un negative_prompt.json partagé) sont disponibles dans le Cosmos cookbook. Pour l’exemple de bordure ci-dessous, téléchargez-les dans un dossier local 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.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)Texte vers image et image vers vidéo distillés (peu d'étapes)
Les points de contrôle distillés en peu d’étapes sont servis par Cosmos3DistilledModularPipeline (blocs : Cosmos3DistilledBlocks) ; les pipelines de base Cosmos3OmniModularPipeline et Cosmos3OmniPipeline ne les prennent pas en charge. num_inference_steps est fixé à la longueur de la configuration distilled_sigmas du pipeline (issue du modular_model_index.json du point de contrôle) et guidance_scale est forcé à 1.0, car le guidage est intégré dans les poids — passer toute autre valeur pour l’un ou l’autre provoque une erreur, et negative_prompt est signalé par un avertissement puis ignoré.
Les points de contrôle distillés en FP8 (revision="fp8") ne déclarent pas de politique de précision mixte, chaque étape reste donc en W8A8 natif.
Les prompts suivent la même structure JSON descriptive que les modèles non distillés, il faut donc d’abord augmenter la résolution du texte court — utilisez --mode text2image (T2I) ou --mode image2video (I2V) comme décrit dans Augmentation de la résolution des prompts, puis transmettez le JSON via 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)classe 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 )
Un ModularPipeline pour la génération omni de Cosmos 3.
- tous
- appel
Cosmos3DistilledModularPipeline
classe 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 )
Un ModularPipeline pour la génération omni de Cosmos 3 distillée (peu d'étapes).
Les points de contrôle distillés intègrent le guidage sans classifieur dans les poids et échantillonnent selon un plan fixe lu depuis la configuration distilled_sigmas du pipeline (remplie à partir de modular_model_index.json), de sorte que guidance_scale et num_inference_steps sont fixes et que negative_prompt n’est pas pris en charge.
- tous
- appel
CosmosActionCondition
classe 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' )
Paramètres
- mode (
str) — La tâche d’action. L’une des valeurs"forward_dynamics"(générer la vidéo future à partir d’une première image et d’une séquence deraw_actionsdonnée),"inverse_dynamics"(inférer les actions reliant les images de conditionnement) ou"policy"(générer conjointement la vidéo future et les actions à partir de la première image). - chunk_size (
int) — Nombre d'étapes de transition d'action dans le chunk. La vidéo de conditionnement associée s'étend surchunk_size + 1images. - domain_name (
str) — Domaine d’incarnation sélectionnant les poids de projection d’action sensibles au domaine. Doit être l’un des domaines d’incarnation Cosmos 3 enregistrés. Il fixe également la largeur d’action non rembourrée utilisée pour extraire les actions prédites, résolue en interne à partir de ce nom (voir_EMBODIMENT_TO_RAW_ACTION_DIM). - resolution_tier (
int, par défaut480) — Niveau de résolution tier pour le conditionnement de l’action (l’un des256,480,704,720). Le niveau sélectionne un canevas prédéfini dont le ratio d’aspect est le plus proche de celui de l’entrée ; l’entrée est réduite (jamais agrandie) et mise en bourre dans ce canevas pour le conditionnement. Il ne s’agit pas de la taille des images de sortie, qui suit le contenu de l’entrée. Faites correspondre le niveau à la résolution native de l’entrée : un niveau inférieur supprime des détails, tandis qu’un niveau supérieur n’ajoute aucune résolution (pas d’agrandissement), gaspille des calculs sur la bourre et constitue un désaccord entraînement/inférence susceptible de nuire à la qualité. - raw_actions (
torch.Tensor, facultatif) — Vecteurs d’action bruts du domaine de forme[T, raw_action_dim]pilotant"forward_dynamics". Les séquences plus courtes quechunk_sizerépètent la dernière action ; les plus longues sont tronquées. Les canaux au-delà deaction_dimdu modèle sont rejetés, et les entrées plus étroites sont complétées par des zéros jusqu’àaction_dim. - image (
PIL.Image.Image,np.ndarrayoutorch.Tensor, facultatif) — Image de conditionnement pour"policy"/"forward_dynamics". Mutuellement exclusif avecvideo. - video (
list,np.ndarray, ortorch.Tensor, optional) — Conditioning video, required for"inverse_dynamics". For"policy"/"forward_dynamics"only its first frame is used. Mutually exclusive withimage. - view_point (
str, par défaut"ego_view") — Étiquette de perspective de caméra utilisée pour remplir le champcinematography.framingde la légende d’action. L’une des valeurs"ego_view","third_person_view","wrist_view"ou"concat_view". Le modèle d’action a été entraîné sur des légendes JSON structurées qui contiennent cette phrase de point de vue ; une étiquette non reconnue supprime le champ framing (avec un avertissement).
Regroupe toutes les entrées requises pour une tâche de génération conditionnée par une action dans Cosmos 3.
Passez ceci à Cosmos3OmniPipeline.__call__() via l'argument action plutôt que via les arguments de niveau supérieur image / height / width, qui sont réservés aux exécutions t2v et i2v.
Cosmos3OmniPipelineOutput
classe diffusers.pipelines.cosmos.pipeline_cosmos3_omni.Cosmos3OmniPipelineOutput
< source >( video: Anysound: typing.Optional[torch.Tensor] = Noneaction: list[torch.Tensor] | None = None )
Paramètres
- video — La vidéo générée. Le type exact dépend de
output_typepassé au pipeline : une liste de frames PIL pour"pil"(par défaut), unnp.ndarrayde forme[T, H, W, C]pour"np", untorch.Tensorde forme[T, C, H, W]pour"pt", ou un tenseur latent brut lorsqueoutput_type="latent". - sound — Forme d'onde audio décodée de forme
[C, N].Nonelorsqueenable_sound=False. - action — Jetons d'action prédits.
Nonesauf si un mode d'action prédit des actions.
Classe de données de sortie pour Cosmos3OmniPipeline.