[Bugfix] Add Cosmos3-Edge preset (fix wrong Nano/Super defaults in T2V/I2V examples) (#5596)

Signed-off-by: lishunyang12 <lishunyang03@gmail.com>
Co-authored-by: lishunyang12 <lishunyang03@gmail.com>

examples/offline_inference/image_to_video/image_to_video.py

@@ -334,6 +334,7 @@ def main():
     is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
     is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
     is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
+    is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower)
 
     image = PIL.Image.open(args.image).convert("RGB") if args.image else None
     last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
@@ -347,7 +348,19 @@ def main():
 
     # Per-model generation defaults, applied only when the matching flag is omitted.
     # Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
-    if is_cosmos:
+    if is_cosmos_edge:
+        # Cosmos3-Edge native defaults: 480x832, gs 5.0, flow_shift 3.0 — NOT the Nano/Super
+        # 720p / gs 6.0 / flow_shift 10.0 below, which produce degenerate output on Edge.
+        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+            24,
+            5.0,
+            189,
+            35,
+            3.0,
+            480 * 832,
+            16,
+        )
+    elif is_cosmos:
         d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
             24,
             6.0,

examples/offline_inference/text_to_video/text_to_video.py

@@ -56,6 +56,19 @@ _MODEL_PRESETS = {
         "flow_shift": 10.0,
         "output": "cosmos3_t2v_output.mp4",
     },
+    # Cosmos3-Edge is a compact checkpoint with its own native resolution / guidance /
+    # flow_shift. It must NOT inherit the Nano/Super "cosmos" numbers above (720p / gs 6.0 /
+    # flow_shift 10.0), which produce degenerate output on Edge.
+    "cosmos3_edge": {
+        "height": 480,
+        "width": 832,
+        "num_frames": 189,
+        "num_inference_steps": 35,
+        "guidance_scale": 5.0,
+        "fps": 24,
+        "flow_shift": 3.0,
+        "output": "cosmos3_edge_t2v_output.mp4",
+    },
     "helios": {
         "height": 384,
         "width": 640,
@@ -103,6 +116,10 @@ def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
         return _MODEL_PRESETS["ltx2"]
     if "vace" in model_lower:
         return _MODEL_PRESETS["vace"]
+    # Edge must be matched before the generic cosmos branch (its "cosmos" substring would
+    # otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
+    if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
+        return _MODEL_PRESETS["cosmos3_edge"]
     if "cosmos" in model_lower:
         return _MODEL_PRESETS["cosmos"]
     if "hunyuan" in model_lower:

recipes/cosmos3/Cosmos3-Edge.md

@@ -17,7 +17,9 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared
 
 - **Transformer**: `Cosmos3EdgeVFMTransformer` (Nemotron-based), auto-selected from
   `backbone_type: cosmos3_edge_nemotron_dense` — nothing to pass.
-- **Defaults** differ; omit `--height` / `--width` to get them:
+- **Defaults** differ. The offline examples auto-detect Edge (from `edge` in the model id /
+  transformer class) and apply the Edge column below, so you can omit `--height` / `--width` /
+  `--guidance-scale` / `--flow-shift` and still get native Edge output:
 
 | | Edge | Nano / Super |
 |---|---|---|
@@ -26,6 +28,10 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared
 | Video `guidance_scale` | 5.0 | 6.0 |
 | Video `flow_shift` | 3.0 | 10.0 |
 
+> **Important:** these auto-defaults live in the example CLIs
+> (`text_to_video.py` / `image_to_video.py`). Do **not** pass the Nano/Super numbers
+> (720p / `guidance_scale` 6.0 / `flow_shift` 10.0) to Edge — they produce degenerate output.
+
 For the larger checkpoints see [Cosmos3-Nano](./Cosmos3-Nano.md) and [Cosmos3-Super](./Cosmos3-Super.md).
 
 ## Support notes