Spaces:
Running on Zero
Running on Zero
File size: 3,205 Bytes
a1833c8 | 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Convert internal training config to public config."""
from cosmos_framework.inference.common.init import init_script
init_script()
from pathlib import Path
from typing import Annotated
import omegaconf
import pydantic
import tyro
from cosmos_framework.inference.common.args import CheckpointOverrides, ResolvedFilePath, ResolvedPath, tyro_cli
from cosmos_framework.inference.common.config import (
CONFIG_DIR,
deserialize_config_dict,
serialize_config,
structure_config,
)
from cosmos_framework.inference.common.init import init_output_dir
def _validate_config_file(v: Path) -> Path:
if v.suffix != ".yaml":
raise ValueError(f"Config file must be a YAML file: {v}")
return v
ConfigFilePath = Annotated[ResolvedFilePath, pydantic.AfterValidator(_validate_config_file)]
class Args(pydantic.BaseModel):
output_dir: Annotated[ResolvedPath, tyro.conf.arg(aliases=("-o",))]
config_file: ConfigFilePath
"""Hydra config yaml file."""
config_overrides: list[str] = pydantic.Field(default_factory=list)
"""Hydra config overrides."""
checkpoint: CheckpointOverrides = CheckpointOverrides.model_construct()
"""Checkpoint arguments."""
def export_train_config(args: Args) -> None:
checkpoint_args = args.checkpoint.build_checkpoint(checkpoints={})
init_output_dir(args.output_dir)
# Load config
config_dict = deserialize_config_dict(args.config_file)
overrides_omegaconf = omegaconf.OmegaConf.from_dotlist(args.config_overrides)
config_omegaconf = omegaconf.OmegaConf.merge(config_dict, overrides_omegaconf)
config_omegaconf.job.wandb_mode = "disabled"
# Set checkpoint
checkpoint_dict = deserialize_config_dict(CONFIG_DIR / "checkpoint/local.yaml")
checkpoint_dict["type"] = config_omegaconf.checkpoint.type
checkpoint_dict["load_path"] = "???"
checkpoint_dict["keys_to_skip_loading"] = ["net_ema."]
config_omegaconf.checkpoint = checkpoint_dict
# Set model
model_dict = checkpoint_args.load_model_config_dict()
model_dict["config"]["ema"] = config_omegaconf.model.config.ema
config_omegaconf.model = model_dict
# Filter callbacks
def process_callback(name: str, callback: omegaconf.DictConfig) -> omegaconf.DictConfig | None:
if hasattr(callback, "save_s3"):
setattr(callback, "save_s3", False)
return callback
callbacks_dict = {}
for k, v in config_omegaconf.trainer.callbacks.items():
v = process_callback(k, v)
if v is not None:
callbacks_dict[k] = v
config_omegaconf.trainer.callbacks = callbacks_dict
omegaconf.OmegaConf.save(config_omegaconf, args.output_dir / "config_raw.yaml")
omegaconf.OmegaConf.resolve(config_omegaconf)
config = structure_config(config_omegaconf)
config.validate()
serialize_config(config, args.output_dir / "config.yaml")
def main() -> None:
args = tyro_cli(Args, description=__doc__, config=(tyro.conf.OmitArgPrefixes,))
export_train_config(args)
if __name__ == "__main__":
main()
|