Spaces:
Running on L40S
Running on L40S
File size: 6,143 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
import json
from pathlib import Path
import attrs
import hydra
import safetensors.torch
import torch
import torch.distributed.checkpoint as dcp
from cosmos_framework.configs.base.defaults.compile import CompileConfig
from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
from cosmos_framework.inference.args import _CHECKPOINTS, DEFAULT_CHECKPOINT
from cosmos_framework.inference.common.args import CheckpointType
from cosmos_framework.inference.common.config import structure_config
from cosmos_framework.inference.model import (
Cosmos3OmniConfig,
_diffusers_to_net_key,
_diffusers_weight_map,
_DiffusersHuggingFaceStorageReader,
_DiffusersLoadPlanner,
_is_diffusers_checkpoint,
_normalize_diffusers_target_key,
)
def test_config():
parallelism = ParallelismConfig(
data_parallel_shard_degree=2,
context_parallel_shard_degree=2,
cfg_parallel_shard_degree=2,
)
compile = CompileConfig(enabled=True, use_cuda_graphs=True)
checkpoint_path = DEFAULT_CHECKPOINT.download()
config = Cosmos3OmniConfig.from_pretrained(
checkpoint_path,
parallelism=attrs.asdict(parallelism),
compile=attrs.asdict(compile),
)
assert hydra.utils.instantiate(structure_config(config.parallelism, ParallelismConfig)) == parallelism
assert hydra.utils.instantiate(structure_config(config.compile, CompileConfig)) == compile
def test_checkpoint_type_from_path_hf_index(tmp_path: Path):
(tmp_path / "config.json").write_text("{}", encoding="utf-8")
(tmp_path / "model.safetensors.index.json").write_text("{}", encoding="utf-8")
assert CheckpointType.from_path(tmp_path) == CheckpointType.HF
def test_normalize_diffusers_target_key():
assert (
_normalize_diffusers_target_key(
"model.net._orig_mod.language_model.model.layers.0._checkpoint_wrapped_module.input_layernorm.weight"
)
== "language_model.model.layers.0.input_layernorm.weight"
)
def test_diffusers_to_net_key():
cases = {
"lm_head.weight": "language_model.lm_head.weight",
"embed_tokens.weight": "language_model.model.embed_tokens.weight",
"norm_moe_gen.weight": "language_model.model.norm_moe_gen.weight",
"layers.18.self_attn.to_q.weight": "language_model.model.layers.18.self_attn.q_proj.weight",
"layers.18.self_attn.to_out.weight": "language_model.model.layers.18.self_attn.o_proj.weight",
"layers.18.self_attn.norm_q.weight": "language_model.model.layers.18.self_attn.q_norm.weight",
"layers.18.self_attn.add_k_proj.weight": "language_model.model.layers.18.self_attn.k_proj_moe_gen.weight",
"layers.18.self_attn.to_add_out.weight": "language_model.model.layers.18.self_attn.o_proj_moe_gen.weight",
"layers.18.self_attn.norm_added_k.weight": "language_model.model.layers.18.self_attn.k_norm_moe_gen.weight",
"language_model.model.layers.18.self_attn.to_q.weight": "language_model.model.layers.18.self_attn.q_proj.weight",
"proj_in.weight": "vae2llm.weight",
"proj_out.bias": "llm2vae.bias",
"time_embedder.linear_1.weight": "time_embedder.mlp.0.weight",
"time_embedder.linear_2.bias": "time_embedder.mlp.2.bias",
"audio_proj_in.weight": "sound2llm.weight",
"audio_proj_out.bias": "llm2sound.bias",
"audio_modality_embed": "sound_modality_embed",
"action_proj_in.fc.weight": "action2llm.fc.weight",
"action_proj_out.bias.weight": "llm2action.bias.weight",
"action_modality_embed": "action_modality_embed",
}
for diffusers_key, net_key in cases.items():
assert _diffusers_to_net_key(diffusers_key, "transformer/diffusion_pytorch_model.safetensors") == net_key
assert (
_diffusers_to_net_key("blocks.0.attn.qkv.weight", "vision_encoder/model.safetensors")
== "language_model.visual.blocks.0.attn.qkv.weight"
)
assert _diffusers_to_net_key("decoder.conv.weight", "vae/diffusion_pytorch_model.safetensors") is None
def test_diffusers_dcp_load_remaps_nested_safetensors(tmp_path: Path):
shard_rel_path = "transformer/diffusion_pytorch_model.safetensors"
shard_path = tmp_path / shard_rel_path
shard_path.parent.mkdir(parents=True)
source = torch.arange(6, dtype=torch.float32).reshape(2, 3)
safetensors.torch.save_file({"proj_in.weight": source}, shard_path)
(tmp_path / "model.safetensors.index.json").write_text(
json.dumps(
{
"metadata": {},
"weight_map": {
"proj_in.weight": shard_rel_path,
"decoder.conv.weight": "vae/diffusion_pytorch_model.safetensors",
},
}
),
encoding="utf-8",
)
target = {"model.net._orig_mod.vae2llm.weight": torch.empty_like(source)}
dcp.load(
state_dict=target,
storage_reader=_DiffusersHuggingFaceStorageReader(tmp_path),
planner=_DiffusersLoadPlanner(tmp_path),
)
torch.testing.assert_close(target["model.net._orig_mod.vae2llm.weight"], source)
def test_diffusers_weight_map_registered_checkpoint():
checkpoint_path = Path(_CHECKPOINTS["Cosmos3-Nano"].hf.download())
assert (checkpoint_path / "model_index.json").exists()
assert (checkpoint_path / "model.safetensors.index.json").exists()
assert CheckpointType.from_path(checkpoint_path) == CheckpointType.HF
assert _is_diffusers_checkpoint(checkpoint_path)
weight_map = _diffusers_weight_map(checkpoint_path)
assert weight_map["proj_in.weight"].startswith("transformer/")
assert weight_map["blocks.0.attn.qkv.weight"] == "vision_encoder/model.safetensors"
assert _diffusers_to_net_key("proj_in.weight", weight_map["proj_in.weight"]) == "vae2llm.weight"
assert (
_diffusers_to_net_key("blocks.0.attn.qkv.weight", weight_map["blocks.0.attn.qkv.weight"])
== "language_model.visual.blocks.0.attn.qkv.weight"
)
|