[BUG] Cosmos3-Edge reasoner-only inference initialises the 2.8 GB video VAE

[chrisvoncsefalvay]

Bug Description

A model_mode: reasoner inference run for nvidia/Cosmos3-Edge initialises the unified video tokenizer and resolves Wan2.2_VAE.pth, although the run only produces reasoner text.

On a cold cache this downloads a 2,818,839,170-byte auxiliary object before the reasoner probe can run. It also makes uvx a launch-time dependency for resolving the asset. The Edge checkpoint itself is about 8.6 GiB, so this hidden dependency materially changes image size, offline deployment and cold-start expectations.

Reproduction Steps

git clone https://github.com/NVIDIA/cosmos-framework.git
cd cosmos-framework
git checkout 058c8c046877f087e07b3472a76cac214f3ad39c
uv sync --all-extras --group=cu130-train

hf download nvidia/Cosmos3-Edge \
  --revision 6f58f6b4c91288838e60b6bcb2cc45d997e961de \
  --local-dir ./Cosmos3-Edge

cat > /tmp/edge-reasoner.json <<'JSON'
{
  "model_mode": "reasoner",
  "prompt": "Describe the image.",
  "vision_path": "https://huggingface.co/nvidia/Cosmos3-Edge/resolve/6f58f6b4c91288838e60b6bcb2cc45d997e961de/assets/example_reasoning_input.png"
}
JSON

env LD_LIBRARY_PATH= .venv/bin/python -m cosmos_framework.scripts.inference \
  --parallelism-preset=latency \
  -i /tmp/edge-reasoner.json \
  -o /tmp/edge-reasoner-output \
  --checkpoint-path ./Cosmos3-Edge \
  --config-file cosmos_framework/inference/configs/model/Cosmos3-Edge.yaml \
  --seed 0 \
  --no-guardrails \
  --no-use-torch-compile \
  --no-use-cuda-graphs \
  --device-memory-utilization 0.5 \
  --max-new-tokens 1

Run with the Wan asset absent from the Hugging Face/Framework cache and observe asset resolution during initialisation.

Reproducibility:

• Always
• Intermittently
• Only once

Expected vs. Actual Behavior

	Description
Expected	Reasoner-only mode initialises the text/vision understanding path without downloading or loading a video-generation VAE.
Actual	OmniMoTModel.__init__ unconditionally instantiates self.config.tokenizer; the Edge YAML points this at Wan2.2_VAE.pth, so reasoner-only startup resolves the 2.8 GB video VAE.

Outputs

Resolved auxiliary object

repository: Wan-AI/Wan2.2-TI2V-5B
revision: 921dbaf3f1674a56f47e83fb80a34bac8a8f203e
filename: Wan2.2_VAE.pth
size: 2818839170 bytes
sha256: 20eb789667fa5e60e7516bf509512f6cb61f01b0aa0695eadaea930c13892b36

The first measured tokenizer setup took 27.82 seconds including this fetch. With the object cached, tokenizer setup took 3.58 seconds.

System Information

Field	Value
Environment	uv sync --all-extras --group=cu130-train
Hardware	NVIDIA GB10, single GPU
OS	Linux aarch64, kernel 6.17.0-1014-nvidia
GPU Driver	580.142
CUDA Version	13.0
Python Version	3.13.11
Package Version / Commit	cosmos-framework 058c8c046877f087e07b3472a76cac214f3ad39c; package 1.2.2; Torch 2.10.0+cu130

Additional Context

The current deployment workaround is to pin and prefetch the exact Wan object and ensure uvx is on the non-interactive launch PATH. A reasoner-only initialisation path could avoid the download/load while retaining the existing tokenizer for generative modes.

[lfengad]

Thx for the feedback! Yeah, currently the inferences in the framework are all using the shared model construction. So even reasoner only case will load the whole tensors.
We have some guide for this in: https://github.com/NVIDIA/cosmos/blob/main/cookbooks/cosmos3/reasoner/README.md#choosing-a-backend-memory-footprint, which could be referred. Also, we could consider the support of selective loading in the framework inference init phase. THX!

[chrisvoncsefalvay]

@lfengad I thought I might as well drop my solution here too if it helps, #132 . Thanks for putting up with my complaining :)

[lfengad]

@chrisvoncsefalvay Yeah, thank you very much for the contribution! That should be helpful. Thx!