Spaces:
Running on L40S
Running on L40S
File size: 2,577 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 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1
"""Export config defaults and schemas."""
from cosmos_framework.inference.common.init import init_script
init_script(default_env={"COSMOS_TRAINING": "0"})
import argparse
import json
import pathlib
import textwrap
import typing
import pydantic
import yaml
from cosmos_framework.inference.args import OmniSampleOverrides, OmniSetupOverrides
MODELS: list[tuple[type[pydantic.BaseModel], dict[str, typing.Any]]] = [
(OmniSetupOverrides, {}),
(OmniSampleOverrides, {}),
]
def _nested_model_cls(field: pydantic.fields.FieldInfo) -> type[pydantic.BaseModel] | None:
"""Return the BaseModel subclass backing *field*, if any."""
ann = field.annotation
if isinstance(ann, type) and issubclass(ann, pydantic.BaseModel):
return ann
for arg in typing.get_args(ann):
if isinstance(arg, type) and issubclass(arg, pydantic.BaseModel):
return arg
return None
def _commented_yaml(model_cls: type, data: dict) -> str:
lines: list[str] = []
for name, field in sorted(model_cls.model_fields.items()):
if name not in data:
continue
if field.description:
for dl in field.description.strip().splitlines():
s = dl.strip()
lines.append(f"# {s}" if s else "#")
value = data[name]
nested = _nested_model_cls(field)
if nested and isinstance(value, dict):
lines.append(f"{name}:")
lines.append(textwrap.indent(_commented_yaml(nested, value).rstrip("\n"), " "))
else:
lines.append(yaml.dump({name: value}, default_flow_style=False, sort_keys=False).rstrip())
return "\n".join(lines) + "\n"
def export_schemas(output_dir: pathlib.Path) -> None:
output_dir.mkdir(parents=True, exist_ok=True)
for cls, defaults in MODELS:
data = cls.model_construct(**defaults).model_dump(mode="json")
(output_dir / f"{cls.__name__}.yaml").write_text(_commented_yaml(cls, data))
(output_dir / f"{cls.__name__}.schema.json").write_text(json.dumps(cls.model_json_schema(), indent=2) + "\n")
print(f"Saved {cls.__name__} -> {output_dir}")
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("-o", "--output", type=pathlib.Path, default="schemas", help="Output directory")
args = parser.parse_args()
export_schemas(args.output.absolute())
if __name__ == "__main__":
main()
|