Support Cosmos3-Super task-specialized (Text2Image / Image2Video) checkpoints

These task-specialized diffusers checkpoints reuse the Cosmos3-Super
architecture but omit unused modality weights and bundle their own VLM
processor. Loading them previously failed, and the processor pulled a
redundant full base-Super download.

- inference/model.py: tolerate absent action/sound projection-head
  weights in the diffusers load planner, mirroring the existing vision
  carve-out. Fixes the masked "TypeError: cannot pickle code objects"
  that surfaced when DCP tried to broadcast the missing-tensor
  ValueError across ranks. No-op for self-consistent base checkpoints:
  Nano/Super provide all modality weights, so the guards never fire.

- inference: add CheckpointConfig.vlm_processor_from_checkpoint. When
  set, the loader sources the VLM processor from the loaded checkpoint's
  own bundled files instead of the repository hardcoded in the model
  config, avoiding a redundant base-Super download. Enabled for the two
  task checkpoints; base Nano/Super keep their configured repository.

- data/vfm/processors: clearer error when build_processor_lazy is given
  neither a repository nor a tokenizer_type source (never fires for
  existing call sites; only improves a previously-TypeError path).

- docs/faq.md: add EADDRINUSE / --master-port entry.
- docs/superpowers/specs: design spec for the processor-source change.

Verified: Text2Image (t2i) and Image2Video (i2v) load and generate; a
full base Cosmos3-Nano t2i run is unchanged with strict weight loading
intact (carve-out never triggers).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: lfengad

cosmos_framework/data/vfm/processors/__init__.py

@@ -174,4 +174,9 @@ def build_processor_lazy(
         if subdir:
             local_path = os.path.join(local_path, subdir)
         return sys.modules[__name__].build_processor(local_path, **kwargs)
+    if not args and "tokenizer_type" not in kwargs:
+        raise ValueError(
+            "build_processor_lazy requires either 'repository' (+ 'revision') or a "
+            "'tokenizer_type' source (e.g. a local checkpoint directory); got neither."
+        )
     return sys.modules[__name__].build_processor(*args, **kwargs)

cosmos_framework/inference/args.py

@@ -869,6 +869,33 @@ def build_sample(self, *, model_config: Any) -> OmniSampleArgs:
             revision="main",
         ),
     ),
+    # Task-specialized Super variants published as diffusers HF checkpoints.
+    # s3_uri is unused for HF-backed checkpoints (kept for parity with the
+    # registry schema); the architecture lives in each model YAML.
+    "Cosmos3-Super-Image2Video": CheckpointConfig(
+        model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"],
+        config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"),
+        s3_uri="s3://bucket1/cosmos3_vfm/cosmos3_ga_image2video/",
+        hf=CheckpointDirHf(
+            repository="nvidia/Cosmos3-Super-Image2Video",
+            revision="main",
+        ),
+        # Self-contained checkpoint: use its bundled processor instead of
+        # downloading the base Cosmos3-Super repo just for the tokenizer.
+        vlm_processor_from_checkpoint=True,
+    ),
+    "Cosmos3-Super-Text2Image": CheckpointConfig(
+        model_memory_bytes=MODEL_MEMORY_BYTES_BY_SIZE["32B"],
+        config_file=str(CONFIG_DIR / "model/Cosmos3-Super.yaml"),
+        s3_uri="s3://bucket1/cosmos3_vfm/cosmos3_ga_text2image/",
+        hf=CheckpointDirHf(
+            repository="nvidia/Cosmos3-Super-Text2Image",
+            revision="main",
+        ),
+        # Self-contained checkpoint: use its bundled processor instead of
+        # downloading the base Cosmos3-Super repo just for the tokenizer.
+        vlm_processor_from_checkpoint=True,
+    ),
 }
 DEFAULT_CHECKPOINT_NAME = "Cosmos3-Nano"
 DEFAULT_CHECKPOINT = _CHECKPOINTS[DEFAULT_CHECKPOINT_NAME]

