[new-model] Add Cosmos Predict2.5 2B distilled T2W inference - #1767
Open
Mister-Raggs wants to merge 16 commits into
Open
[new-model] Add Cosmos Predict2.5 2B distilled T2W inference#1767Mister-Raggs wants to merge 16 commits into
Mister-Raggs wants to merge 16 commits into
Conversation
Cosmos-Predict2.5 2B text encoding crashes on load because
`apply_chat_template(tokenize=True)` returns a transformers `BatchEncoding`
-- a `collections.UserDict`, i.e. a `Mapping` but NOT a `dict` -- which the
`isinstance(tokenizer_output, dict)` check missed, so the id extraction fell
through to `raise RuntimeError("Unexpected chat_template output type")`.
Match on `Mapping` instead (every `dict` is a `Mapping`, so existing
plain-dict and list outputs are unchanged) and lift the normalization into a
tested `_normalize_chat_template_ids` helper. Add weight-free parametrized
regression tests over every shape apply_chat_template can return.
Address review: the helper documents a flat list[int] contract, but a nested list with batch>1 slipped through the len==1 unwrap guard and was returned as-is, violating the contract and breaking downstream padding. Raise on batch>1 instead of silently returning a nested list; add tests.
mergify
Bot
added
type: new-model
Contributor
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI
🔴 PR merge requirementsWaiting for
This rule is failing.
|
6 tasks
Mister-Raggs
changed the title
Mister-Raggs
marked this pull request as ready for review
August 27, 2026 20:54
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Add initial FastVideo inference support for NVIDIA's released Cosmos Predict2.5 2B distilled Text2World student. The port follows the student's official fixed-noise four-step TrigFlow rollout rather than reusing the full model's UniPC path.
This is a draft because a public converted FastVideo package ID is still pending. It also depends on #1607; this branch currently carries the same two Reason1 commits so the end-to-end validation is reproducible, and they can be dropped once #1607 lands.
Changes
Cosmos25DistilledSchedulerwith the official preconditioning, x0 conversion, fixed initial noise, and four-step schedule.base/distilledstudent checkpoint while reusing the existing Cosmos Predict2.5 tokenizer, VAE, and Reason1 components.cosmos25_distilled_full.mp4
Test plan and results
Spark / GB10, BF16 with Torch SDPA:
0.000655, final relative mean error0.038397, final mean absolute error0.037414.256x448x9T2W smoke passed; 2.19 seconds after model load.(256, 448, 3).704x1280x77generation passed in 143.53 seconds after model load and passed visual inspection.Local static validation:
git diff --check: passed.codespell, PyMarkdown, and filename hooks: passed.torchvisionis not installed; the same suites and GPU gates above passed in the FastVideo environment on Spark.Scope and limitations
Checklist