#!/usr/bin/env python3 # SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """ Convert a Cosmos3 DCP checkpoint to diffusers format. Example: CUDA_VISIBLE_DEVICES=0 python scripts/convert_cosmos3_to_diffusers.py \ --checkpoint-path Cosmos3-Nano \ --output converted/cosmos3-nano-pipeline \ --save-pipeline A raw Cosmos3 Edge DCP checkpoint is detected automatically. The Edge path loads its EMA denoiser weights directly from DCP, packages the matching Wan 2.2 VAE, and writes a shared-weight Omni checkpoint: the Diffusers transformer is reused by the reasoner and only the missing vision tower is stored separately. Edge checkpoints trained with action generation are also detected automatically: the `action2llm`/`llm2action` domain-aware projections and the action modality embedding are exported onto the transformer's `action_proj_in`/`action_proj_out`/`action_modality_embed` modules, with `action_dim` and `num_embodiment_domains` derived from the DCP tensor shapes. """ import argparse import contextlib import json import pathlib import re import shutil import torch from diffusers.models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer DEFAULT_SOUND_TOKENIZER_CONFIG = { "sampling_rate": 48000, "vocoder_input_dim": 64, "dec_dim": 320, "dec_c_mults": [1, 2, 4, 8, 16], "dec_strides": [2, 4, 5, 6, 8], "dec_out_channels": 2, } COSMOS3_EDGE_REASONER = "nvidia/Cosmos3-Edge-Reasoner" COSMOS3_EDGE_REASONER_REVISION = "590c1c0f1cd7146162d478a2180556055c1a252b" COSMOS3_EDGE_VAE = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" # Keep the Transformers metadata, but avoid copying its two source shards: # Edge's language tower is already byte-identical to the base part of the # converted Diffusers transformer. Only the visual tower and projector have # to be extracted from source weights. COSMOS3_EDGE_REASONER_METADATA_FILES = ( "chat_template.jinja", "config.json", "generation_config.json", "preprocessor_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json", "video_preprocessor_config.json", ) COSMOS3_EDGE_REASONER_INDEX_FILE = "model.safetensors.index.json" COSMOS3_EDGE_REASONER_VISION_SHARD = "00000.safetensors" COSMOS3_EDGE_REASONER_FILES = ( *COSMOS3_EDGE_REASONER_METADATA_FILES, COSMOS3_EDGE_REASONER_INDEX_FILE, COSMOS3_EDGE_REASONER_VISION_SHARD, ) _EDGE_ATTN_KEY_REMAP = ( (".q_proj_moe_gen.", ".add_q_proj."), (".k_proj_moe_gen.", ".add_k_proj."), (".v_proj_moe_gen.", ".add_v_proj."), (".o_proj_moe_gen.", ".to_add_out."), (".q_norm_moe_gen.", ".norm_added_q."), (".k_norm_moe_gen.", ".norm_added_k."), (".q_proj.", ".to_q."), (".k_proj.", ".to_k."), (".v_proj.", ".to_v."), (".o_proj.", ".to_out."), ) _TIME_EMBEDDER_REMAP = { "mlp.0.weight": "linear_1.weight", "mlp.0.bias": "linear_1.bias", "mlp.2.weight": "linear_2.weight", "mlp.2.bias": "linear_2.bias", } def _get_config_value(*configs, name, default=None): for config in configs: if config is None: continue if hasattr(config, name): value = getattr(config, name) if value is not None: return value if isinstance(config, dict) and config.get(name) is not None: return config[name] return default def _find_edge_dcp_model_dir(checkpoint_path: pathlib.Path) -> pathlib.Path | None: candidates = [checkpoint_path, checkpoint_path / "model"] if checkpoint_path.is_dir(): candidates.extend(metadata_path.parent for metadata_path in checkpoint_path.glob("iter_*/model/.metadata")) candidates.extend(metadata_path.parent for metadata_path in checkpoint_path.glob("*/iter_*/model/.metadata")) for candidate in candidates: if not (candidate / ".metadata").is_file(): continue from torch.distributed.checkpoint.filesystem import FileSystemReader state_dict_metadata = FileSystemReader(str(candidate)).read_metadata().state_dict_metadata has_edge_mlp = "net.language_model.model.layers.0.mlp.up_proj.weight" in state_dict_metadata has_qwen_mlp = "net.language_model.model.layers.0.mlp.gate_proj.weight" in state_dict_metadata if has_edge_mlp and not has_qwen_mlp: return candidate return None def _remap_edge_dcp_key(key: str) -> str: if key.startswith("language_model."): key = key.removeprefix("language_model.").removeprefix("model.") for old, new in _EDGE_ATTN_KEY_REMAP: if old in key: return key.replace(old, new) return key if key.startswith("vae2llm."): return f"proj_in.{key.removeprefix('vae2llm.')}" if key.startswith("llm2vae."): return f"proj_out.{key.removeprefix('llm2vae.')}" if key.startswith("time_embedder."): time_embedder_key = key.removeprefix("time_embedder.") return f"time_embedder.{_TIME_EMBEDDER_REMAP[time_embedder_key]}" if key.startswith("action2llm."): return f"action_proj_in.{key.removeprefix('action2llm.')}" if key.startswith("llm2action."): return f"action_proj_out.{key.removeprefix('llm2action.')}" if key == "action_modality_embed": return key raise ValueError(f"Unsupported Cosmos3 Edge DCP key: {key!r}.") def _detect_edge_action_config(checkpoint_path: pathlib.Path, use_ema: bool) -> dict | None: from torch.distributed.checkpoint.filesystem import FileSystemReader prefix = "net_ema." if use_ema else "net." metadata = FileSystemReader(str(checkpoint_path)).read_metadata().state_dict_metadata action_suffixes = ( "action2llm.fc.weight", "action2llm.bias.weight", "llm2action.fc.weight", "llm2action.bias.weight", "action_modality_embed", ) action_metadata = {suffix: metadata.get(f"{prefix}{suffix}") for suffix in action_suffixes} present_suffixes = [suffix for suffix, value in action_metadata.items() if value is not None] if not present_suffixes: return None # `action2llm`/`llm2action` are DomainAwareLinear modules: `fc` is an # Embedding(num_domains, output_size * input_size) and `bias` an # Embedding(num_domains, output_size), so the shapes fix every config value. missing_suffixes = [suffix for suffix, value in action_metadata.items() if value is None] if missing_suffixes: raise ValueError( f"Cosmos3 Edge {'EMA' if use_ema else 'regular'} action weights are incomplete: " f"present={present_suffixes}, missing={missing_suffixes}." ) action_shapes = {suffix: tuple(value.size) for suffix, value in action_metadata.items()} expected_ranks = { "action2llm.fc.weight": 2, "action2llm.bias.weight": 2, "llm2action.fc.weight": 2, "llm2action.bias.weight": 2, "action_modality_embed": 1, } invalid_shapes = { suffix: shape for suffix, shape in action_shapes.items() if len(shape) != expected_ranks[suffix] } if invalid_shapes: raise ValueError(f"Cosmos3 Edge action tensors have invalid ranks: {invalid_shapes}.") num_embodiment_domains, proj_in_flat = action_shapes["action2llm.fc.weight"] proj_in_domains, hidden_size = action_shapes["action2llm.bias.weight"] proj_out_domains, proj_out_flat = action_shapes["llm2action.fc.weight"] proj_out_bias_domains, action_dim = action_shapes["llm2action.bias.weight"] (action_modality_embed_size,) = action_shapes["action_modality_embed"] domain_counts = { "action2llm.fc.weight": num_embodiment_domains, "action2llm.bias.weight": proj_in_domains, "llm2action.fc.weight": proj_out_domains, "llm2action.bias.weight": proj_out_bias_domains, } if len(set(domain_counts.values())) != 1: raise ValueError(f"Cosmos3 Edge action tensors disagree on embodiment domains: {domain_counts}.") if min(num_embodiment_domains, hidden_size, action_dim) < 1: raise ValueError(f"Cosmos3 Edge action tensor dimensions must be positive: {action_shapes}.") expected_projection_size = hidden_size * action_dim if proj_in_flat != expected_projection_size or proj_out_flat != expected_projection_size: raise ValueError( "Cosmos3 Edge action projection shapes are inconsistent: " f"action2llm.fc.weight={action_shapes['action2llm.fc.weight']}, " f"llm2action.fc.weight={action_shapes['llm2action.fc.weight']}; both must factor into " f"hidden_size={hidden_size} x action_dim={action_dim}." ) if action_modality_embed_size != hidden_size: raise ValueError( "Cosmos3 Edge action modality embedding has the wrong width: " f"action_modality_embed={action_shapes['action_modality_embed']}, expected ({hidden_size},)." ) return { "action_gen": True, "action_dim": action_dim, "num_embodiment_domains": num_embodiment_domains, } def _build_edge_transformer(dtype: torch.dtype, action_config: dict | None = None): import inspect from accelerate import init_empty_weights from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer action_kwargs = dict(action_config or {}) init_params = inspect.signature(Cosmos3OmniTransformer.__init__).parameters edge_required_params = { "hidden_act", "qk_norm_for_text", "backbone_type", "temporal_compression_factor", } action_required_params = set(action_kwargs) missing_edge_params = sorted(edge_required_params - set(init_params)) missing_action_params = sorted(action_required_params - set(init_params)) if missing_edge_params or missing_action_params: missing_descriptions = [] if missing_edge_params: missing_descriptions.append(f"Edge backbone parameters {missing_edge_params}") if missing_action_params: missing_descriptions.append(f"action-generation parameters {missing_action_params}") raise RuntimeError( "The installed diffusers build cannot construct this Cosmos3 Edge checkpoint; it is missing " f"{', '.join(missing_descriptions)}. Use a build that combines the Edge Nemotron-dense backbone " "with upstream action generation support (commit bcc20e452)." ) with init_empty_weights(): transformer = Cosmos3OmniTransformer( **action_kwargs, attention_bias=False, attention_dropout=0.0, base_fps=24, enable_fps_modulation=True, head_dim=128, hidden_act="relu2", hidden_size=2048, intermediate_size=9216, latent_channel=48, latent_patch_size=2, num_attention_heads=16, num_hidden_layers=28, num_key_value_heads=8, patch_latent_dim=192, qk_norm_for_text=False, rms_norm_eps=1e-5, rope_scaling={"mrope_section": [24, 20, 20]}, rope_theta=100_000_000.0, timestep_scale=0.001, unified_3d_mrope_reset_spatial_ids=True, unified_3d_mrope_temporal_modality_margin=15000, vocab_size=131072, backbone_type="cosmos3_edge_nemotron_dense", temporal_compression_factor=4, ) transformer = transformer.to(dtype=dtype) transformer.time_embedder.to(dtype=torch.float32) return transformer.to_empty(device="cpu") def _validate_edge_action_pipeline_support() -> None: import inspect from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline pipeline_params = inspect.signature(Cosmos3OmniPipeline.__call__).parameters required_params = {"action", "action_latents"} missing_params = sorted(required_params - set(pipeline_params)) if missing_params: raise RuntimeError( "The checkpoint has action generation weights and --save-pipeline was requested, but this diffusers " f"build's Cosmos3OmniPipeline does not accept {missing_params}. Use a build that also includes the " "Cosmos3 action pipeline support from commit bcc20e452." ) def _load_edge_dcp_weights(transformer, checkpoint_path: pathlib.Path, use_ema: bool) -> None: import torch.distributed.checkpoint as dcp from torch.distributed.checkpoint.filesystem import FileSystemReader prefix = "net_ema." if use_ema else "net." state_dict_metadata = FileSystemReader(str(checkpoint_path)).read_metadata().state_dict_metadata target_state_dict = transformer.state_dict() dcp_state_dict: dict[str, torch.Tensor] = {} for key in state_dict_metadata: if not key.startswith(prefix): continue target_key = _remap_edge_dcp_key(key.removeprefix(prefix)) if target_key not in target_state_dict: raise RuntimeError(f"Cosmos3 Edge DCP key {key!r} maps to unknown transformer key {target_key!r}.") dcp_state_dict[key] = target_state_dict[target_key] mapped_keys = {_remap_edge_dcp_key(key.removeprefix(prefix)) for key in dcp_state_dict} missing_keys = sorted(set(target_state_dict) - mapped_keys) if missing_keys: raise RuntimeError(f"Cosmos3 Edge DCP is missing transformer weights: {missing_keys}.") print(f"Loading {'EMA' if use_ema else 'regular'} Cosmos3 Edge weights from {checkpoint_path} …") dcp.load(state_dict=dcp_state_dict, storage_reader=FileSystemReader(str(checkpoint_path)), no_dist=True) del dcp_state_dict, target_state_dict def _resolve_edge_reasoner_path(args) -> pathlib.Path: if args.reasoner_path is not None: reasoner_path = pathlib.Path(args.reasoner_path).expanduser().absolute() else: from huggingface_hub import snapshot_download print( "Downloading the pinned Cosmos3 Edge reasoner snapshot " f"({args.reasoner_repo_id}@{args.reasoner_revision}) …" ) reasoner_path = pathlib.Path( snapshot_download( repo_id=args.reasoner_repo_id, revision=args.reasoner_revision, allow_patterns=list(COSMOS3_EDGE_REASONER_FILES), ) ) if not reasoner_path.is_dir(): raise FileNotFoundError(f"Cosmos3 Edge reasoner directory not found: {reasoner_path}") missing_files = [filename for filename in COSMOS3_EDGE_REASONER_FILES if not (reasoner_path / filename).is_file()] if missing_files: raise FileNotFoundError( f"Cosmos3 Edge reasoner at {reasoner_path} is missing required files: {missing_files}" ) return reasoner_path def _load_json(path: pathlib.Path) -> dict: with open(path, encoding="utf-8") as f: return json.load(f) def _save_json(payload: dict, path: pathlib.Path) -> None: path.parent.mkdir(parents=True, exist_ok=True) with open(path, "w", encoding="utf-8") as f: json.dump(payload, f, indent=2, sort_keys=True) f.write("\n") def _edge_reasoner_language_to_generator_key_map() -> dict[str, str]: """Map source Nemotron reasoner keys to their shared Diffusers tensor keys.""" mapping = { "lm_head.weight": "lm_head.weight", "model.language_model.embeddings.weight": "embed_tokens.weight", "model.language_model.norm_f.weight": "norm.weight", } attention_key_map = { "q_proj": "to_q", "k_proj": "to_k", "v_proj": "to_v", "o_proj": "to_out", } for layer_idx in range(28): attention_layer_idx = 2 * layer_idx mlp_layer_idx = attention_layer_idx + 1 reasoner_attention_prefix = f"model.language_model.layers.{attention_layer_idx}" generator_attention_prefix = f"layers.{layer_idx}" mapping[f"{reasoner_attention_prefix}.norm.weight"] = f"{generator_attention_prefix}.input_layernorm.weight" for source_name, generator_name in attention_key_map.items(): mapping[f"{reasoner_attention_prefix}.mixer.{source_name}.weight"] = ( f"{generator_attention_prefix}.self_attn.{generator_name}.weight" ) reasoner_mlp_prefix = f"model.language_model.layers.{mlp_layer_idx}" generator_mlp_prefix = f"layers.{layer_idx}" mapping[f"{reasoner_mlp_prefix}.norm.weight"] = f"{generator_mlp_prefix}.post_attention_layernorm.weight" for source_name in ("up_proj", "down_proj"): mapping[f"{reasoner_mlp_prefix}.mixer.{source_name}.weight"] = ( f"{generator_mlp_prefix}.mlp.{source_name}.weight" ) return mapping def _native_edge_text_config(source_config: dict) -> dict: """Strip inactive hybrid fields and express the dense reasoner layout explicitly.""" layers_block_type = source_config.get("layers_block_type") if layers_block_type is None: legacy_pattern = source_config.get("hybrid_override_pattern") if legacy_pattern is None: num_hidden_layers = source_config.get("num_hidden_layers", 56) if num_hidden_layers % 2: raise ValueError("Cosmos3 Edge requires an even number of alternating attention/MLP layers.") layers_block_type = ["full_attention", "mlp"] * (num_hidden_layers // 2) else: layer_type_map = {"*": "full_attention", "-": "mlp"} try: layers_block_type = [layer_type_map[layer] for layer in legacy_pattern] except KeyError as exc: raise ValueError(f"Unsupported Cosmos3 Edge legacy layer type: {exc.args[0]!r}.") from exc else: layers_block_type = list(layers_block_type) if set(layers_block_type) - {"full_attention", "mlp"}: raise ValueError(f"Cosmos3 Edge only supports dense attention/MLP layers, got {layers_block_type!r}.") native_config = {"model_type": "cosmos3_edge_text", "layers_block_type": layers_block_type} for key in ( "attention_bias", "attention_dropout", "bos_token_id", "dtype", "eos_token_id", "head_dim", "hidden_size", "initializer_range", "intermediate_size", "layer_norm_epsilon", "max_position_embeddings", "mlp_bias", "mlp_hidden_act", "num_attention_heads", "num_key_value_heads", "num_logits_to_keep", "pad_token_id", "rope_theta", "use_cache", "vocab_size", ): if key in source_config: native_config[key] = source_config[key] native_config["num_hidden_layers"] = len(layers_block_type) native_config["rope_parameters"] = { "rope_type": "default", "rope_theta": native_config.get("rope_theta", 100_000_000.0), "mrope_section": source_config.get("mrope_section", [24, 20, 20]), } return native_config def _native_edge_projector_config(source_config: dict) -> dict: merger_intermediate_size = source_config.get("merger_intermediate_size") if merger_intermediate_size is None: merger_intermediate_size = source_config["merger_intermedia"] native_config = { "model_type": "cosmos3_edge_projector", "input_hidden_size": source_config["input_hidden_size"], "merger_intermediate_size": merger_intermediate_size, "out_hidden_size": source_config["out_hidden_size"], "spatial_merge_size": source_config["spatial_merge_size"], "use_postshuffle_norm": source_config["use_postshuffle_norm"], } return native_config def _native_edge_vision_config(source_config: dict) -> dict: native_config = {"model_type": "cosmos3_edge_vision"} for key in ( "attention_dropout", "hidden_act", "hidden_size", "intermediate_size", "layer_norm_eps", "num_attention_heads", "num_channels", "num_hidden_layers", "num_patches", "patch_size", "spatial_merge_size", ): if key in source_config: native_config[key] = source_config[key] return native_config def _native_edge_config(source_config: dict) -> dict: native_config = { key: value for key, value in source_config.items() if key not in {"architectures", "auto_map", "model_type", "projector_config", "text_config", "vision_config"} } native_config["architectures"] = ["Cosmos3EdgeForConditionalGeneration"] native_config["model_type"] = "cosmos3_edge" native_config["allow_patterns_overrides"] = ["*/*.safetensors"] native_config["text_config"] = _native_edge_text_config(source_config["text_config"]) native_config["vision_config"] = _native_edge_vision_config(source_config["vision_config"]) native_config["projector_config"] = _native_edge_projector_config(source_config["projector_config"]) return native_config def _native_edge_image_processor_config(source_config: dict) -> dict: native_config = dict(source_config) native_config.pop("auto_map", None) native_config["processor_class"] = "Cosmos3EdgeProcessor" native_config["image_processor_type"] = "Cosmos3EdgeImageProcessor" return native_config def _native_edge_video_processor_config(source_config: dict) -> dict: native_config = dict(source_config) native_config.pop("auto_map", None) native_config["processor_class"] = "Cosmos3EdgeProcessor" native_config["video_processor_type"] = "Cosmos3EdgeVideoProcessor" return native_config def _validate_edge_reasoner_mapping(source_weight_map: dict[str, str], generator_weight_map: dict[str, str]) -> dict[str, str]: language_keys = { key for key in source_weight_map if key == "lm_head.weight" or key.startswith("model.language_model.") } vision_keys = { key for key in source_weight_map if key.startswith(("model.visual.", "model.projector.")) } unexpected_keys = set(source_weight_map) - language_keys - vision_keys if unexpected_keys: raise RuntimeError(f"Unexpected Cosmos3 Edge reasoner keys: {sorted(unexpected_keys)}") mapping = _edge_reasoner_language_to_generator_key_map() if set(mapping) != language_keys: missing = sorted(language_keys - set(mapping)) extra = sorted(set(mapping) - language_keys) raise RuntimeError(f"Incomplete Edge language mapping: missing={missing}, extra={extra}") missing_generator_keys = sorted(set(mapping.values()) - set(generator_weight_map)) if missing_generator_keys: raise RuntimeError(f"Diffusers transformer is missing shared reasoner tensors: {missing_generator_keys}") if len(vision_keys) != 443: raise RuntimeError(f"Expected 443 Edge vision/projector tensors, found {len(vision_keys)}.") return mapping def _write_edge_vision_encoder( reasoner_path: pathlib.Path, output_dir: pathlib.Path, source_weight_map: dict[str, str], ) -> set[str]: try: from safetensors import safe_open from safetensors.torch import save_file except ImportError as exc: raise ImportError("Saving the Cosmos3 Edge vision encoder requires safetensors.") from exc vision_keys = { key for key in source_weight_map if key.startswith(("model.visual.", "model.projector.")) } non_vision_shards = {source_weight_map[key] for key in vision_keys} - {COSMOS3_EDGE_REASONER_VISION_SHARD} if non_vision_shards: raise RuntimeError(f"Edge visual tensors unexpectedly span shards: {sorted(non_vision_shards)}") vision_dir = output_dir / "vision_encoder" vision_dir.mkdir(parents=True, exist_ok=True) vision_path = vision_dir / "model.safetensors" print(f"Extracting {len(vision_keys)} Cosmos3 Edge vision/projector tensors to {vision_path} …") with safe_open(reasoner_path / COSMOS3_EDGE_REASONER_VISION_SHARD, framework="pt", device="cpu") as source_file: state_dict = {key: source_file.get_tensor(key).contiguous() for key in sorted(vision_keys)} save_file(state_dict, str(vision_path), metadata={"format": "pt"}) del state_dict # This is an auxiliary shard consumed through the root reasoner manifest; # it is not a standalone Transformers model. Do not leave a stale legacy # config beside it when converting into an existing output directory. (vision_dir / "config.json").unlink(missing_ok=True) return vision_keys def _write_edge_reasoner_manifest(reasoner_path: pathlib.Path, output_dir: pathlib.Path) -> None: source_index = _load_json(reasoner_path / COSMOS3_EDGE_REASONER_INDEX_FILE) source_weight_map = source_index["weight_map"] source_config = _load_json(reasoner_path / "config.json") generator_index = _load_json(output_dir / "transformer" / "diffusion_pytorch_model.safetensors.index.json") generator_weight_map = generator_index["weight_map"] language_mapping = _validate_edge_reasoner_mapping(source_weight_map, generator_weight_map) vision_keys = _write_edge_vision_encoder(reasoner_path, output_dir, source_weight_map) weight_map = { generator_key: f"transformer/{generator_weight_map[generator_key]}" for generator_key in language_mapping.values() } weight_map.update({key: "vision_encoder/model.safetensors" for key in vision_keys}) if len(weight_map) != len(source_weight_map): raise RuntimeError( f"Edge unified weight index has {len(weight_map)} entries, expected {len(source_weight_map)}." ) _save_json( {"metadata": source_index["metadata"], "weight_map": weight_map}, output_dir / COSMOS3_EDGE_REASONER_INDEX_FILE, ) _save_json(_native_edge_config(source_config), output_dir / "config.json") def _copy_edge_reasoner_metadata(reasoner_path: pathlib.Path, output_dir: pathlib.Path) -> None: print(f"Writing the shared Cosmos3 Edge reasoner into {output_dir} …") for filename in COSMOS3_EDGE_REASONER_METADATA_FILES: if filename == "config.json": continue if filename == "preprocessor_config.json": _save_json( _native_edge_image_processor_config(_load_json(reasoner_path / filename)), output_dir / filename, ) continue if filename == "video_preprocessor_config.json": _save_json( _native_edge_video_processor_config(_load_json(reasoner_path / filename)), output_dir / filename, ) continue shutil.copy2(reasoner_path / filename, output_dir / filename) for filename in ( "configuration_nemotron_siglip2_h.py", "modeling_cosmos3_edge_omni.py", "modeling_nemotron_siglip2_h.py", "processing.py", ): (output_dir / filename).unlink(missing_ok=True) _write_edge_reasoner_manifest(reasoner_path, output_dir) for filename in ("00000.safetensors", "00001.safetensors"): (output_dir / filename).unlink(missing_ok=True) def _copy_edge_conversion_script(output_dir: pathlib.Path) -> None: source = pathlib.Path(__file__).resolve() destination = output_dir / source.name if source != destination: shutil.copy2(source, destination) def _add_edge_reasoner_to_pipeline(args) -> None: output_dir = pathlib.Path(args.output).expanduser().absolute() expected_paths = ("model_index.json", "scheduler", "text_tokenizer", "transformer", "vae") missing_paths = [str(output_dir / path) for path in expected_paths if not (output_dir / path).exists()] if missing_paths: raise FileNotFoundError( "Expected an existing Cosmos3 Edge Diffusers pipeline before adding its reasoner; " f"missing paths: {missing_paths}" ) reasoner_path = _resolve_edge_reasoner_path(args) _copy_edge_reasoner_metadata(reasoner_path, output_dir) _copy_edge_conversion_script(output_dir) print("Done.") def _convert_edge_dcp(args, checkpoint_path: pathlib.Path, dtype: torch.dtype) -> None: from transformers import PreTrainedTokenizerFast from diffusers import AutoencoderKLWan, UniPCMultistepScheduler from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline if args.include_sound_tokenizer or args.sound_tokenizer_path is not None: raise ValueError("The supplied Cosmos3 Edge checkpoint is video-only and cannot include a sound tokenizer.") if args.include_reasoner and not args.save_pipeline: raise ValueError( "A Cosmos3 Edge reasoner can only be included with --save-pipeline because its root Transformers " "config.json would conflict with a transformer-only Diffusers save. Use --no-include-reasoner instead." ) action_config = _detect_edge_action_config(checkpoint_path, args.use_ema) if action_config is not None: print( "Detected Cosmos3 Edge action generation weights " f"(action_dim={action_config['action_dim']}, " f"num_embodiment_domains={action_config['num_embodiment_domains']})." ) transformer = _build_edge_transformer(dtype, action_config) if action_config is not None and args.save_pipeline: _validate_edge_action_pipeline_support() _load_edge_dcp_weights(transformer, checkpoint_path, args.use_ema) reasoner_path = _resolve_edge_reasoner_path(args) if args.include_reasoner else None output_dir = pathlib.Path(args.output).expanduser().absolute() output_dir.mkdir(parents=True, exist_ok=True) if not args.save_pipeline: print(f"Saving Cosmos3 Edge transformer to {output_dir} …") transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") print("Done.") return tokenizer_source = str(reasoner_path) if reasoner_path is not None else args.reasoner_repo_id tokenizer_kwargs = {} if reasoner_path is None: tokenizer_kwargs["revision"] = args.reasoner_revision text_tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_source, **tokenizer_kwargs) for token in ("<|vision_start|>", "<|vision_end|>"): token_id = text_tokenizer.convert_tokens_to_ids(token) if token_id is None or token_id < 0 or token_id >= transformer.config.vocab_size: raise ValueError( f"Cosmos3 Edge tokenizer token {token!r} has invalid ID {token_id!r} for " f"vocab_size={transformer.config.vocab_size}." ) diffusers_vae = AutoencoderKLWan.from_pretrained(COSMOS3_EDGE_VAE, subfolder="vae", torch_dtype=torch.bfloat16) scheduler = UniPCMultistepScheduler( num_train_timesteps=1000, solver_order=2, predict_x0=True, solver_type="bh2", lower_order_final=True, final_sigmas_type="zero", use_flow_sigmas=True, prediction_type="flow_prediction", flow_shift=5.0, ) pipeline = Cosmos3OmniPipeline( transformer=transformer, text_tokenizer=text_tokenizer, vae=diffusers_vae, scheduler=scheduler, enable_safety_checker=False, default_use_system_prompt=False, use_native_flow_schedule=True, ) print(f"Saving Cosmos3 Edge pipeline to {output_dir} …") pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") if reasoner_path is not None: _copy_edge_reasoner_metadata(reasoner_path, output_dir) _copy_edge_conversion_script(output_dir) print("Done.") def _load_sound_tokenizer_state_dict(checkpoint_path: pathlib.Path) -> dict[str, torch.Tensor]: if checkpoint_path.suffix == ".safetensors": try: from safetensors.torch import load_file except ImportError as exc: raise ImportError("Loading AVAE .safetensors checkpoints requires safetensors.") from exc checkpoint = load_file(str(checkpoint_path), device="cpu") else: checkpoint = torch.load(checkpoint_path, map_location="cpu") if not isinstance(checkpoint, dict): raise TypeError(f"AVAE checkpoint must be a dict, got {type(checkpoint)!r}.") for key in ("generator", "state_dict", "model"): value = checkpoint.get(key) if isinstance(value, dict): checkpoint = value break state_dict = { key: value.detach().cpu().contiguous() for key, value in checkpoint.items() if isinstance(value, torch.Tensor) } if not state_dict: raise RuntimeError(f"No tensor state dict found in AVAE checkpoint keys: {list(checkpoint.keys())[:16]}") return state_dict def _load_sound_tokenizer_config(config_path: pathlib.Path | None, fallback_config_path: pathlib.Path) -> dict: selected_config_path = config_path if selected_config_path is None and fallback_config_path.exists(): selected_config_path = fallback_config_path if selected_config_path is None: return dict(DEFAULT_SOUND_TOKENIZER_CONFIG) with open(selected_config_path, encoding="utf-8") as f: return json.load(f) _SOUND_TOKENIZER_PER_KEY_PREFIXES = ("module.", "generator.", "model.", "state_dict.") _SOUND_TOKENIZER_RES_UNIT_INNER_NAMES = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"} def _sound_tokenizer_strip_per_key_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: out = dict(state_dict) changed = True while changed: changed = False for prefix in _SOUND_TOKENIZER_PER_KEY_PREFIXES: if any(key.startswith(prefix) for key in out): out = {(key[len(prefix) :] if key.startswith(prefix) else key): value for key, value in out.items()} changed = True break if any(key.startswith(("decoder.", "encoder.", "bottleneck.")) for key in out): break return out def _sound_tokenizer_filter_decoder(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: return {key: value for key, value in state_dict.items() if key.startswith("decoder.")} def _sound_tokenizer_infer_num_blocks(state_dict: dict[str, torch.Tensor]) -> int: block_indices: set[int] = set() for key in state_dict: match = re.match(r"decoder\.layers\.(\d+)\.layers\.\d+\.", key) if match: block_indices.add(int(match.group(1))) return len(block_indices) def _sound_tokenizer_remap_flat_layout(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Convert legacy AVAE `decoder.layers.*` keys to OobleckDecoder attribute keys.""" if not any(re.match(r"decoder\.layers\.\d+\.", key) for key in state_dict): return state_dict num_blocks = _sound_tokenizer_infer_num_blocks(state_dict) if num_blocks == 0: raise RuntimeError("Detected flat `decoder.layers.*` layout but no decoder blocks were found; cannot remap.") snake1_idx = num_blocks + 1 conv2_idx = num_blocks + 2 def _remap(key: str) -> str: match = re.fullmatch(r"decoder\.layers\.(\d+)\.layers\.(\d+)\.layers\.(\d+)\.(.+)", key) if match: block_n, res_n, inner_n, rest = ( int(match.group(1)), int(match.group(2)), int(match.group(3)), match.group(4), ) if res_n not in (2, 3, 4): raise RuntimeError(f"Unexpected residual position {res_n} in {key!r}.") inner_name = _SOUND_TOKENIZER_RES_UNIT_INNER_NAMES.get(inner_n) if inner_name is None: raise RuntimeError(f"Unexpected residual inner index {inner_n} in {key!r}.") return f"decoder.block.{block_n - 1}.res_unit{res_n - 1}.{inner_name}.{rest}" match = re.fullmatch(r"decoder\.layers\.(\d+)\.layers\.(\d+)\.(.+)", key) if match: block_n, sub_n, rest = int(match.group(1)), int(match.group(2)), match.group(3) block_idx = block_n - 1 if sub_n == 0: return f"decoder.block.{block_idx}.snake1.{rest}" if sub_n == 1: return f"decoder.block.{block_idx}.conv_t1.{rest}" raise RuntimeError(f"Unexpected decoder block sub-index {sub_n} in {key!r}.") match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", key) if match: layer_n, rest = int(match.group(1)), match.group(2) if layer_n == 0: return f"decoder.conv1.{rest}" if layer_n == snake1_idx: return f"decoder.snake1.{rest}" if layer_n == conv2_idx: return f"decoder.conv2.{rest}" raise RuntimeError( f"Unexpected decoder leaf layer index {layer_n} (expected 0, {snake1_idx}, or {conv2_idx}) in {key!r}." ) return key return {_remap(key): value for key, value in state_dict.items()} def _sound_tokenizer_reshape_snake_params(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: out: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): if (key.endswith(".alpha") or key.endswith(".beta")) and value.ndim == 1: value = value.unsqueeze(0).unsqueeze(-1).contiguous() out[key] = value return out def _sound_tokenizer_reapply_weight_norm(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Reconstruct weight-norm parameters if the source checkpoint has folded conv weights.""" out = dict(state_dict) candidate_keys = [ key for key in state_dict if key.endswith(".weight") and any(f".{layer}." in key for layer in ("conv1", "conv2", "conv_t1")) ] for key in candidate_keys: stem = key[: -len(".weight")] weight_g_key = f"{stem}.weight_g" weight_v_key = f"{stem}.weight_v" if weight_g_key in state_dict or weight_v_key in state_dict: continue weight = state_dict[key] norm_dims = tuple(range(1, weight.ndim)) out.pop(key) out[weight_g_key] = weight.norm(p=2, dim=norm_dims, keepdim=True).contiguous() out[weight_v_key] = weight.contiguous() return out def _remap_avae_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: """Convert a legacy AVAE state dict into the Cosmos3AVAEAudioTokenizer state dict.""" state_dict = _sound_tokenizer_strip_per_key_prefixes(state_dict) state_dict = _sound_tokenizer_filter_decoder(state_dict) if not state_dict: raise RuntimeError("Sound tokenizer state dict has no `decoder.*` keys after prefix stripping.") state_dict = _sound_tokenizer_remap_flat_layout(state_dict) state_dict = _sound_tokenizer_reshape_snake_params(state_dict) state_dict = _sound_tokenizer_reapply_weight_norm(state_dict) if any(re.match(r"decoder\.layers\.\d+", key) for key in state_dict): raise RuntimeError("Flat `decoder.layers.*` keys remain after remap; conversion is incomplete.") return state_dict def _build_sound_tokenizer( checkpoint_path: pathlib.Path, config_path: pathlib.Path | None, ) -> Cosmos3AVAEAudioTokenizer: config = _load_sound_tokenizer_config(config_path, fallback_config_path=pathlib.Path()) print(f"Loading AVAE sound tokenizer weights from {checkpoint_path} …") raw_state_dict = _load_sound_tokenizer_state_dict(checkpoint_path) state_dict = _remap_avae_state_dict(raw_state_dict) print(f" Remapped {len(raw_state_dict)} → {len(state_dict)} decoder keys.") sound_tokenizer = Cosmos3AVAEAudioTokenizer( sampling_rate=config.get("sampling_rate", DEFAULT_SOUND_TOKENIZER_CONFIG["sampling_rate"]), vocoder_input_dim=config.get("vocoder_input_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["vocoder_input_dim"]), dec_dim=config.get("dec_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_dim"]), dec_c_mults=tuple(config.get("dec_c_mults", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_c_mults"])), dec_strides=tuple(config.get("dec_strides", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_strides"])), dec_out_channels=config.get("dec_out_channels", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_out_channels"]), ) load_result = sound_tokenizer.load_state_dict(state_dict, strict=True) if load_result.missing_keys or load_result.unexpected_keys: raise RuntimeError( "Cosmos3 AVAE sound tokenizer load did not match strictly: " f"missing={load_result.missing_keys}, unexpected={load_result.unexpected_keys}." ) return sound_tokenizer @contextlib.contextmanager def _skip_source_sound_tokenizer_load(omni_mot_model_cls): original_set_up_tokenizers = omni_mot_model_cls.set_up_tokenizers def set_up_tokenizers_without_sound(self): if not getattr(self.config, "sound_gen", False): return original_set_up_tokenizers(self) sound_gen = self.config.sound_gen self.config.sound_gen = False try: return original_set_up_tokenizers(self) finally: self.config.sound_gen = sound_gen omni_mot_model_cls.set_up_tokenizers = set_up_tokenizers_without_sound try: yield finally: omni_mot_model_cls.set_up_tokenizers = original_set_up_tokenizers def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--checkpoint-path", default="Cosmos3-Nano", help="Named checkpoint (e.g. 'Cosmos3-Nano') or path to a DCP checkpoint directory.", ) parser.add_argument("--output", required=True, help="Directory to save the converted diffusers model.") parser.add_argument( "--save-pipeline", action="store_true", help="Save the full pipeline (transformer + VAE + tokenizer + scheduler).", ) parser.add_argument( "--dtype", default="bf16", choices=["fp32", "fp16", "bf16"], help="Dtype to save the transformer in." ) parser.add_argument( "--sound-tokenizer-path", help="Optional AVAE sound tokenizer checkpoint to save under sound_tokenizer/." ) parser.add_argument( "--sound-tokenizer-config-path", help="Optional AVAE config JSON to save under sound_tokenizer/config.json." ) parser.add_argument( "--include-sound-tokenizer", action="store_true", help="Require saving sound_tokenizer/ even if the source transformer is video-only.", ) parser.add_argument( "--use-ema", action=argparse.BooleanOptionalAction, default=True, help="Use `net_ema` weights when converting a raw Cosmos3 Edge DCP checkpoint.", ) parser.add_argument( "--include-reasoner", action=argparse.BooleanOptionalAction, default=True, help=( "Add the pinned Cosmos3 Edge reasoner in shared-weight form so the pipeline output is usable by " "Transformers as well as Diffusers." ), ) parser.add_argument( "--reasoner-repo-id", default=COSMOS3_EDGE_REASONER, help="Hugging Face repository containing the Cosmos3 Edge reasoner checkpoint.", ) parser.add_argument( "--reasoner-revision", default=COSMOS3_EDGE_REASONER_REVISION, help="Pinned revision of the Cosmos3 Edge reasoner checkpoint.", ) parser.add_argument( "--reasoner-path", help="Optional local Cosmos3 Edge reasoner snapshot, used instead of downloading --reasoner-repo-id.", ) parser.add_argument( "--copy-edge-reasoner", action="store_true", help=( "Add the pinned shared-weight reasoner and this converter to an existing Cosmos3 Edge Diffusers " "pipeline at --output." ), ) args = parser.parse_args() if args.copy_edge_reasoner: _add_edge_reasoner_to_pipeline(args) return dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype] raw_checkpoint_path = pathlib.Path(args.checkpoint_path).expanduser() edge_dcp_model_dir = _find_edge_dcp_model_dir(raw_checkpoint_path) if raw_checkpoint_path.exists() else None if edge_dcp_model_dir is not None: _convert_edge_dcp(args, edge_dcp_model_dir, dtype) return from cosmos3.common.init import init_script init_script() from accelerate import init_empty_weights from cosmos3.args import _CHECKPOINTS from cosmos3.model import Cosmos3OmniModel from projects.cosmos3.vfm.models.omni_mot_model import OmniMoTModel from transformers import AutoTokenizer from diffusers import AutoencoderKLWan, UniPCMultistepScheduler from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline sound_tokenizer_path = ( pathlib.Path(args.sound_tokenizer_path).expanduser().absolute() if args.sound_tokenizer_path else None ) sound_tokenizer_config_path = ( pathlib.Path(args.sound_tokenizer_config_path).expanduser().absolute() if args.sound_tokenizer_config_path else None ) if args.include_sound_tokenizer and sound_tokenizer_path is None: raise ValueError("Sound tokenizer output was requested, but --sound-tokenizer-path was not provided.") if sound_tokenizer_path is not None and not sound_tokenizer_path.exists(): raise FileNotFoundError(f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}") if sound_tokenizer_config_path is not None and not sound_tokenizer_config_path.exists(): raise FileNotFoundError(f"Sound tokenizer config not found: {sound_tokenizer_config_path}") checkpoint_name = args.checkpoint_path if checkpoint_name in _CHECKPOINTS: checkpoint_path = pathlib.Path(_CHECKPOINTS[checkpoint_name].download()) else: checkpoint_path = pathlib.Path(checkpoint_name).expanduser().absolute() print(f"Resolved checkpoint path: {checkpoint_path}") print("Instantiating model and loading weights from DCP checkpoint …") print("Skipping source AVAE tokenizer instantiation during converter-only model load …") with _skip_source_sound_tokenizer_load(OmniMoTModel): _tmp = Cosmos3OmniModel.from_pretrained_dcp(checkpoint_path).model # Extract network components and architecture config from DCP model language_model = _tmp.net.language_model vae2llm = _tmp.net.vae2llm llm2vae = _tmp.net.llm2vae time_embedder = _tmp.net.time_embedder lm_cfg = _tmp.net.language_model.config net_cfg = _tmp.net.config model_cfg = _tmp.config patch_latent_dim = _tmp.net.patch_latent_dim hidden_size = _tmp.net.hidden_size num_attention_heads = _tmp.net.num_heads num_key_value_heads = _tmp.net.num_kv_heads head_dim = _tmp.net.head_dim num_hidden_layers = _tmp.net.num_hidden_layers latent_patch_size = _tmp.net.latent_patch_size latent_channel = _tmp.net.latent_channel timestep_scale = _tmp.net.timestep_scale base_fps = int(net_cfg.base_fps) enable_fps_modulation = net_cfg.enable_fps_modulation unified_3d_mrope_reset_spatial_ids = _tmp.config.diffusion_expert_config.unified_3d_mrope_reset_spatial_ids unified_3d_mrope_temporal_modality_margin = ( _tmp.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin ) sound2llm = getattr(_tmp.net, "sound2llm", None) llm2sound = getattr(_tmp.net, "llm2sound", None) sound_modality_embed = getattr(_tmp.net, "sound_modality_embed", None) has_sound_projection_weights = any(module is not None for module in (sound2llm, llm2sound, sound_modality_embed)) sound_gen = bool( _get_config_value(net_cfg, model_cfg, name="sound_gen", default=False) or has_sound_projection_weights ) sound_dim = _get_config_value(net_cfg, model_cfg, name="sound_dim", default=None) if sound_dim is None and sound2llm is not None: sound_dim = sound2llm.in_features sound_latent_fps = _get_config_value(net_cfg, model_cfg, name="sound_latent_fps", default=25.0) if sound_gen: missing_sound_modules = [ name for name, module in ( ("sound2llm", sound2llm), ("llm2sound", llm2sound), ("sound_modality_embed", sound_modality_embed), ) if module is None ] if missing_sound_modules: raise RuntimeError( "Source checkpoint is configured for sound generation but is missing " f"sound projection weights: {missing_sound_modules}." ) if sound_dim is None: raise RuntimeError("Source checkpoint is configured for sound generation but sound_dim is missing.") del _tmp torch.cuda.empty_cache() # Init diffusers Cosmos3OmniTransformer with full architecture config from DCP with init_empty_weights(): transformer = Cosmos3OmniTransformer( attention_bias=lm_cfg.attention_bias, attention_dropout=lm_cfg.attention_dropout, base_fps=base_fps, enable_fps_modulation=enable_fps_modulation, head_dim=head_dim, hidden_size=hidden_size, intermediate_size=lm_cfg.intermediate_size, latent_channel=latent_channel, latent_patch_size=latent_patch_size, num_attention_heads=num_attention_heads, num_hidden_layers=num_hidden_layers, num_key_value_heads=num_key_value_heads, patch_latent_dim=patch_latent_dim, rms_norm_eps=lm_cfg.rms_norm_eps, rope_scaling=lm_cfg.rope_scaling, rope_theta=lm_cfg.rope_theta, sound_dim=sound_dim, sound_gen=sound_gen, sound_latent_fps=sound_latent_fps, timestep_scale=timestep_scale, unified_3d_mrope_reset_spatial_ids=unified_3d_mrope_reset_spatial_ids, unified_3d_mrope_temporal_modality_margin=unified_3d_mrope_temporal_modality_margin, vocab_size=lm_cfg.vocab_size, backbone_type="cosmos3_edge_nemotron_dense", ) # The source language_model nests its transformer stack under a `model.` attribute # (HF Qwen-style). Diffusers Cosmos3OmniTransformer holds those layers flat, so # strip the leading `model.` prefix from the language-model state-dict keys. state_dict = { (k[len("model.") :] if k.startswith("model.") else k): v for k, v in language_model.state_dict().items() } # Remap PackedAttentionMoT attribute names from the source (Qwen-style q_proj/k_proj/... # plus cosmos-specific *_moe_gen) to the diffusers AttentionModuleMixin canonical names. # Order matters: the *_moe_gen substrings must be substituted before the plain ones. _ATTN_KEY_REMAP = [ (".q_proj_moe_gen.", ".add_q_proj."), (".k_proj_moe_gen.", ".add_k_proj."), (".v_proj_moe_gen.", ".add_v_proj."), (".o_proj_moe_gen.", ".to_add_out."), (".q_norm_moe_gen.", ".norm_added_q."), (".k_norm_moe_gen.", ".norm_added_k."), (".q_proj.", ".to_q."), (".k_proj.", ".to_k."), (".v_proj.", ".to_v."), (".o_proj.", ".to_out."), (".q_norm.", ".norm_q."), (".k_norm.", ".norm_k."), ] remapped_state_dict: dict[str, torch.Tensor] = {} for k, v in state_dict.items(): for old, new in _ATTN_KEY_REMAP: if old in k: k = k.replace(old, new) break remapped_state_dict[k] = v state_dict = remapped_state_dict for k, v in vae2llm.state_dict().items(): state_dict[f"proj_in.{k}"] = v for k, v in llm2vae.state_dict().items(): state_dict[f"proj_out.{k}"] = v _TIME_EMBEDDER_REMAP = { "mlp.0.weight": "linear_1.weight", "mlp.0.bias": "linear_1.bias", "mlp.2.weight": "linear_2.weight", "mlp.2.bias": "linear_2.bias", } for k, v in time_embedder.state_dict().items(): state_dict[f"time_embedder.{_TIME_EMBEDDER_REMAP[k]}"] = v if sound_gen: for k, v in sound2llm.state_dict().items(): state_dict[f"audio_proj_in.{k}"] = v for k, v in llm2sound.state_dict().items(): state_dict[f"audio_proj_out.{k}"] = v state_dict["audio_modality_embed"] = sound_modality_embed transformer.load_state_dict(state_dict, strict=True, assign=True) del ( language_model, vae2llm, llm2vae, time_embedder, sound2llm, llm2sound, sound_modality_embed, state_dict, ) torch.cuda.empty_cache() transformer = transformer.to(dtype=dtype) output_dir = pathlib.Path(args.output) output_dir.mkdir(parents=True, exist_ok=True) include_sound_tokenizer = ( args.include_sound_tokenizer or sound_tokenizer_path is not None or (sound_gen and args.save_pipeline) ) if include_sound_tokenizer and sound_tokenizer_path is None: raise ValueError( "The source checkpoint is configured for sound generation, so --sound-tokenizer-path " "is required when saving a full pipeline." ) if args.save_pipeline: text_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL-8B-Instruct") diffusers_vae = AutoencoderKLWan.from_pretrained( "Wan-AI/Wan2.2-TI2V-5B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16 ) sound_tokenizer = None if include_sound_tokenizer: assert sound_tokenizer_path is not None sound_tokenizer = _build_sound_tokenizer(sound_tokenizer_path, sound_tokenizer_config_path) # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps. # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281. # EDM sigma = flow_sigma / (1 - flow_sigma), so: # sigma_max = 0.9998 / 0.0002 = 4999 (but capped at 200 to avoid duplicate # integer timesteps from Karras clustering near the top) # sigma_min = 0.1281 / (1 - 0.1281) = 0.1281 / 0.8719 ≈ 0.147 scheduler = UniPCMultistepScheduler( use_karras_sigmas=True, use_flow_sigmas=True, prediction_type="flow_prediction", sigma_max=200.0, sigma_min=0.147, ) pipeline = Cosmos3OmniPipeline( transformer=transformer, text_tokenizer=text_tokenizer, vae=diffusers_vae, scheduler=scheduler, sound_tokenizer=sound_tokenizer, ) print(f"Saving full pipeline to {output_dir} …") pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") else: print(f"Saving transformer to {output_dir} …") transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") if include_sound_tokenizer: print("Skipping sound_tokenizer/ save because --save-pipeline was not set.") print("Done.") if __name__ == "__main__": main()