convert_model_to_diffusers: export via official diffusers, drop diffusers-cosmos3 shim

cosmos_framework/scripts/convert_model_to_diffusers.py

@@ -40,18 +40,6 @@ class Args(pydantic.BaseModel):
     config_only: bool = False
     """If True, only save config."""
 
-    remap_sound_tokenizer_keys: bool = True
-    """If True, remap the legacy AVAE sound tokenizer state dict into the
-    diffusers OobleckDecoder layout (and save the result as
-    `diffusion_pytorch_model.safetensors`). Off by default — the sound
-    tokenizer is written verbatim under `model.safetensors`."""
-
-    remap_time_embedder_keys: bool = True
-    """If True, remap the transformer's `time_embedder` state dict from the
-    legacy `nn.Sequential` layout (`mlp.0.*` / `mlp.2.*`) to the diffusers
-    `TimestepEmbedding` layout (`linear_1.*` / `linear_2.*`). Off by default —
-    keys are forwarded verbatim."""
-
 
 class SafetensorsIndexMetadata(pydantic.BaseModel):
     total_size: int = 0
@@ -93,9 +81,17 @@ def convert_model_to_diffusers(args: Args):
     sound_tokenizer_dir, sound_tokenizer_name = model_dict["config"]["sound_tokenizer"]["avae_path"].rsplit("/", 1)
     sound_tokenizer_checkpoint = CheckpointConfig.maybe_from_uri(f"s3://bucket/{sound_tokenizer_dir}")
     assert sound_tokenizer_checkpoint is not None
-    sound_tokenizer_path = Path(sound_tokenizer_checkpoint.hf.download()) / sound_tokenizer_name
+    sound_tokenizer_local = Path(sound_tokenizer_checkpoint.hf.download())
+    # HF-published checkpoints ship sound_tokenizer/ in the diffusers layout
+    # (config.json + diffusion_pytorch_model.safetensors), which the converter
+    # consumes directly; fall back to the legacy AVAE file pair named by the
+    # model config's avae_path.
+    sound_tokenizer_path = sound_tokenizer_local / "diffusion_pytorch_model.safetensors"
+    sound_tokenizer_config_path = sound_tokenizer_local / "config.json"
+    if not sound_tokenizer_path.is_file():
+        sound_tokenizer_path = sound_tokenizer_local / sound_tokenizer_name
+        sound_tokenizer_config_path = sound_tokenizer_path.with_suffix(".json")
     assert sound_tokenizer_path.is_file(), f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}"
-    sound_tokenizer_config_path = sound_tokenizer_path.with_suffix(".json")
     assert sound_tokenizer_config_path.is_file(), f"Sound tokenizer config not found: {sound_tokenizer_config_path}"
 
     vision_encoder_model = model_dict["config"]["vlm_config"]["tokenizer"]["pretrained_model_name"]
@@ -116,8 +112,6 @@ def convert_model_to_diffusers(args: Args):
             sound_tokenizer_path=str(sound_tokenizer_path),
             sound_tokenizer_config_path=str(sound_tokenizer_config_path),
             include_sound_tokenizer=True,
-            remap_sound_tokenizer_keys=args.remap_sound_tokenizer_keys,
-            remap_time_embedder_keys=args.remap_time_embedder_keys,
             vision_encoder_model=vision_encoder_model,
             skip_vision_encoder=False,
         )

docs/code_structure.md

@@ -38,7 +38,7 @@ Cosmos/
 │   ├── inference/      # Inference subpackage (model, args, defaults, Ray serving, common helpers, SFT experiment configs)
 │   └── ...             # Training-infra subpackages: data, model, trainer, callbacks, checkpoint, …
 │   └── scripts/        # CLI entry-point scripts: train.py, _train.py, inference.py, export_model.py, …
-├── packages/           # Backend shim packages: diffusers-cosmos3, transformers-cosmos3, vllm-cosmos3
+├── packages/           # Backend shim packages: transformers-cosmos3, vllm-cosmos3
 ├── docs/               # User documentation (you are here)
 ├── docker/             # Dockerfiles for reproducible environments
 ├── examples/           # Runnable training / fine-tuning / inference examples
@@ -49,7 +49,7 @@ Cosmos/
 └── .python-version     # Python version pin (used by uv)
 ```
 
-`cosmos_framework/` is the single Python package. Training infrastructure (data, model, trainer, callbacks, checkpoint, utils, …) lives in top-level subpackages; inference (Diffusers / Transformers / vLLM-friendly inference core, Ray serving, per-modality defaults, training-side experiment YAMLs) lives under `cosmos_framework/inference/`. The library-style backend shims that load Cosmos3 checkpoints into upstream ecosystems live under `packages/{diffusers,transformers,vllm}-cosmos3/`.
+`cosmos_framework/` is the single Python package. Training infrastructure (data, model, trainer, callbacks, checkpoint, utils, …) lives in top-level subpackages; inference (Diffusers / Transformers / vLLM-friendly inference core, Ray serving, per-modality defaults, training-side experiment YAMLs) lives under `cosmos_framework/inference/`. The library-style backend shims that load Cosmos3 checkpoints into upstream ecosystems live under `packages/{transformers,vllm}-cosmos3/` (Diffusers needs no shim: diffusers ≥ 0.39 ships `Cosmos3OmniPipeline` natively).
 
 ## The `cosmos_framework/` Package
 
@@ -154,7 +154,7 @@ The full inference subpackage:
 
 Training-side experiment SKUs live separately at `cosmos_framework/configs/base/experiment/sft/*.py` (see [`cosmos_framework/configs/`](#cosmos_frameworkconfigs)) — not under `inference/`.
 
-Library-style backend shims that adapt Cosmos3 checkpoints to the Diffusers / Transformers / vLLM ecosystems live separately under `packages/{diffusers,transformers,vllm}-cosmos3/`.
+Library-style backend shims that adapt Cosmos3 checkpoints to the Transformers / vLLM ecosystems live separately under `packages/{transformers,vllm}-cosmos3/`. Diffusers needs no shim: diffusers ≥ 0.39 ships `Cosmos3OmniPipeline` / `Cosmos3OmniTransformer` natively.
 
 ### `cosmos_framework/launcher/`
 
@@ -197,7 +197,7 @@ Add new worker types as sibling subpackages — each owns its own startup, messa
 - `examples/` — runnable end-to-end examples; see `examples/README.md`.
 - `docker/` — Dockerfiles and image build helpers; see `docker/README.md`.
 - `cosmos_framework/scripts/` — CLI entry-point scripts (`train.py`, `inference.py`, `export_model.py`, …); invoke as `python -m cosmos_framework.scripts.<name>`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=<recipe-toml>`); recipe TOMLs live under [`examples/toml/sft_config/`](../examples/toml/sft_config/) and the schema is defined in [`cosmos_framework/configs/toml_config/sft_config.py`](../cosmos_framework/configs/toml_config/sft_config.py) — see [`examples/README.md`](../examples/README.md) and [`docs/training.md`](./training.md).
-- `packages/` — library-style backend shims: `packages/{diffusers,transformers,vllm}-cosmos3/`, each installable independently.
+- `packages/` — library-style backend shims: `packages/{transformers,vllm}-cosmos3/`, each installable independently.
 
 ## Where to Add New Code