Skip to content

Export Cosmos3 Edge scheduler metadata#130

Merged
lfengad merged 2 commits into
mainfrom
liangf/edge-scheduler-metadata
Jul 23, 2026
Merged

Export Cosmos3 Edge scheduler metadata#130
lfengad merged 2 commits into
mainfrom
liangf/edge-scheduler-metadata

Conversation

@lfengad

@lfengad lfengad commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

Port of imaginaire4 MR !10668 into the OSS cosmos-framework repo.

What

Exports Cosmos3 Edge scheduler metadata into both model_index.json and modular_model_index.json during Diffusers conversion.

  • _convert_model_to_diffusers.py
    • Add an elif is_edge_model branch that instantiates Cosmos3EdgeUniPCMultistepScheduler(flow_shift=3.0). Previously Edge checkpoints fell into the generic Karras UniPCMultistepScheduler branch.
    • _normalize_edge_model_index now reads flow_shift back from the already-written scheduler/scheduler_config.json to populate native_flow_shift, so the exported value cannot drift from the scheduler constructor argument (addresses the reviewer P2 comment on the original MR).
  • convert_model_to_diffusers.py
    • Add PIPELINE_BEHAVIOR_FIELDS and propagate those metadata fields from model_index.json into modular_model_index.json, gated on key presence so pipelines lacking them are unaffected.

Notes

  • Requires a Diffusers build that exports Cosmos3EdgeUniPCMultistepScheduler. The import is ungated at module top-level, consistent with the existing Cosmos3OmniTransformer import — the converter already requires a bleeding-edge Diffusers build.
  • Diff matches the upstream MR (+28 lines); path is cosmos_framework/scripts/ here vs packages/cosmos3/cosmos3/scripts/ upstream.

🤖 Generated with Claude Code

lfengad and others added 2 commits July 22, 2026 22:58
Port of imaginaire4 MR !10668.

- Instantiate Cosmos3EdgeUniPCMultistepScheduler for Edge models so the
  exported scheduler matches inference (previously Edge fell into the generic
  Karras UniPC branch).
- Populate model_index.json native_flow_shift by reading flow_shift back from
  the saved scheduler/scheduler_config.json rather than hardcoding it, so the
  two values cannot drift.
- Propagate pipeline-behavior fields (default_use_system_prompt,
  enable_safety_checker, use_native_flow_schedule, native_flow_shift) from
  model_index.json into modular_model_index.json, gated on key presence.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move the Cosmos3EdgeUniPCMultistepScheduler import into a try/except (mirroring
the existing Cosmos3OmniPipeline guard) so older Diffusers builds that predate
the Edge scheduler can still run non-Edge conversions. Edge conversions on such
builds now raise a clear, actionable ImportError at use time instead of failing
at module load.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@lfengad
lfengad enabled auto-merge (squash) July 23, 2026 12:05
@lfengad
lfengad merged commit f734253 into main Jul 23, 2026
8 checks passed
@lfengad
lfengad deleted the liangf/edge-scheduler-metadata branch July 23, 2026 12:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants