Export Cosmos3 Edge scheduler metadata

cosmos_framework/scripts/_convert_model_to_diffusers.py

@@ -27,6 +27,14 @@
 except ImportError:
     # Some intermediate main-branch revisions used this renamed export.
     from diffusers import Cosmos3OmniDiffusersPipeline as Cosmos3OmniPipeline
+
+try:
+    from diffusers import Cosmos3EdgeUniPCMultistepScheduler
+except ImportError:
+    # Older Diffusers builds predate the Edge scheduler. Keep the import optional
+    # so non-Edge conversions still work; Edge conversions raise a clear error at
+    # use time (see the is_edge_model scheduler branch below).
+    Cosmos3EdgeUniPCMultistepScheduler = None
 from transformers import AutoConfig, AutoTokenizer
 
 from cosmos_framework.inference.model import Cosmos3OmniModel
@@ -1151,10 +1159,19 @@ def _write_edge_transformer_config(output_dir: pathlib.Path, model_cfg: Any) ->
 def _normalize_edge_model_index(output_dir: pathlib.Path) -> None:
     """Keep Edge pipeline metadata consistent with the existing Hub export."""
     model_index_path = output_dir / "model_index.json"
+    scheduler_config_path = output_dir / "scheduler" / "scheduler_config.json"
     model_index = _load_json(model_index_path)
+    scheduler_config = _load_json(scheduler_config_path)
+    try:
+        native_flow_shift = float(scheduler_config["flow_shift"])
+    except (KeyError, TypeError, ValueError) as error:
+        raise ValueError(
+            f"Cosmos3 Edge scheduler config must define a numeric flow_shift: {scheduler_config_path}"
+        ) from error
     model_index.update(
         {
             "default_use_system_prompt": False,
+            "native_flow_shift": native_flow_shift,
             "text_tokenizer": ["transformers", "PreTrainedTokenizerFast"],
             "use_native_flow_schedule": True,
         }
@@ -1592,6 +1609,19 @@ def convert_model_to_diffusers(args: Args) -> None:
                 },
                 fixed_step_requires_explicit_sigmas=True,
             )
+        elif is_edge_model:
+            if Cosmos3EdgeUniPCMultistepScheduler is None:
+                raise ImportError(
+                    "Converting a Cosmos3 Edge checkpoint requires a Diffusers build that exports "
+                    "Cosmos3EdgeUniPCMultistepScheduler; the installed build does not. Upgrade Diffusers "
+                    "to a revision that ships the Edge scheduler."
+                )
+            scheduler = Cosmos3EdgeUniPCMultistepScheduler(
+                use_karras_sigmas=False,
+                use_flow_sigmas=True,
+                prediction_type="flow_prediction",
+                flow_shift=3.0,
+            )
         else:
             # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps.
             # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281.

cosmos_framework/scripts/convert_model_to_diffusers.py

@@ -185,6 +185,13 @@ def _load_edge_policy_metadata(checkpoint_path: Path) -> EdgePolicyMetadata:
     True: ("Cosmos3DistilledModularPipeline", "Cosmos3DistilledBlocks"),
 }
 
+PIPELINE_BEHAVIOR_FIELDS = (
+    "default_use_system_prompt",
+    "enable_safety_checker",
+    "use_native_flow_schedule",
+    "native_flow_shift",
+)
+
 
 def _write_modular_model_index(output_path: Path, repo_id: str) -> None:
     """Write modular_model_index.json next to the task-based model_index.json.
@@ -228,6 +235,10 @@ def _write_modular_model_index(output_path: Path, repo_id: str) -> None:
         # consumers (e.g. vllm-omni).
         modular_index["distilled_sigmas"] = distilled_sigmas
 
+    for name in PIPELINE_BEHAVIOR_FIELDS:
+        if name in model_index:
+            modular_index[name] = model_index[name]
+
     for name, value in model_index.items():
         # Keep only saved components: [library, class] pairs with non-null entries.
         # Skips meta keys (_class_name, ...) and unset components ([null, null]).