cosmos_framework/inference/common/args.py

@@ -385,6 +385,17 @@ class CheckpointConfig(pydantic.BaseModel):
     hf: CheckpointDirHf
     """Config for checkpoint on Hugging Face."""
 
+    vlm_processor_from_checkpoint: bool = False
+    """When True, load the VLM text/vision processor from the checkpoint's own
+    bundled files (its local download directory) instead of the repository
+    hardcoded in the model config's ``vlm_config.tokenizer`` node.
+
+    Set this only for self-contained checkpoints that ship their own processor
+    at the repository root (e.g. the task-specialized Text2Image / Image2Video
+    diffusers checkpoints). Avoids a redundant download of the base model repo
+    just to obtain the tokenizer.
+    """
+
     def download(self) -> str:
         return self.hf.download()
 
@@ -404,6 +415,7 @@ class CheckpointArgs(ConfigArgs):
     model_memory_bytes: int | None
 
     checkpoint_hf: CheckpointDirHf | None
+    vlm_processor_from_checkpoint: bool = False
 
     credential_path: str
     use_ema_weights: bool
@@ -443,6 +455,8 @@ class CheckpointOverrides(ConfigOverrides):
 
     checkpoint_hf: Suppress[CheckpointDirHf | None] = None
     """Hugging Face checkpoint directory."""
+    vlm_processor_from_checkpoint: Suppress[bool] = False
+    """Load the VLM processor from the loaded checkpoint instead of a hardcoded repo."""
 
     credential_path: Training[str] = "credentials/gcp_checkpoint.secret"
     """Path to S3 credentials file for remote checkpoint loading."""
@@ -459,6 +473,7 @@ def _build_checkpoint(self, checkpoints: dict[str, CheckpointConfig]):
             self.model_memory_bytes = checkpoint.model_memory_bytes
             self.config_file = checkpoint.config_file
             self.checkpoint_hf = checkpoint.hf
+            self.vlm_processor_from_checkpoint = checkpoint.vlm_processor_from_checkpoint
         elif self.checkpoint_path.startswith("s3://"):
             self.checkpoint_type = CheckpointType.DCP
             self.checkpoint_path = self.checkpoint_path.rstrip("/")

cosmos_framework/inference/inference.py

@@ -1050,6 +1050,15 @@ def _create(cls, setup_args: SetupArgs, **kwargs: Any) -> Self:
                 config = None
             else:
                 model_dict = setup_args.load_model_config_dict()
