convert_model_to_diffusers: export via official diffusers, drop diffusers-cosmos3 shim#98
Merged
Conversation
…sers-cosmos3 shim diffusers >= 0.39.0 ships Cosmos3OmniPipeline / Cosmos3OmniTransformer / Cosmos3AVAEAudioTokenizer natively, so the pre-upstream diffusers_cosmos3 shim is obsolete. Its own denoising __call__ also had a compounding per-step numeric drift that produced soft t2i images, so nothing should run inference through it. - Rewrite scripts/_convert_model_to_diffusers.py against the official classes, adopting the key remap from upstream scripts/convert_cosmos3_to_diffusers.py (strip the language model's `model.` nesting; q_proj/*_moe_gen -> to_q/add_q_proj/... attention names; vae2llm/llm2vae -> proj_in/proj_out; time_embedder mlp.0/2 -> linear_1/2; sound2llm/llm2sound -> audio_proj_in/out). This also fixes the exporter's stale native-layout key handling, which could not strict-load into the current flat-layout model classes. - Keep the repo-specific extensions upstream lacks: action projection export (action2llm/llm2action/action_modality_embed -> action_proj_in/out) and the vision_encoder/ sidecar save for the transformers/vLLM consumers of the exported repository. - Build sound_tokenizer/ as a real Cosmos3AVAEAudioTokenizer pipeline component instead of hand-writing safetensors + patching model_index.json; drop the now-unconditional remap_* CLI flags. - Construct the pipeline with enable_safety_checker=False (the converter env does not ship cosmos_guardrail); consumers can opt back in at load time. - Remove packages/diffusers-cosmos3, bump diffusers to >=0.39.0 (safetensors follows to 0.8.0), and drop the now-fixed diffusers audit ignores. - Add a key-remap test that round-trips exporter keys through the inference loader's inverse mapping; the full remap was additionally verified against the real diffusers 0.39.0 Cosmos3OmniTransformer state-dict key set (action+sound enabled, exact match). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ckpoints Three issues surfaced by an end-to-end Cosmos3-Nano export run: - The wrapper resolved the AVAE sound tokenizer via `.hf.download()`, which bypasses the CheckpointConfig post_download hook that synthesizes the legacy .ckpt — and the HF-published sound_tokenizer/ is already in the diffusers layout anyway. Consume config.json + diffusion_pytorch_model.safetensors directly (the converter passes that layout through unchanged), falling back to the legacy file pair named by the model config's avae_path. - Vision-encoder weights in diffusers-layout source checkpoints live under vision_encoder/ with bare Qwen3VLVisionModel keys, not under the native `model.visual.` prefix; add that fallback to _load_vision_encoder. - language_model.config is a nested VL config (Qwen3VLConfig) whose text-model fields live on text_config; resolve via get_text_config(). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The OSS CI only runs the curated GPU suites under tests/ and does not collect package-internal *_test.py files, so this test never runs here. The remap was already verified against the real diffusers 0.39.0 Cosmos3OmniTransformer state-dict key set and by the end-to-end Cosmos3-Nano export + inference run. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…kip-vision-encoder Surfaced by exporting a vision-generation-only DPO DCP checkpoint (no Qwen3-VL visual weights): vision_gen gates the generation branch, so such checkpoints still set it, and the vision-encoder load then fails with an unhelpful error unless --skip-vision-encoder is passed. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
lfengad
approved these changes
Jul 9, 2026
Xuanmeng-Zhang
approved these changes
Jul 9, 2026
pengcuo
approved these changes
Jul 9, 2026
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
4 participants
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.
What
Replace the
packages/diffusers-cosmos3shim with official diffusers (≥ 0.39.0) in the checkpoint exporter, and delete the package.Why
Cosmos3OmniPipeline,Cosmos3OmniTransformer,Cosmos3AVAEAudioTokenizer); the shim was a pre-upstream stand-in.model.-nestedq_proj/vae2llmvs flatto_q/proj_in) and could not strict-load.Changes
scripts/_convert_model_to_diffusers.pyagainst the official classes; keep the repo-specific extensions (action-projection export,vision_encoder/sidecar for the transformers/vLLM consumers).sound_tokenizer/as a realCosmos3AVAEAudioTokenizerpipeline component instead of hand-writing safetensors + patchingmodel_index.json.sound_tokenizer/directly, load vision weights fromvision_encoder/-subfolder sources, resolve the text config viaget_text_config()for nestedQwen3VLConfig.Relation to upstream
scripts/convert_cosmos3_to_diffusers.pyUpstream diffusers ships a converter for the same task, and this rewrite adopts its verified key-remap table. It cannot be used directly here:
projects.cosmos3.vfm.*), so it does not run against this repo;vision_encoder/sidecar, no top-level vLLM config / unified safetensors index).Verification
Cosmos3OmniTransformerstate-dict key set (action + sound enabled), and inverts cleanly through the native inference loader's mapping.Cosmos3OmniPipelinet2i sharp; native inference re-read of the export sharp.net.*/net_ema.*keys, t2i-only) converts and generates sharp t2i through the official pipeline (driven via the inner exporter; the CLI wrapper assertsaction_gen/sound_genand targets full omni checkpoints).Compat notes
model_index: Cosmos3OmniDiffusersPipeline) still load via an explicitCosmos3OmniPipeline.from_pretrained(...).cosmos_guardrailinstalled requiresenable_safety_checker=False(matches upstream semantics).