+                if setup_args.vlm_processor_from_checkpoint:
+                    # Source the VLM processor from the loaded checkpoint's own
+                    # bundled files instead of the repository hardcoded in the
+                    # model config. Drops the redundant base-model download.
+                    tokenizer_cfg = model_dict["config"]["vlm_config"]["tokenizer"]
+                    tokenizer_cfg.pop("repository", None)
+                    tokenizer_cfg.pop("revision", None)
+                    tokenizer_cfg.pop("subdir", None)
+                    tokenizer_cfg["tokenizer_type"] = str(checkpoint_path)
                 config = Cosmos3OmniConfig(model=model_dict)
             model = Cosmos3OmniModel.from_pretrained_dcp(
                 checkpoint_path,

cosmos_framework/inference/model.py

@@ -318,6 +318,16 @@ def set_up_planner(
         missing_keys = set(target_state_dict) - loaded_keys
         if not self.has_vision_weights:
             missing_keys = {key for key in missing_keys if not key.startswith("language_model.visual.")}
+        # Task-specialized checkpoints (e.g. Text2Image, Image2Video) omit the
+        # optional generative-modality projection heads (action, sound). They
+        # are unused for those tasks, so tolerate their absence the same way
+        # vision weights are tolerated when the checkpoint provides none of them.
+        for modality_prefixes in (
+            ("action2llm.", "llm2action.", "action_modality_embed"),
+            ("sound2llm.", "llm2sound.", "sound_modality_embed"),
+        ):
+            if not any(key.startswith(modality_prefixes) for key in loaded_keys):
+                missing_keys = {key for key in missing_keys if not key.startswith(modality_prefixes)}
         if missing_keys:
             sample = sorted(missing_keys)[:10]
             raise ValueError(

docs/faq.md

@@ -225,6 +225,25 @@ Delete the torchinductor cache under the /tmp directory, `rm -rf /tmp/torchinduc
 
 ---
 
+### Q: I get `torch.distributed.DistNetworkError: ... port: 29500 ... EADDRINUSE, address already in use`
+
+`torchrun` defaults its rendezvous to port `29500`. The error means that port is already taken on the node — usually because another `torchrun` job (yours or someone else's on a shared node) is still using it.
+
+Pass a different free port with `--master-port`, placed **before** `-m` (it is a `torchrun` argument, not an inference argument):
+
+```shell
+torchrun --nproc-per-node=8 --master-port=29501 -m cosmos_framework.scripts.inference \
+  --parallelism-preset=throughput \
+  -i "inputs/omni/t2i.json" \
+  -o outputs/omni_t2i \
+  --checkpoint-path Cosmos3-Super-Text2Image \
+  --seed=0
+```
+
+Any free port works (e.g. `29501`, `29510`); give each concurrent job on the same node a distinct port. Alternatively, `--rdzv-endpoint=localhost:0` lets `torchrun` auto-pick a free port.
+
+---
+
 ## Training
 
 ### Q: I get `torch.cuda.OutOfMemoryError` during training (SFT)

docs/superpowers/specs/2026-06-03-vlm-processor-from-checkpoint-design.md

@@ -0,0 +1,184 @@
+# Design: Optional VLM-Processor Source ("unset → use loaded checkpoint")
+
+**Date:** 2026-06-03
+**Status:** Approved (pending spec review)
+**Area:** Cosmos3 inference — model/processor loading
+
+## Problem
+
+When loading a checkpoint, the VLM text/vision processor (tokenizer, chat
+template, image/video preprocessor) is sourced from a repository hardcoded in
+the model config:
+
+```yaml
+# cosmos_framework/inference/configs/model/Cosmos3-Super.yaml
+vlm_config:
+  tokenizer:
+    _target_: cosmos3._src.vfm.processors.build_processor_lazy
+    repository: nvidia/Cosmos3-Super
+    revision: main
+```
+
+`build_processor_lazy(repository="nvidia/Cosmos3-Super", revision="main")`
+downloads that repo via `CheckpointDirHf(...).download()` (with `--include '*'`,
+i.e. **all 27 transformer shards of the 32B base model**) just to read the
+processor files from its root.
+
+This is wasteful and conceptually wrong for the **task-specialized** checkpoints
+(`Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`). Those are
+self-contained HF checkpoints that already bundle their own processor at the
+repository root (`preprocessor_config.json`, `video_preprocessor_config.json`,
+`tokenizer.json`, `tokenizer_config.json`, `chat_template.json`, `merges.txt`,
+`vocab.json`, plus a `text_tokenizer/` subdir). They reuse `Cosmos3-Super.yaml`
+only because they share the Super architecture — but doing so drags in a full,
+redundant download of the base Super repo solely for the processor.
+
+## Goal
+
+Make the processor source configurable with the following semantics, confirmed
+with the requester:
+
+- **`repository` unset in the tokenizer config → load the processor from the
+  checkpoint currently being loaded** (its already-downloaded local directory).
+- **`repository` set (today's `Cosmos3-Nano` / `Cosmos3-Super` case) → load from
+  that repository, exactly as today.**
+
+Base `Cosmos3-Nano` / `Cosmos3-Super` behavior must be unchanged. The
+task-specialized checkpoints must stop downloading base Super for the processor.
+
+## Key Facts (verified)
+
+- `OmniMoTModel.set_up_tokenizers()` builds the processor via
+  `self.vlm_processor = lazy_instantiate(self.vlm_config.tokenizer)`
+  (`cosmos_framework/model/vfm/omni_mot_model.py:140`). This runs inside
+  `Cosmos3OmniModel(config)` construction.
+- `build_processor_lazy(*args, repository=None, revision=None, subdir="", **kwargs)`
+  (`cosmos_framework/data/vfm/processors/__init__.py:141`) already has two modes:
+  - `repository` set → download repo, then `build_processor(local_path, **kwargs)`.
+  - `repository` is `None` → `build_processor(*args, **kwargs)` (needs a
+    positional/`tokenizer_type` source).
+- `build_processor(tokenizer_type, ...)`
+  (`cosmos_framework/data/vfm/processors/__init__.py:100`) already loads from a
+  **local directory**: `if os.path.isdir(tokenizer_type): return
+  Qwen3VLProcessor(tokenizer_type, ...)`.
+- `OmniInference._create` (`cosmos_framework/inference/inference.py`, the
+  `else` / non-`MODULE` branch ~line 1047) has, in the same scope:
+  `checkpoint_path = setup_args.download_checkpoint()` (the already-downloaded
+  local dir) and the raw `model_dict = setup_args.load_model_config_dict()`,
+  immediately before `config = Cosmos3OmniConfig(model=model_dict)`.
+- The dict path to the processor node is
+  `model_dict["config"]["vlm_config"]["tokenizer"]` (verified against
+  `Cosmos3-Super.yaml` structure: `model.config.vlm_config.tokenizer`).
+- The active checkpoint's registry entry is the `CheckpointConfig` defined in
+  `cosmos_framework/inference/common/args.py:370` (fields: `model_memory_bytes`,
+  `config_file`, `s3_uri`, `hf`, methods `download()`, `pretrained_kwargs()`).
+  The registry dict is `OmniSetupArgs.CHECKPOINTS` (`= _CHECKPOINTS` in
+  `args.py`), keyed by checkpoint name.
+
+## Design
+
+### 1. Registry flag
+
+Add one field to `CheckpointConfig` (`cosmos_framework/inference/common/args.py`):
+
+```python
+vlm_processor_from_checkpoint: bool = False
+"""When True, load the VLM text/vision processor from the checkpoint's own
+bundled files (the local download directory) instead of the repository
+hardcoded in the model config's ``vlm_config.tokenizer`` node."""
+```
+
+Set `vlm_processor_from_checkpoint=True` on the two task-specialized entries in
+`_CHECKPOINTS` (`cosmos_framework/inference/args.py`):
+`Cosmos3-Super-Text2Image` and `Cosmos3-Super-Image2Video`.
+`Cosmos3-Nano` and `Cosmos3-Super` keep the default `False`.
+
+### 2. Loader injection
+
+In `OmniInference._create`, in the non-`MODULE` branch, after `checkpoint_path`
+and `model_dict` are obtained and **before** `Cosmos3OmniConfig(model=model_dict)`:
+
+```python
+ckpt_cfg = setup_args.CHECKPOINTS.get(setup_args.checkpoint_path)
+if ckpt_cfg is not None and ckpt_cfg.vlm_processor_from_checkpoint:
+    tok = model_dict["config"]["vlm_config"]["tokenizer"]
+    tok.pop("repository", None)
+    tok.pop("revision", None)
+    tok.pop("subdir", None)
+    tok["tokenizer_type"] = str(checkpoint_path)
+```
+
+This drives `build_processor_lazy` down its existing `repository is None` path →
+`build_processor(tokenizer_type=<checkpoint_dir>)` → `os.path.isdir` → 
+`Qwen3VLProcessor(<checkpoint_dir>)`. The processor loads from the
+already-downloaded checkpoint; **no additional download** occurs.
+
+Notes:
+- The rewrite is a no-op for checkpoints not in the registry
+  (`setup_args.checkpoint_path` is a custom path/URI) — `.get` returns `None`.
+- The rewrite mutates the raw `model_dict` (cleaner than mutating the
+  post-transform `Cosmos3OmniConfig.model`).
+- This branch is the one the task-specialized checkpoints take (they use a YAML
+  `config_file`, i.e. `config_file_type != MODULE`). The `MODULE` /
+  `load_model_from_checkpoint` branch is out of scope (base DCP experiments).
+
+### 3. No behavioral change to `build_processor_lazy` / `build_processor`
+
+Both already support the local-dir mode. Only one small robustness improvement:
+tighten the error when `build_processor_lazy` is called with neither
+`repository` nor a `tokenizer_type` source, so a future misconfiguration fails
+with a clear message instead of a `TypeError` on a missing positional arg.
+
+### 4. Behavior matrix
+
+| Checkpoint                 | `repository` in YAML | flag  | Processor source                          | Extra download |
+| -------------------------- | -------------------- | ----- | ----------------------------------------- | -------------- |
+| `Cosmos3-Nano`             | `nvidia/Cosmos3-Nano`| False | `nvidia/Cosmos3-Nano` (as today)          | as today       |
+| `Cosmos3-Super`            | `nvidia/Cosmos3-Super`| False| `nvidia/Cosmos3-Super` (as today)         | as today       |
+| `Cosmos3-Super-Text2Image` | (overridden away)    | True  | the loaded checkpoint dir                 | none           |
+| `Cosmos3-Super-Image2Video`| (overridden away)    | True  | the loaded checkpoint dir                 | none           |
+
+### 5. Safety / failure mode
+
+If a flagged checkpoint directory does not contain the processor files,
+`Qwen3VLProcessor(<dir>)` raises a clear file-not-found / load error at
+`set_up_tokenizers` time. The flag is therefore only set for self-contained
+checkpoints (the task HF repos qualify — verified file listings). Documented as
+a precondition of the flag.
+
+## Components & boundaries
+
+- **`CheckpointConfig` (common/args.py):** declares the capability flag. Pure
+  data; no behavior change.
+- **Checkpoint registry (`args.py`):** declares which checkpoints are
+  self-contained. Data only.
+- **`OmniInference._create` (inference.py):** the single place that translates
+  the flag + `checkpoint_path` into a processor-config rewrite. The only new
+  control flow.
+- **`build_processor_lazy` / `build_processor`:** unchanged dispatch; consume a
+  local dir. (One error-message hardening.)
+
+## Testing
+
+1. **Unit — config rewrite.** Construct `model_dict` (or a minimal stand-in with
+   the `config.vlm_config.tokenizer` node) and invoke the rewrite logic for:
+   - a flagged checkpoint → assert the tokenizer node has `tokenizer_type ==
+     str(checkpoint_path)` and no `repository`/`revision`/`subdir`;
+   - an unflagged checkpoint → assert the node is untouched (still has
+     `repository`).
+   If the rewrite is extracted into a small helper, test the helper directly.
+2. **Unit — `build_processor` local dir.** Given a temp dir with the minimal
+   processor files, assert `build_processor(<dir>)` returns a `Qwen3VLProcessor`
+   (mirrors existing local-artifact behavior; may already be covered).
+3. **End-to-end smoke (manual / existing).** The already-validated t2i run
+   (`Cosmos3-Super-Text2Image`, 8 GPUs) completes and the logs show **no**
+   `hf download ... nvidia/Cosmos3-Super ... --include '*'` line for the
+   processor.
+
+## Out of scope
+
+- The `MODULE` / `load_model_from_checkpoint` (base DCP experiment) path.
+- Changing the base `Cosmos3-Nano` / `Cosmos3-Super` processor source.
+- The separate i2v sampling host-RAM OOM issue (tracked elsewhere).
+- Subdir selection within the checkpoint (root layout is sufficient and matches
+  the current `subdir: ""` behavior).