diff --git a/cosmos_framework/callbacks/learning_rate_logger.py b/cosmos_framework/callbacks/learning_rate_logger.py
index 223d827f..657e3181 100644
--- a/cosmos_framework/callbacks/learning_rate_logger.py
+++ b/cosmos_framework/callbacks/learning_rate_logger.py
@@ -4,21 +4,21 @@
 import torch
 import wandb
 
+from cosmos_framework.utils import distributed
 from cosmos_framework.utils.callback import Callback
 
 
 class LearningRateLogger(Callback):
-    """Logs per-model-part learning rate every ``every_n × logging_iter`` steps.
+    """Log reasoner default and LR-multiplier groups.
 
-    Designed for VLM training where the optimizer is an
-    ``OptimizersContainer`` exposing ``.optimizers`` (list of single-element
-    optimizer lists) paired with ``.model_part_names``. Silently no-ops when
-    those attributes are absent so it can be registered alongside plain
-    ``torch.optim.Optimizer`` setups without harm.
+    Parameters that do not match a configured ``optimizer.lr_multipliers`` key
+    are logged as ``optim/lr_default``. Named multiplier groups, such as
+    ``model.visual``, are logged as ``optim/lr_model.visual`` every
+    ``every_n × logging_iter`` steps.
     """
 
-    def __init__(self, every_n: int = 10):
-        self.every_n = every_n
+    def __init__(self, every_n: int = 10) -> None:
+        self.every_n: int = every_n
 
     def on_before_optimizer_step(
         self,
@@ -32,16 +32,30 @@ def on_before_optimizer_step(
         gate = self.config.trainer.logging_iter * self.every_n
         if not (iteration == 1 or (gate > 0 and iteration % gate == 0)):
             return
-        if not wandb.run:
+        if not distributed.is_rank0() or not wandb.run:
             return
-        if not (hasattr(optimizer, "optimizers") and hasattr(optimizer, "model_part_names")):
+        if not (hasattr(optimizer, "optimizers") and hasattr(optimizer, "model")):
             return
-        unique_lr: dict[str, float] = {}
-        for optim_per_model, name in zip(optimizer.optimizers, optimizer.model_part_names):
-            if not optim_per_model:
-                continue
-            for pg in optim_per_model[0].param_groups:
-                unique_lr[f"optim/lr_{name}"] = pg["lr"]
+
+        lr_multiplier_keys = list(self.config.optimizer.get("lr_multipliers", {}))
+        optimizer_net = getattr(optimizer.model, "net", None)
+        if optimizer_net is None:
+            return
+
+        lr_key_by_param_id: dict[int, str] = {}
+        for param_name, param in optimizer_net.named_parameters():
+            lr_key_by_param_id[id(param)] = next(
+                (lr_key for lr_key in lr_multiplier_keys if lr_key in param_name),
+                "default",
+            )
+
+        unique_lr: dict[str, float | torch.Tensor] = {}
+        for inner_optimizer in optimizer.optimizers:
+            for param_group in inner_optimizer.param_groups:
+                for param in param_group["params"]:
+                    lr_key = lr_key_by_param_id.get(id(param))
+                    if lr_key is not None:
+                        unique_lr[f"optim/lr_{lr_key}"] = param_group["lr"]
         if not unique_lr:
             return
         wandb.log(unique_lr, step=iteration)
diff --git a/cosmos_framework/callbacks/norm_monitor.py b/cosmos_framework/callbacks/norm_monitor.py
index 36e6ea00..39a2ec27 100644
--- a/cosmos_framework/callbacks/norm_monitor.py
+++ b/cosmos_framework/callbacks/norm_monitor.py
@@ -229,6 +229,12 @@ def _compute_l2_stats(self, tensor: torch.Tensor, detach: bool = True) -> dict[s
         if isinstance(data, DTensor):
             data = data.to_local()
 
+        if data.numel() == 0:
+            return {
+                "sq_sum": torch.zeros((), device=data.device, dtype=torch.float32),  # []
+                "max": torch.zeros((), device=data.device, dtype=data.dtype),  # []
+            }
+
         return {
             "sq_sum": (data.float() ** 2).sum(),
             "max": data.abs().max(),
diff --git a/cosmos_framework/configs/base/reasoner/defaults/config.py b/cosmos_framework/configs/base/reasoner/defaults/config.py
index 6da4659a..f6c3e21c 100644
--- a/cosmos_framework/configs/base/reasoner/defaults/config.py
+++ b/cosmos_framework/configs/base/reasoner/defaults/config.py
@@ -52,6 +52,7 @@ class DataSetting:
     num_data_workers: int = 8
     data_prefetch_factor: int = 1
     val_split_ratio: float = 0.0
+    recipe_name: str | None = None
 
 
 @attrs.define(slots=False)
diff --git a/cosmos_framework/configs/base/reasoner/defaults/optimizer.py b/cosmos_framework/configs/base/reasoner/defaults/optimizer.py
index dbd5829c..e73066f6 100644
--- a/cosmos_framework/configs/base/reasoner/defaults/optimizer.py
+++ b/cosmos_framework/configs/base/reasoner/defaults/optimizer.py
@@ -20,7 +20,7 @@
     betas=(0.9, 0.95),
     fused=True,
     keys_to_select=[],
-    lr_multipliers={"vision_encoder": 0.1},
+    lr_multipliers={"model.visual": 0.1},
 )
 
 # ``f_start`` / ``f_min`` are ratios of the optimizer's base ``lr``:
@@ -32,7 +32,7 @@
     cycle_lengths=["${trainer.max_iter}"],
     f_start=[0.01],
     f_max=[1.0],
-    f_min=[0.5],
+    f_min=[0.1],
 )
 
 
diff --git a/cosmos_framework/data/generator/augmentors/image_editing_transform.py b/cosmos_framework/data/generator/augmentors/image_editing_transform.py
index 1571fd44..d0a513c1 100644
--- a/cosmos_framework/data/generator/augmentors/image_editing_transform.py
+++ b/cosmos_framework/data/generator/augmentors/image_editing_transform.py
@@ -375,13 +375,19 @@ class ImageEditingToTrainingFormat(Augmentor):
     (e.g. ``OmniInterleavedMediaResize``).  This augmentor only normalises the
     PIL images to tensors and assembles the remaining metadata fields.
 
+    Supports both single-source image editing and multi-reference generation:
+    ``source_image`` may be a single ``PIL.Image`` (one reference) or a
+    ``list[PIL.Image]`` (N references). The output ``images`` always places the
+    target last, so the resulting layout is ``[ref_1, ..., ref_N, target]`` with
+    ``num_frames = N + 1``.
+
     Input (from data_dict):
-        - source_image: PIL.Image (already resized by upstream augmentor)
+        - source_image: PIL.Image | list[PIL.Image] (already resized by upstream augmentor)
         - target_image: PIL.Image (already resized by upstream augmentor)
         - editing_instruction: str
 
     Output (added to data_dict):
-        - images: list[torch.Tensor]  — ``[source (C,H_s,W_s), target (C,H_t,W_t)]``
+        - images: list[torch.Tensor]  — ``[ref_1, ..., ref_N, target]`` (each C,H,W)
         - ai_caption: str
         - selected_caption_type: str
         - fps: float
@@ -415,21 +421,28 @@ def __call__(self, data_dict: dict) -> dict | None:
         Returns:
             Updated data_dict with training-compatible fields, or None on error.
         """
-        source_image: Image.Image = data_dict.get("source_image")
+        source_image = data_dict.get("source_image")
         target_image: Image.Image = data_dict.get("target_image")
         editing_instruction: str = data_dict.get("editing_instruction", "")
 
         if source_image is None or target_image is None:
             return None
 
+        # Support both single-source editing (one PIL image) and multi-reference
+        # generation (a list of PIL images). Normalise to a list of references.
+        source_images = [source_image] if isinstance(source_image, Image.Image) else list(source_image)
+        if not source_images:
+            return None
+
         try:
-            # Normalize PIL images to tensors (upstream augmentor already handled resizing)
-            source_tensor = self._normalize_image(source_image)  # [C,H_s,W_s]
-            target_tensor = self._normalize_image(target_image)  # [C,H_t,W_t]
+            # Normalize PIL images to tensors (upstream augmentor already handled resizing).
+            # Each image keeps its own spatial size; the model encodes them separately.
+            # The target is placed last: [ref_1, ..., ref_N, target].
+            images = [self._normalize_image(src) for src in source_images]  # each [C,H_s,W_s]
+            images.append(self._normalize_image(target_image))  # [C,H_t,W_t]
 
             # Store as list of tensors for the batch collation.
-            # Each image keeps its own spatial size; the model encodes them separately.
-            data_dict["images"] = [source_tensor, target_tensor]
+            data_dict["images"] = images
 
             # Set text fields
             data_dict["ai_caption"] = editing_instruction
@@ -437,16 +450,13 @@ def __call__(self, data_dict: dict) -> dict | None:
 
             # Set metadata
             data_dict["fps"] = 30.0  # Same as standard image training
-            data_dict["num_frames"] = 2  # Source + target = 2 frames
+            data_dict["num_frames"] = len(images)  # N references + target
             data_dict["image_size"] = [
                 torch.tensor(
-                    [source_image.height, source_image.width, source_image.height, source_image.width],
-                    dtype=torch.float,
-                ),  # [4]
-                torch.tensor(
-                    [target_image.height, target_image.width, target_image.height, target_image.width],
+                    [img.height, img.width, img.height, img.width],
                     dtype=torch.float,
-                ),  # [4]
+                )  # [4]
+                for img in (*source_images, target_image)
             ]
             # Set the dataset name if not already present
             if "dataset_name" not in data_dict:
diff --git a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py
index 8faa9721..24f61a07 100644
--- a/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py
+++ b/cosmos_framework/data/generator/augmentors/interleaved_image_transform.py
@@ -15,6 +15,8 @@
 from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor
 from cosmos_framework.data.imaginaire.webdataset.augmentors.image.misc import obtain_image_size
 
+Image.MAX_IMAGE_PIXELS = 933120000
+
 
 class ResizeToPaddingDivisor(Augmentor):
     """Resize images so that both width and height are multiples of padding_divisor."""
@@ -264,3 +266,76 @@ def _resize_image(
                 resized_image = resized_image.resize((final_width, final_height), Image.Resampling.LANCZOS)
 
         return resized_image
+
+
+class InterleavedMediaResizeByMaxPixels(Augmentor):
+    """Resize interleaved media by constraining total pixel area (max_pixels), preserving aspect ratio.
+
+    Unlike :class:`InterleavedMediaResize` (which caps the longest side), this augmentor scales
+    each image so its total pixel area does not exceed ``max_pixels`` while keeping the aspect
+    ratio, then aligns both dimensions down to multiples of ``padding_divisor`` (default 16).
+    Images are only scaled down, never up.
+
+    Mirrors the FLUX2 max-pixel resize behavior. Produces ``diffusion_media_list`` keyed the same
+    way as the input ``media_list`` dict, compatible with the downstream
+    ``ExtractMultiReferenceConversation`` / ``ExtractImageEditingConversation`` augmentors.
+
+    Args:
+        input_keys: List with the single key holding the media dict. Defaults to ``["media_list"]``.
+        max_pixels: Maximum total pixel area per image after resize.
+        padding_divisor: Dimension alignment divisor (both width and height become multiples of it).
+        args: Additional arguments passed to the parent class.
+    """
+
+    def __init__(
+        self,
+        input_keys: Optional[List] = None,
+        max_pixels: int = 1048576,
+        padding_divisor: int = 16,
+        args: Optional[dict] = None,
+    ) -> None:
+        input_keys = input_keys or ["media_list"]
+        super().__init__(input_keys, None, args)
+        self.max_pixels = max_pixels
+        self.padding_divisor = padding_divisor
+
+    def _compute_target_size(self, width: int, height: int) -> tuple[int, int]:
+        """Scale to fit within max_pixels, then align down to padding_divisor."""
+        total_pixels = width * height
+        if total_pixels > self.max_pixels:
+            scale = math.sqrt(self.max_pixels / total_pixels)
+            width = int(width * scale)
+            height = int(height * scale)
+
+        width = max(self.padding_divisor, (width // self.padding_divisor) * self.padding_divisor)
+        height = max(self.padding_divisor, (height // self.padding_divisor) * self.padding_divisor)
+        return width, height
+
+    def _resize_image(self, img: Image.Image) -> Image.Image:
+        w, h = img.size
+        target_w, target_h = self._compute_target_size(w, h)
+        if (target_w, target_h) != (w, h):
+            img = img.resize((target_w, target_h), Image.Resampling.LANCZOS)
+        return img
+
+    def _process_media(self, media) -> Image.Image | list:
+        if isinstance(media, list):
+            return [self._resize_image(frame) for frame in media]
+        return self._resize_image(media)
+
+    def __call__(self, data_dict: Dict) -> Optional[Dict]:
+        assert len(self.input_keys) == 1, (
+            "This transform only supports one input key. Try to organize all the media contents under one key."
+        )
+        media_key = self.input_keys[0]
+        media_list = data_dict.get(media_key)
+        if media_list is None:
+            print(f"Input key {media_key} not found in data_dict: {data_dict.get('__key__', 'unknown')}")
+            return None
+
+        diffusion_media_content = {}
+        for key, media in media_list.items():
+            diffusion_media_content[key] = self._process_media(media)
+
+        data_dict["diffusion_media_list"] = diffusion_media_content
+        return data_dict
diff --git a/cosmos_framework/data/generator/augmentors/multi_reference_transform.py b/cosmos_framework/data/generator/augmentors/multi_reference_transform.py
new file mode 100644
index 00000000..02d38e3f
--- /dev/null
+++ b/cosmos_framework/data/generator/augmentors/multi_reference_transform.py
@@ -0,0 +1,447 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Augmentors for multi-reference image generation.
+
+Multi-reference data layout (per sample):
+    annotations: dict with keys {instruction, raw_instruction, in_instruction,
+        input_images, output_image, editing_type, dataset_name, split}
+    images: dict mapping {input_01, input_02, ..., input_NN, output} -> raw JPEG bytes
+
+The ``instruction`` field may be either a single string (legacy format) or a
+dict of prompt variants ``{original, short, medium, detailed}`` (new format).
+For the dict form, one variant is sampled uniformly at random per sample.
+
+These augmentors transform that on-disk format into the same in-memory layout
+expected by ``ImageEditingToTrainingFormat`` (i.e. ``source_image`` is a list of
+PIL images, ``target_image`` is a single PIL image, ``editing_instruction`` is a
+string), so that the downstream image-editing augmentors can be reused unchanged.
+"""
+
+from __future__ import annotations
+
+import io
+import random
+import re
+from typing import Optional
+
+from PIL import Image, UnidentifiedImageError
+
+from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor
+from cosmos_framework.utils import log
+
+Image.MAX_IMAGE_PIXELS = 933120000
+
+_INPUT_KEY_RE = re.compile(r"^input_(\d+)$")
+_OUTPUT_KEY = "output"
+
+# The ``instruction`` annotation field may be either a single string (legacy
+# format) or a dict of prompt variants at different lengths (new format). When
+# it is a dict, one variant is sampled uniformly at random per sample. These are
+# the variant keys we sample over, in their canonical order.
+_PROMPT_VARIANT_KEYS = ("original", "short", "medium", "detailed")
+
+
+def _sorted_input_keys(keys: list[str]) -> list[str]:
+    """Return ``input_NN`` keys sorted by their numeric index."""
+    indexed: list[tuple[int, str]] = []
+    for key in keys:
+        match = _INPUT_KEY_RE.match(key)
+        if match is not None:
+            indexed.append((int(match.group(1)), key))
+    indexed.sort(key=lambda x: x[0])
+    return [key for _, key in indexed]
+
+
+class MultiReferencePKLToMedia(Augmentor):
+    """Decode the multi-reference image bundle into PIL images.
+
+    Reads ``data_dict[input_key]`` (a ``dict[str, bytes]`` whose keys are
+    ``input_01``..``input_NN`` plus ``output``) and writes the decoded PIL
+    images into ``data_dict[output_key]`` as ``dict[str, PIL.Image]`` with the
+    same keys.
+
+    Returns ``None`` (skip sample) if any image fails to decode or if the
+    ``output`` image is missing — both are required for training.
+    """
+
+    def __init__(
+        self,
+        input_keys: list | None = None,
+        input_key: str = "images",
+        output_key: str = "media_list",
+        args: Optional[dict] = None,
+    ) -> None:
+        super().__init__(input_keys or [], args=args)
+        self.input_key = input_key
+        self.output_key = output_key
+
+    def _bytes_to_pil(self, image_bytes: bytes, identifier: str) -> Image.Image | None:
+        try:
+            with io.BytesIO(image_bytes) as stream:
+                img = Image.open(stream)
+                img.load()
+                return img.convert("RGB")
+        except UnidentifiedImageError:
+            log.warning(
+                f"Skipping item '{identifier}': cannot identify image bytes.",
+                rank0_only=False,
+            )
+        except Exception as e:
+            log.warning(
+                f"Skipping item '{identifier}': error decoding image bytes: {e}",
+                rank0_only=False,
+            )
+        return None
+
+    def __call__(self, data_dict: dict) -> dict | None:
+        if self.input_key not in data_dict:
+            log.warning(
+                f"Input key '{self.input_key}' not found in data_dict (keys={list(data_dict.keys())})",
+                rank0_only=False,
+            )
+            return None
+
+        bundle = data_dict[self.input_key]
+        if not isinstance(bundle, dict):
+            log.warning(
+                f"Expected dict at data_dict['{self.input_key}'], got {type(bundle)}",
+                rank0_only=False,
+            )
+            return None
+
+        media: dict[str, Image.Image] = {}
+        for name, item in bundle.items():
+            if not isinstance(item, (bytes, bytearray)):
+                log.warning(
+                    f"Skipping item '{name}': expected bytes, got {type(item)}",
+                    rank0_only=False,
+                )
+                return None
+            decoded = self._bytes_to_pil(bytes(item), identifier=f"{self.input_key}['{name}']")
+            if decoded is None:
+                return None
+            media[name] = decoded
+
+        if _OUTPUT_KEY not in media:
+            log.warning(
+                f"Multi-reference sample missing '{_OUTPUT_KEY}' image: {data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+        if not _sorted_input_keys(list(media.keys())):
+            log.warning(
+                f"Multi-reference sample has no 'input_NN' images: {data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+
+        data_dict[self.output_key] = media
+        if self.input_key != self.output_key:
+            del data_dict[self.input_key]
+
+        return data_dict
+
+
+class ExtractMultiReferenceConversation(Augmentor):
+    """Extract source/target images and instruction from multi-reference annotation.
+
+    Expected ``data_dict`` state on entry:
+        annotations: dict with at least ``instruction`` and (implicitly via the
+            image bundle) ``input_NN``/``output`` images. ``in_instruction`` is
+            present but intentionally ignored — per pipeline design we keep all
+            input images in the order indicated by their ``input_NN`` key.
+        diffusion_media_list: dict[str, PIL.Image] keyed by the same
+            ``input_NN``/``output`` names, populated by
+            ``InterleavedMediaResize`` upstream.
+
+    The ``instruction`` field supports two formats:
+        - str: a single instruction (legacy format), used directly.
+        - dict: prompt variants keyed by ``original``/``short``/``medium``/
+          ``detailed`` (new format); one non-empty variant is sampled uniformly
+          at random per sample.
+
+    Output additions to ``data_dict``:
+        source_image: list[PIL.Image] in input-index order, truncated to
+            ``max_reference_images`` if necessary.
+        target_image: PIL.Image (the ``output`` image).
+        editing_instruction: str (the sampled ``instruction``, with
+            ``<img-1>, <img-2>, ...`` markers preserved).
+        dataset_name: ``f"{annotations.dataset_name}/{annotations.split}"``
+            (used for logging only).
+
+    Returns ``None`` when required fields are missing.
+    """
+
+    def __init__(
+        self,
+        input_keys: list | None = None,
+        max_reference_images: int = 20,
+        annotation_key: str = "annotations",
+        media_key: str = "diffusion_media_list",
+        instruction_key: str = "instruction",
+        prompt_variant_keys: tuple[str, ...] | list[str] = _PROMPT_VARIANT_KEYS,
+        args: Optional[dict] = None,
+    ) -> None:
+        super().__init__(input_keys or [], args=args)
+        if max_reference_images <= 0:
+            raise ValueError(f"max_reference_images must be positive, got {max_reference_images}")
+        self.max_reference_images = max_reference_images
+        self.annotation_key = annotation_key
+        self.media_key = media_key
+        self.instruction_key = instruction_key
+        self.prompt_variant_keys = tuple(prompt_variant_keys)
+
+    def _resolve_instruction(self, annotation: dict) -> str | None:
+        """Resolve the instruction string, sampling a variant when given a dict.
+
+        Returns a non-empty, stripped instruction string, or ``None`` when no
+        usable instruction is present.
+        """
+        raw = annotation.get(self.instruction_key)
+
+        if isinstance(raw, str):
+            return raw.strip() or None
+
+        if isinstance(raw, dict):
+            # Prefer the canonical variant keys (in order), then fall back to any
+            # other string values so we never silently drop a usable prompt.
+            candidates = [
+                raw[key].strip()
+                for key in self.prompt_variant_keys
+                if isinstance(raw.get(key), str) and raw[key].strip()
+            ]
+            if not candidates:
+                candidates = [value.strip() for value in raw.values() if isinstance(value, str) and value.strip()]
+            if candidates:
+                return random.choice(candidates)
+
+        return None
+
+    def __call__(self, data_dict: dict) -> dict | None:
+        for required_key in (self.annotation_key, self.media_key):
+            if required_key not in data_dict:
+                log.warning(
+                    f"'{required_key}' not found in data_dict: {data_dict.get('__key__', 'unknown')}",
+                    rank0_only=False,
+                )
+                return None
+
+        annotation = data_dict[self.annotation_key]
+        if not isinstance(annotation, dict):
+            log.warning(
+                f"Expected dict for '{self.annotation_key}', got {type(annotation)}: "
+                f"{data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+
+        instruction = self._resolve_instruction(annotation)
+        if not instruction:
+            log.warning(
+                f"Missing/empty '{self.instruction_key}' in annotation: {data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+
+        media_dict = data_dict[self.media_key]
+        if not isinstance(media_dict, dict):
+            log.warning(
+                f"Expected dict for '{self.media_key}', got {type(media_dict)}: {data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+
+        target_image = media_dict.get(_OUTPUT_KEY)
+        if isinstance(target_image, list):
+            target_image = target_image[0] if target_image else None
+        if target_image is None:
+            log.warning(
+                f"Missing '{_OUTPUT_KEY}' image after resize: {data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+
+        ordered_input_keys = _sorted_input_keys(list(media_dict.keys()))
+        if not ordered_input_keys:
+            log.warning(
+                f"No 'input_NN' images after resize: {data_dict.get('__key__', 'unknown')}",
+                rank0_only=False,
+            )
+            return None
+
+        if len(ordered_input_keys) > self.max_reference_images:
+            ordered_input_keys = ordered_input_keys[: self.max_reference_images]
+
+        source_images: list[Image.Image] = []
+        for key in ordered_input_keys:
+            ref = media_dict[key]
+            if isinstance(ref, list):
+                ref = ref[0] if ref else None
+            if ref is None:
+                log.warning(
+                    f"Reference image '{key}' is None after resize: {data_dict.get('__key__', 'unknown')}",
+                    rank0_only=False,
+                )
+                return None
+            source_images.append(ref)
+
+        data_dict["source_image"] = source_images
+        data_dict["target_image"] = target_image
+        data_dict["editing_instruction"] = instruction
+
+        ds_name = annotation.get("dataset_name")
+        split = annotation.get("split")
+        if isinstance(ds_name, str) and ds_name:
+            data_dict["dataset_name"] = f"{ds_name}/{split}" if isinstance(split, str) and split else ds_name
+
+        return data_dict
+
+
+class RandomResizeReferenceImages(Augmentor):
+    """Randomly rescale each reference image (aspect ratio preserved); leave the target untouched.
+
+    Operates on the ``media_list`` dict produced by :class:`MultiReferencePKLToMedia`,
+    i.e. before :class:`InterleavedMediaResize`. With probability ``resize_prob`` (one
+    Bernoulli draw per sample), every key matching ``input_NN`` gets its own independently
+    sampled ratio in ``[min_resize_ratio, max_resize_ratio]`` and is resized via LANCZOS.
+    The ``output`` key (target image) is skipped so the target resolution is unchanged.
+
+    Padding-divisor alignment and the side-length cap are handled by the downstream
+    :class:`InterleavedMediaResize` stage, so this augmentor does not need to worry
+    about VAE divisibility.
+    """
+
+    def __init__(
+        self,
+        input_keys: list | None = None,
+        min_resize_ratio: float = 0.5,
+        max_resize_ratio: float = 1.5,
+        resize_prob: float = 0.0,
+        media_key: str = "media_list",
+        output_key_pattern: str = _OUTPUT_KEY,
+        args: Optional[dict] = None,
+    ) -> None:
+        if not 0.0 <= resize_prob <= 1.0:
+            raise ValueError(f"resize_prob must be in [0, 1], got {resize_prob}")
+        if min_resize_ratio <= 0:
+            raise ValueError(f"min_resize_ratio must be > 0, got {min_resize_ratio}")
+        if max_resize_ratio < min_resize_ratio:
+            raise ValueError(f"max_resize_ratio ({max_resize_ratio}) must be >= min_resize_ratio ({min_resize_ratio})")
+        super().__init__(input_keys or [media_key], args=args)
+        self.min_resize_ratio = float(min_resize_ratio)
+        self.max_resize_ratio = float(max_resize_ratio)
+        self.resize_prob = float(resize_prob)
+        self.media_key = media_key
+        self.output_key_pattern = output_key_pattern
+
+    def _resize_one(self, img: Image.Image) -> Image.Image:
+        ratio = random.uniform(self.min_resize_ratio, self.max_resize_ratio)
+        w, h = img.size
+        new_w = max(1, round(w * ratio))
+        new_h = max(1, round(h * ratio))
+        if (new_w, new_h) == (w, h):
+            return img
+        return img.resize((new_w, new_h), Image.LANCZOS)
+
+    def __call__(self, data_dict: dict) -> dict | None:
+        if self.resize_prob <= 0.0 or random.random() >= self.resize_prob:
+            return data_dict
+
+        media_list = data_dict.get(self.media_key)
+        if not isinstance(media_list, dict):
+            return data_dict
+
+        for key, media in media_list.items():
+            if key == self.output_key_pattern or _INPUT_KEY_RE.match(key) is None:
+                continue
+            if isinstance(media, list):
+                media_list[key] = [self._resize_one(frame) for frame in media]
+            elif isinstance(media, Image.Image):
+                media_list[key] = self._resize_one(media)
+
+        return data_dict
+
+
+_MARKER_RE: re.Pattern[str] = re.compile(r"<img-(\d+)>")
+
+
+class ReorderReferenceImages(Augmentor):
+    """Shuffle the reference list and rewrite ``<img-N>`` markers in the instruction.
+
+    Operates on the ``source_image`` list and ``editing_instruction`` string produced by
+    :class:`ExtractMultiReferenceConversation`. With probability ``shuffle_prob`` (one
+    Bernoulli draw per sample), shuffles ``source_image`` and consistently renumbers the
+    ``<img-N>`` markers inside ``editing_instruction`` so that the same physical image is
+    still referenced after the shuffle.
+
+    Shuffling is **only** applied when ``editing_instruction`` contains at least one
+    ``<img-N>`` marker. Some samples refer to references by natural-language phrases
+    (``"Image 1"``, ``"first image"``, etc.) that we cannot reliably renumber, so
+    shuffling those would silently desync the instruction from the image order. Such
+    samples are returned unchanged.
+
+    The marker rewrite is done in a single ``re.sub`` pass to avoid cascading rewrites
+    (e.g. ``<img-1>`` -> ``<img-2>`` -> ``<img-3>``). Markers whose index falls outside
+    ``[1, len(source_image)]`` are left untouched.
+    """
+
+    def __init__(
+        self,
+        input_keys: list | None = None,
+        shuffle_prob: float = 0.0,
+        source_key: str = "source_image",
+        instruction_key: str = "editing_instruction",
+        args: Optional[dict] = None,
+    ) -> None:
+        if not 0.0 <= shuffle_prob <= 1.0:
+            raise ValueError(f"shuffle_prob must be in [0, 1], got {shuffle_prob}")
+        super().__init__(input_keys or [], args=args)
+        self.shuffle_prob = float(shuffle_prob)
+        self.source_key = source_key
+        self.instruction_key = instruction_key
+
+    def __call__(self, data_dict: dict) -> dict | None:
+        if self.shuffle_prob <= 0.0 or random.random() >= self.shuffle_prob:
+            return data_dict
+
+        source_image = data_dict.get(self.source_key)
+        instruction = data_dict.get(self.instruction_key)
+        if not isinstance(source_image, list) or len(source_image) <= 1:
+            return data_dict
+        if not isinstance(instruction, str) or not instruction:
+            return data_dict
+
+        # Only shuffle when the instruction actually uses ``<img-N>`` markers.
+        # Otherwise the references are by natural-language phrases (e.g. "Image 1",
+        # "the first image") that we can't reliably renumber, so reordering would
+        # silently desync the instruction from the image order.
+        if _MARKER_RE.search(instruction) is None:
+            return data_dict
+
+        n = len(source_image)
+        # Try a few times to get a non-identity permutation; otherwise accept identity.
+        perm = list(range(n))
+        for _ in range(5):
+            candidate = random.sample(range(n), n)
+            if candidate != list(range(n)):
+                perm = candidate
+                break
+
+        if perm == list(range(n)):
+            return data_dict
+
+        inv = [0] * n
+        for new_idx, old_idx in enumerate(perm):
+            inv[old_idx] = new_idx
+
+        def _remap(match: re.Match) -> str:
+            old_idx = int(match.group(1)) - 1
+            if 0 <= old_idx < n:
+                return f"<img-{inv[old_idx] + 1}>"
+            return match.group(0)
+
+        data_dict[self.source_key] = [source_image[perm[i]] for i in range(n)]
+        data_dict[self.instruction_key] = _MARKER_RE.sub(_remap, instruction)
+
+        return data_dict
diff --git a/cosmos_framework/data/generator/augmentors/pkl_to_media.py b/cosmos_framework/data/generator/augmentors/pkl_to_media.py
index c66ea953..4b0fda79 100644
--- a/cosmos_framework/data/generator/augmentors/pkl_to_media.py
+++ b/cosmos_framework/data/generator/augmentors/pkl_to_media.py
@@ -18,6 +18,7 @@
 from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor
 from cosmos_framework.utils import log
 from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images
+from cosmos_framework.utils.generator.torchcodec_video import decode_frames_tchw_uint8, probe_video
 
 Image.MAX_IMAGE_PIXELS = 933120000
 _VIDEO_EXTENSIONS = "mp4 avi webm mov".split()
@@ -64,7 +65,7 @@ def _video_decoder_qwen_func(
         target_fps (float, optional): Target FPS. Defaults to 2.0.
         min_video_token_length (int, optional): Minimum token length. Defaults to 16.
         max_video_token_length (int, optional): Maximum token length. Defaults to 8192.
-        num_threads (int, optional): Number of threads for decord. Defaults to 0.
+        num_threads (int, optional): Number of video decoding threads. Defaults to 0.
         random_augmentation (bool, optional): Whether to randomize the FPS and max_video_token_length. Defaults to False.
         fps_random_range (list[float], optional): Random FPS range. Defaults to [10.0, 24.0].
         max_video_token_length_random_range (list[float], optional): Random max_video_token_length range. Defaults to [0.75, 1.25].
@@ -80,17 +81,14 @@ def _video_decoder_qwen_func(
     Returns:
         dict | None: Dictionary with video frames tensor and target FPS
     """
-    import decord
-
     # Check video extension
     extension = re.sub(r".*[.]", "", key)
     if extension.lower() not in _VIDEO_EXTENSIONS:
         return None
 
     # Read video
-    video_buffer = io.BytesIO(data)
-    video_reader = decord.VideoReader(video_buffer, num_threads=num_threads)
-    total_frames, video_fps = len(video_reader), video_reader.get_avg_fps()
+    metadata = probe_video(data, num_threads=num_threads)
+    total_frames, video_fps = metadata.num_frames, metadata.average_fps
 
     if start_frame is not None and end_frame is not None:
         total_frames = end_frame - start_frame
@@ -138,8 +136,7 @@ def _video_decoder_qwen_func(
         idx = torch.linspace(start_frame, end_frame - 1, nframes).round().long().tolist()  # [nframes]
     else:
         idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()  # [nframes]
-    video_frames = video_reader.get_batch(idx).asnumpy()
-    video_frames = torch.tensor(video_frames).permute(0, 3, 1, 2)  # [T,C,H,W]
+    video_frames, _ = decode_frames_tchw_uint8(data, idx, num_threads=num_threads)  # [T,C,H,W]
     sample_fps = nframes / max(total_frames, 1e-6) * video_fps
 
     # recompute max_pixels based on number of sampled frames
@@ -156,13 +153,9 @@ def _video_decoder_qwen_func(
         [resized_height, resized_width],
         interpolation=InterpolationMode.BICUBIC,
         antialias=True,
-    ).float()
+    ).float()  # [T,C,H,W]
     video_frames = video_frames.permute(1, 0, 2, 3)  # [C,T,H,W]
 
-    # Clean up
-    video_reader.seek(0)  # set video reader point back to 0 to clean up cache
-    del video_reader  # delete the reader to avoid memory leak
-
     return dict(videos=video_frames, fps=sample_fps)
 
 
diff --git a/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py
index 03a61f5e..dddbadc2 100644
--- a/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py
+++ b/cosmos_framework/data/generator/augmentors/reasoner/bytes_to_media.py
@@ -19,6 +19,7 @@
 from cosmos_framework.data.generator.reasoner.video_decoder_qwen import _video_decoder_qwen_func
 from cosmos_framework.data.generator.processors.qwen3vl_processor import Qwen3VLProcessor
 from cosmos_framework.utils.generator.video_preprocess import tensor_to_pil_images
+from cosmos_framework.utils.generator.torchcodec_video import probe_video
 
 
 class BytesToMedia(Augmentor):
@@ -104,21 +105,16 @@ def _probe_video_duration_seconds(
     ) -> float | None:
         """Probe the effective video duration in seconds used for proportional budget allocation."""
         try:
-            import decord
-
-            with io.BytesIO(video_bytes) as stream:
-                video_reader = decord.VideoReader(stream, num_threads=self.video_decoder_params["num_threads"])
-                frame_count = (
-                    max(end_frame - start_frame, 0)
-                    if start_frame is not None and end_frame is not None
-                    else len(video_reader)
-                )
-                video_fps = video_reader.get_avg_fps()
-                video_reader.seek(0)
-                del video_reader
-                if frame_count <= 0 or video_fps <= 0:
-                    return None
-                return frame_count / video_fps
+            metadata = probe_video(video_bytes, num_threads=self.video_decoder_params["num_threads"])
+            frame_count = (
+                max(end_frame - start_frame, 0)
+                if start_frame is not None and end_frame is not None
+                else metadata.num_frames
+            )
+            video_fps = metadata.average_fps
+            if frame_count <= 0 or video_fps <= 0:
+                return None
+            return frame_count / video_fps
         except Exception as e:
             log.warning(f"Could not probe video duration for '{identifier}': {e}")
             return None
diff --git a/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py b/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py
new file mode 100644
index 00000000..6ddebd71
--- /dev/null
+++ b/cosmos_framework/data/generator/augmentors/torchcodec_callers_test.py
@@ -0,0 +1,208 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Tests for Cosmos3 dataset call sites that use TorchCodec video helpers."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.L1, pytest.mark.CPU]
+
+
+def _import_or_skip(module_name: str) -> Any:
+    try:
+        return import_module(module_name)
+    except ModuleNotFoundError as exc:
+        pytest.skip(f"Optional dependency unavailable while importing {module_name}: {exc.name}")
+
+
+def test_pkl_qwen_decoder_uses_probe_and_decodes_selected_window(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _import_or_skip("cosmos_framework.data.generator.augmentors.pkl_to_media")
+    video_bytes = b"video-bytes"
+    calls: dict[str, object] = {}
+
+    def fake_probe_video(source: object, num_threads: int = 0, **_: object) -> SimpleNamespace:
+        calls["probe"] = (source, num_threads)
+        return SimpleNamespace(num_frames=10, average_fps=12.0, height=8, width=10)
+
+    def fake_decode_frames_tchw_uint8(
+        source: object,
+        indices: list[int],
+        num_threads: int = 0,
+        **_: object,
+    ) -> tuple[torch.Tensor, SimpleNamespace]:
+        calls["decode"] = (source, indices, num_threads)
+        frames = torch.arange(3 * 3 * 8 * 10, dtype=torch.uint8).reshape(3, 3, 8, 10)  # [T,C,H,W]
+        return frames, fake_probe_video(source, num_threads=num_threads)
+
+    monkeypatch.setattr(module, "probe_video", fake_probe_video)
+    monkeypatch.setattr(module, "decode_frames_tchw_uint8", fake_decode_frames_tchw_uint8)
+    monkeypatch.setattr(module, "smart_nframes", lambda *_args, **_kwargs: 3)
+    monkeypatch.setattr(module, "smart_resize", lambda height, width, **_kwargs: (height, width))
+
+    result = module._video_decoder_qwen_func(
+        key="clip.mp4",
+        data=video_bytes,
+        min_fps_thres=1,
+        max_fps_thres=60,
+        target_fps=3.0,
+        min_video_token_length=1,
+        max_video_token_length=1024,
+        num_threads=5,
+        start_frame=2,
+        end_frame=8,
+    )
+
+    assert calls["probe"] == (video_bytes, 5)
+    assert calls["decode"] == (video_bytes, [2, 4, 7], 5)
+    assert result is not None
+    assert tuple(result["videos"].shape) == (3, 3, 8, 10)
+    assert result["videos"].dtype == torch.float32
+    assert result["fps"] == pytest.approx(6.0)
+
+
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        "cosmos_framework.data.generator.augmentors.reasoner.bytes_to_media",
+        "cosmos_framework.utils.datasets.augmentors.bytes_to_media",
+    ],
+)
+def test_bytes_to_media_uses_probe_durations_for_token_budget(
+    monkeypatch: pytest.MonkeyPatch,
+    module_name: str,
+) -> None:
+    module = _import_or_skip(module_name)
+    seen_sources: list[tuple[object, int]] = []
+
+    def fake_probe_video(source: object, num_threads: int = 0, **_: object) -> SimpleNamespace:
+        seen_sources.append((source, num_threads))
+        frame_counts = {b"first": 80, b"second": 40}
+        return SimpleNamespace(num_frames=frame_counts[source], average_fps=20.0, height=16, width=16)
+
+    monkeypatch.setattr(module, "probe_video", fake_probe_video)
+    augmentor = module.BytesToMedia(
+        min_video_token_length=8,
+        max_video_token_length=80,
+        num_threads=3,
+        is_input_pickle_byptes=False,
+    )
+
+    durations = augmentor._get_video_durations(
+        {
+            "video_first.mp4": b"first",
+            "control_input_depth": b"second",
+            "image.jpg": b"image",
+        },
+        {},
+    )
+
+    assert durations == {"video_first.mp4": 4.0, "control_input_depth": 2.0}
+    assert seen_sources == [(b"first", 3), (b"second", 3)]
+    weighted_params = augmentor._get_decoder_params(
+        video_count=2,
+        video_duration=durations["video_first.mp4"],
+        total_video_duration=sum(durations.values()),
+    )
+    assert weighted_params["max_video_token_length"] == 53
+
+
+@pytest.mark.parametrize(
+    "module_name",
+    [
+        "cosmos_framework.data.generator.augmentors.reasoner.bytes_to_media",
+        "cosmos_framework.utils.datasets.augmentors.bytes_to_media",
+    ],
+)
+def test_bytes_to_media_probe_duration_respects_start_end_frame(
+    monkeypatch: pytest.MonkeyPatch,
+    module_name: str,
+) -> None:
+    module = _import_or_skip(module_name)
+    monkeypatch.setattr(
+        module,
+        "probe_video",
+        lambda *_args, **_kwargs: SimpleNamespace(num_frames=100, average_fps=25.0, height=16, width=16),
+    )
+    augmentor = module.BytesToMedia(num_threads=2, is_input_pickle_byptes=False)
+
+    duration = augmentor._probe_video_duration_seconds(
+        b"video",
+        identifier="media['video.mp4']",
+        start_frame=7,
+        end_frame=32,
+    )
+    empty_duration = augmentor._probe_video_duration_seconds(
+        b"video",
+        identifier="media['video.mp4']",
+        start_frame=7,
+        end_frame=7,
+    )
+
+    assert duration == pytest.approx(1.0)
+    assert empty_duration is None
+
+
+def test_multiview_extract_frames_resizes_torchcodec_frames(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _import_or_skip("cosmos_framework.data.generator.multiview.multiview_dataset")
+    calls: list[tuple[object, list[int]]] = []
+
+    def fake_decode_frames_tchw_uint8(
+        source: object,
+        indices: list[int],
+        **_: object,
+    ) -> tuple[torch.Tensor, SimpleNamespace]:
+        calls.append((source, indices))
+        frames = torch.arange(len(indices) * 3 * 4 * 6, dtype=torch.uint8).reshape(len(indices), 3, 4, 6)  # [T,C,H,W]
+        return frames, SimpleNamespace(average_fps=29.97)
+
+    monkeypatch.setattr(module, "decode_frames_tchw_uint8", fake_decode_frames_tchw_uint8)
+
+    frames, fps, original_hw = module.ExtractFramesAndCaptions._extract_frames(  # [T,C,H,W]
+        b"video",
+        [1, 3, 5],
+        (8, 12),
+    )
+
+    assert calls == [(b"video", [1, 3, 5])]
+    assert tuple(frames.shape) == (3, 3, 8, 12)
+    assert frames.dtype == torch.uint8
+    assert fps == pytest.approx(29.97)
+    assert original_hw == (4, 6)
+
+
+def test_sekai_frame_count_uses_probe_video(monkeypatch: pytest.MonkeyPatch) -> None:
+    module = _import_or_skip("projects.cosmos3.sil.omnidreams.datasets.sekai")
+    seen: list[bytes] = []
+
+    def fake_probe_video(source: bytes) -> SimpleNamespace:
+        seen.append(source)
+        return SimpleNamespace(num_frames=37, average_fps=30.0, height=16, width=16)
+
+    monkeypatch.setattr(module, "probe_video", fake_probe_video)
+
+    assert module._num_available_video_frames(b"sekai-video") == 37
+    assert seen == [b"sekai-video"]
+
+
+def test_sekai_fit_window_clamps_to_available_span() -> None:
+    module = _import_or_skip("projects.cosmos3.sil.omnidreams.datasets.sekai")
+
+    assert module._fit_window_to_available_video(
+        frame_start=20,
+        num_video_frames=8,
+        stride=2,
+        num_available_frames=30,
+    ) == (20, 5)
+    assert module._fit_window_to_available_video(
+        frame_start=100,
+        num_video_frames=8,
+        stride=3,
+        num_available_frames=10,
+    ) == (0, 4)
diff --git a/cosmos_framework/data/generator/augmentors/video_parsing.py b/cosmos_framework/data/generator/augmentors/video_parsing.py
index a443948e..fb5808a6 100644
--- a/cosmos_framework/data/generator/augmentors/video_parsing.py
+++ b/cosmos_framework/data/generator/augmentors/video_parsing.py
@@ -35,7 +35,7 @@ class VideoParsing(Augmentor):
     This augmentor is used to parse the video bytes and get the video frames.
     the return dict is back-compatible with old datasets, which video decoding happens in the decoder stage.
 
-    Now uses torchcodec instead of decord for video decoding, with optional audio extraction.
+    Uses TorchCodec for video decoding, with optional audio extraction.
     """
 
     def __init__(self, input_keys: list, output_keys: Optional[list] = None, args: Optional[dict] = None) -> None:
diff --git a/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py
index 706b4d4e..714268e2 100644
--- a/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py
+++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step.py
@@ -53,6 +53,8 @@ def __call__(
         num_steps: int | None = None,
         shift: float | None = None,
         seed: int | list[int] | None = None,
+        condition_reference: torch.Tensor | list[torch.Tensor] | None = None,
+        condition_mask: torch.Tensor | list[torch.Tensor] | None = None,
     ) -> torch.Tensor | list[torch.Tensor]:
         """Run the fixed-step sampling loop.
 
@@ -77,15 +79,37 @@ def __call__(
                 ``len(t_list) - 1`` when provided).
             shift: When set, derive the sigma schedule dynamically using the
                 flow-matching shift formula instead of ``self.t_list``.
+            condition_reference: Optional clean reference tensor(s) to preserve
+                where ``condition_mask`` is 1.
+            condition_mask: Optional mask tensor(s), same shape as ``noise``,
+                where 1 marks clean conditioning values and 0 marks generated
+                values.
 
         Returns:
             Denoised sample(s).  A single ``torch.Tensor`` when ``noise`` is a
             tensor, or a ``list[torch.Tensor]`` when ``noise`` is a list.
         """
+        assert (condition_reference is None) == (condition_mask is None), (
+            "condition_reference and condition_mask must be both set or both None"
+        )
         if isinstance(noise, list):
             device = noise[0].device
+            if seed is None:
+                seed = [None] * len(noise)
+            assert isinstance(seed, list), "seed must be a list when noise is a list"
+            if condition_reference is None:
+                condition_reference = [None] * len(noise)
+                condition_mask = [None] * len(noise)
+            else:
+                assert isinstance(condition_reference, list), "condition_reference must be a list when noise is a list"
+                assert isinstance(condition_mask, list), "condition_mask must be a list when noise is a list"
         else:
             device = noise.device
+            assert not isinstance(seed, list), "seed must not be a list when noise is a tensor"
+            assert not isinstance(condition_reference, list), (
+                "condition_reference must not be a list when noise is a tensor"
+            )
+            assert not isinstance(condition_mask, list), "condition_mask must not be a list when noise is a tensor"
 
         if shift is not None:
             assert num_steps is not None, "num_steps is required when shift is provided"
@@ -105,27 +129,43 @@ def __call__(
             timestep = torch.tensor(sigma_cur * self.num_train_timesteps, device=device)
             v_pred = velocity_fn(latent, timestep.reshape(1, 1))
 
-            def _sde_step(seed: int | None, latent: torch.Tensor, v_pred: torch.Tensor) -> torch.Tensor:
-                x0_pred = latent - sigma_cur * v_pred
+            def _sde_step(
+                seed: int | None,
+                latent: torch.Tensor,
+                v_pred: torch.Tensor,
+                condition_reference: torch.Tensor | None,
+                condition_mask: torch.Tensor | None,
+            ) -> torch.Tensor:
+                x0_pred = latent - sigma_cur * v_pred  # [...,D]
 
                 if sigma_next > 0:
                     if self.sample_type == "ode":
                         # Euler ODE step
-                        latent = latent + (sigma_next - sigma_cur) * v_pred
+                        latent_next = latent + (sigma_next - sigma_cur) * v_pred  # [...,D]
                     else:
                         if seed is not None:
                             torch.manual_seed(seed + step_idx)
-                        eps_fresh = torch.randn_like(x0_pred)
-                        latent = (1.0 - sigma_next) * x0_pred + sigma_next * eps_fresh
+                        eps_fresh = torch.randn_like(x0_pred)  # [...,D]
+                        latent_next = (1.0 - sigma_next) * x0_pred + sigma_next * eps_fresh  # [...,D]
                 else:
-                    latent = x0_pred
-                return latent
+                    latent_next = x0_pred  # [...,D]
+
+                if condition_reference is not None:
+                    assert condition_mask is not None, "condition_mask is required when condition_reference is set"
+                    condition_reference = condition_reference.to(
+                        dtype=latent_next.dtype, device=latent_next.device
+                    )  # [...,D]
+                    condition_mask = condition_mask.to(dtype=latent_next.dtype, device=latent_next.device)  # [...,D]
+                    latent_next = condition_mask * condition_reference + (1.0 - condition_mask) * latent_next  # [...,D]
+                return latent_next
 
             latent = run_multiseed(
                 _sde_step,
                 seed=seed,
                 latent=latent,
                 v_pred=v_pred,
+                condition_reference=condition_reference,
+                condition_mask=condition_mask,
             )
 
         return latent
diff --git a/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py
index dd214018..a6aff60e 100644
--- a/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py
+++ b/cosmos_framework/model/generator/diffusion/samplers/fixed_step_test.py
@@ -90,6 +90,45 @@ def velocity_fn(x, _t):
     assert torch.allclose(result_a, result_b)
 
 
+@pytest.mark.L0
+def test_sde_preserves_conditioned_entries():
+    sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="sde")
+    noise = [torch.tensor([5.0, 1.0])]  # [N]
+    condition_reference = [torch.tensor([5.0, 0.0])]  # [N]
+    condition_mask = [torch.tensor([1.0, 0.0])]  # [N]
+
+    def velocity_fn(x, _t):
+        return [torch.zeros_like(x_i) for x_i in x]
+
+    result = sampler(
+        velocity_fn,
+        noise,
+        seed=[123],
+        condition_reference=condition_reference,
+        condition_mask=condition_mask,
+    )
+    assert torch.allclose(result[0][0], condition_reference[0][0])
+
+
+@pytest.mark.L0
+def test_ode_preserves_conditioned_entries():
+    sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.0], sample_type="ode")
+    noise = torch.tensor([5.0, 1.0])  # [N]
+    condition_reference = torch.tensor([5.0, 0.0])  # [N]
+    condition_mask = torch.tensor([1.0, 0.0])  # [N]
+
+    def velocity_fn(x, _t):
+        return torch.ones_like(x)  # [N]
+
+    result = sampler(
+        velocity_fn,
+        noise,
+        condition_reference=condition_reference,
+        condition_mask=condition_mask,
+    )
+    assert torch.allclose(result[0], condition_reference[0])
+
+
 @pytest.mark.L0
 def test_output_shape_preserved():
     sampler = FixedStepSampler(t_list=[1.0, 0.5, 0.25, 0.0])
diff --git a/cosmos_framework/model/generator/omni_mot_model.py b/cosmos_framework/model/generator/omni_mot_model.py
index 771427b2..d3a544c4 100644
--- a/cosmos_framework/model/generator/omni_mot_model.py
+++ b/cosmos_framework/model/generator/omni_mot_model.py
@@ -1625,6 +1625,8 @@ def _prepare_inference_data(
         list[list[int]],
         list[list[int]],
         list[torch.Tensor],
+        list[torch.Tensor],
+        list[torch.Tensor],
     ]:
         """
         Prepare all data needed for inference sampling.
@@ -1651,6 +1653,10 @@ def _prepare_inference_data(
                 - uncond_text_tokens: Unconditional text tokens (for CFG)
                 - initial_noise: List of noise tensors (one per sample), each containing
                   flattened vision (and optionally action) noise concatenated
+                - condition_reference: List of clean reference tensors flattened
+                  in the same order as initial_noise
+                - condition_mask: List of masks flattened in the same order as
+                  initial_noise, where 1 keeps condition_reference values fixed
         """
         # 1. Build sequence plans (same as training)
         sequence_plans = build_sequence_plans_from_data_batch(
@@ -1786,28 +1792,67 @@ def _prepare_inference_data(
         # noise_action_list and noise_sound_list are dense (only modality-having samples),
         # so we use separate indexes.
         initial_noise: list[torch.Tensor] = []
+        condition_reference: list[torch.Tensor] = []
+        condition_mask: list[torch.Tensor] = []
         idx_vision = 0
         idx_action = 0
         idx_sound = 0
 
         for i in range(n_sample):
             parts = []
+            condition_reference_parts = []
+            condition_mask_parts = []
 
             # Flatten and concatenate all vision items for this sample
             num_vis = num_items_per_sample[i] if num_items_per_sample is not None else 1
             for _ in range(num_vis):
                 parts.append(noise_vision_list[idx_vision].reshape(-1))
+                x0_vision = gen_data_clean.x0_tokens_vision[idx_vision]  # [C,T,H,W]
+                mask_vision = packed_sequence.vision.condition_mask[idx_vision].to(  # [T,1,1]
+                    dtype=x0_vision.dtype, device=x0_vision.device
+                )
+                condition_reference_parts.append(x0_vision.reshape(-1))  # [N_vision]
+                condition_mask_parts.append((mask_vision * torch.ones_like(x0_vision)).reshape(-1))  # [N_vision]
                 idx_vision += 1
 
             if noise_action_list is not None and sequence_plans[i].has_action:
+                assert packed_sequence.action is not None
+                assert packed_sequence.action.condition_mask is not None
+                assert gen_data_clean.x0_tokens_action is not None
                 parts.append(noise_action_list[idx_action].reshape(-1))
+                x0_action = gen_data_clean.x0_tokens_action[idx_action].to(  # [T,D]
+                    dtype=noise_action_list[idx_action].dtype,
+                    device=noise_action_list[idx_action].device,
+                )
+                if gen_data_clean.raw_action_dim is not None and gen_data_clean.raw_action_dim[idx_action] is not None:
+                    x0_action = x0_action.clone()  # [T,D]
+                    x0_action[:, gen_data_clean.raw_action_dim[idx_action] :] = 0  # [T,D]
+                mask_action = packed_sequence.action.condition_mask[idx_action].to(  # [T,1]
+                    dtype=x0_action.dtype, device=x0_action.device
+                )
+                condition_reference_parts.append(x0_action.reshape(-1))  # [N_action]
+                condition_mask_parts.append((mask_action * torch.ones_like(x0_action)).reshape(-1))  # [N_action]
                 idx_action += 1
 
             if noise_sound_list is not None and sequence_plans[i].has_sound:
+                assert packed_sequence.sound is not None
+                assert packed_sequence.sound.condition_mask is not None
+                assert gen_data_clean.x0_tokens_sound is not None
                 parts.append(noise_sound_list[idx_sound].reshape(-1))
+                x0_sound = gen_data_clean.x0_tokens_sound[idx_sound].to(  # [C_sound,T_sound]
+                    dtype=noise_sound_list[idx_sound].dtype,
+                    device=noise_sound_list[idx_sound].device,
+                )
+                mask_sound = packed_sequence.sound.condition_mask[idx_sound].T.to(  # [1,T_sound]
+                    dtype=x0_sound.dtype, device=x0_sound.device
+                )
+                condition_reference_parts.append(x0_sound.reshape(-1))  # [N_sound]
+                condition_mask_parts.append((mask_sound * torch.ones_like(x0_sound)).reshape(-1))  # [N_sound]
                 idx_sound += 1
 
             initial_noise.append(torch.cat(parts, dim=0))  # [N_tokens_flat]
+            condition_reference.append(torch.cat(condition_reference_parts, dim=0))  # [N_tokens_flat]
+            condition_mask.append(torch.cat(condition_mask_parts, dim=0))  # [N_tokens_flat]
 
         return (
             sequence_plans,
@@ -1815,6 +1860,8 @@ def _prepare_inference_data(
             cond_text_tokens,
             uncond_text_tokens,
             initial_noise,
+            condition_reference,
+            condition_mask,
         )
 
     def _get_velocity(
@@ -2304,6 +2351,8 @@ def generate_samples_from_batch(
             cond_tokens,
             uncond_tokens,
             initial_noise,
+            condition_reference,
+            condition_mask,
         ) = self._prepare_inference_data(data_batch, seed, has_negative_prompt)
 
         if n_sample is not None:
@@ -2477,6 +2526,13 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool):
         else:
             log.info(f"Using sampler: EDM (sigma_max={sigma_max}, num_steps={num_steps})")
 
+        fixed_step_sampler_kwargs = {}
+        if isinstance(sampler, FixedStepSampler):
+            fixed_step_sampler_kwargs = {
+                "condition_reference": condition_reference,
+                "condition_mask": condition_mask,
+            }
+
         if isinstance(sampler, FixedStepSampler) or scheduler_type == "unipc":
             latents = sampler(
                 velocity_fn,
@@ -2484,6 +2540,7 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool):
                 num_steps=num_steps,
                 shift=shift,
                 seed=seed,
+                **fixed_step_sampler_kwargs,
             )
             if _extra_num_steps > 0:
                 # Dummy sampler call to issue (_extra_num_steps × per-step)
@@ -2501,6 +2558,7 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool):
                     num_steps=_extra_num_steps,
                     shift=shift,
                     seed=seed,
+                    **fixed_step_sampler_kwargs,
                 )
         else:
             # EDM Sampler
@@ -2534,7 +2592,7 @@ def x0_fn(noise_x: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
                 # run. Avoids two EDM-specific footguns:
                 #   (1) ``EDMSampler._forward_impl`` always runs an extra
                 #       ``sample_clean`` denoiser forward (see
-                #       ``cosmos_framework/model/generator/diffusion/samplers/edm.py``).
+                #       ``cosmos_framework/model/vfm/diffusion/samplers/edm.py``).
                 #       A nested sampler call would add one too many
                 #       forwards on fast ranks, since the slow rank's
                 #       single call also pays the ``sample_clean`` cost.
diff --git a/cosmos_framework/utils/generator/data_utils.py b/cosmos_framework/utils/generator/data_utils.py
index 4ffb1e24..91d83193 100644
--- a/cosmos_framework/utils/generator/data_utils.py
+++ b/cosmos_framework/utils/generator/data_utils.py
@@ -45,6 +45,11 @@ def get_vision_data_resolution(spatial_shape: tuple[int, int]) -> str:
         return "480"
     elif min_dim <= 960:
         return "720"
+    elif min_dim <= 2048:
+        # Free-form inputs above the 720 tier (e.g. multi-reference generation
+        # producing shapes like (992, 1024)) that are not a canonical 768 shape:
+        # route to the closest defined higher tier "768".
+        return "768"
     else:
         raise ValueError(f"Unsupported resolution: {spatial_shape}")
 
diff --git a/cosmos_framework/utils/generator/torchcodec_video.py b/cosmos_framework/utils/generator/torchcodec_video.py
new file mode 100644
index 00000000..6a27ac0d
--- /dev/null
+++ b/cosmos_framework/utils/generator/torchcodec_video.py
@@ -0,0 +1,192 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""TorchCodec helpers for Cosmos3 video decoding."""
+
+from __future__ import annotations
+
+import io
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, BinaryIO
+
+import numpy as np
+import torch
+
+VideoSource = str | Path | bytes | io.BytesIO | BinaryIO
+
+
+@dataclass(frozen=True)
+class VideoMetadata:
+    num_frames: int
+    average_fps: float
+    height: int | None = None
+    width: int | None = None
+
+
+def _normalize_source(source: VideoSource) -> VideoSource:
+    if isinstance(source, Path):
+        return source.as_posix()
+    if isinstance(source, io.BytesIO):
+        source.seek(0)
+    return source
+
+
+def _get_video_decoder_cls() -> Any:
+    try:
+        from torchcodec.decoders import VideoDecoder
+    except ImportError as e:
+        raise ImportError("TorchCodec is required for Cosmos3 video decoding. Install the torchcodec package.") from e
+    return VideoDecoder
+
+
+def _build_decoder(
+    source: VideoSource,
+    *,
+    num_threads: int = 1,
+    seek_mode: str = "exact",
+    device: str = "cpu",
+) -> Any:
+    normalized_source = _normalize_source(source)
+    # Preserve FFmpeg/TorchCodec's 0 sentinel so callers can request automatic thread selection.
+    num_ffmpeg_threads = 0 if num_threads == 0 else max(num_threads, 1)
+    kwargs: dict[str, Any] = {"seek_mode": seek_mode, "num_ffmpeg_threads": num_ffmpeg_threads}
+    if device != "cpu":
+        kwargs["device"] = device
+    video_decoder_cls = _get_video_decoder_cls()
+    return video_decoder_cls(normalized_source, **kwargs)
+
+
+def _read_basic_metadata(decoder: Any) -> tuple[int, float]:
+    metadata = decoder.metadata
+    num_frames = metadata.num_frames
+    average_fps = metadata.average_fps
+    if num_frames is None or average_fps is None:
+        raise ValueError(f"TorchCodec missing metadata (num_frames={num_frames}, average_fps={average_fps})")
+    return int(num_frames), float(average_fps)
+
+
+def _metadata_from_frame(
+    decoder: Any,
+    first_frame_tchw: torch.Tensor | None = None,
+    *,
+    include_dimensions: bool = True,
+) -> VideoMetadata:
+    num_frames, average_fps = _read_basic_metadata(decoder)
+    if not include_dimensions:
+        return VideoMetadata(num_frames=num_frames, average_fps=average_fps)
+    if first_frame_tchw is None:
+        first_frame_tchw = decoder.get_frames_at([0]).data.cpu()  # [1,C,H,W]
+    _, _, height, width = first_frame_tchw.shape  # [T,C,H,W]
+    return VideoMetadata(num_frames=num_frames, average_fps=average_fps, height=int(height), width=int(width))
+
+
+def probe_video(
+    source: VideoSource,
+    *,
+    num_threads: int = 1,
+    seek_mode: str = "exact",
+    device: str = "cpu",
+    include_dimensions: bool = False,
+) -> VideoMetadata:
+    """Read video metadata, optionally decoding frame 0 to get frame dimensions."""
+    decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device)
+    return _metadata_from_frame(decoder, include_dimensions=include_dimensions)
+
+
+class TorchCodecVideoReader:
+    """Reusable indexed video reader backed by one TorchCodec decoder."""
+
+    metadata: VideoMetadata
+
+    def __init__(
+        self,
+        source: VideoSource,
+        *,
+        num_threads: int = 1,
+        seek_mode: str = "exact",
+        device: str = "cpu",
+        include_dimensions: bool = False,
+    ) -> None:
+        self._decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device)
+        self.metadata = _metadata_from_frame(self._decoder, include_dimensions=include_dimensions)
+
+    def __len__(self) -> int:
+        return self.metadata.num_frames
+
+    def __getitem__(self, index: int) -> np.ndarray:
+        return self.get_frame_nhwc_uint8(index)  # [H,W,C]
+
+    def get_avg_fps(self) -> float:
+        return self.metadata.average_fps
+
+    def get_frames_tchw_uint8(self, indices: list[int]) -> torch.Tensor:
+        frames_tchw = self._decoder.get_frames_at(indices).data.cpu()  # [T,C,H,W]
+        return frames_tchw  # [T,C,H,W]
+
+    def get_frames_nhwc_uint8(self, indices: list[int]) -> np.ndarray:
+        frames_tchw = self.get_frames_tchw_uint8(indices)  # [T,C,H,W]
+        frames_nhwc = frames_tchw.permute(0, 2, 3, 1).contiguous().numpy()  # [T,H,W,C]
+        return frames_nhwc  # [T,H,W,C]
+
+    def get_frame_nhwc_uint8(self, index: int) -> np.ndarray:
+        frames_nhwc = self.get_frames_nhwc_uint8([index])  # [1,H,W,C]
+        return frames_nhwc[0]  # [H,W,C]
+
+
+def decode_frames_tchw_uint8(
+    source: VideoSource,
+    indices: list[int],
+    *,
+    num_threads: int = 1,
+    seek_mode: str = "exact",
+    device: str = "cpu",
+) -> tuple[torch.Tensor, VideoMetadata]:
+    decoder = _build_decoder(source, num_threads=num_threads, seek_mode=seek_mode, device=device)
+    frames_tchw = decoder.get_frames_at(indices).data.cpu()  # [T,C,H,W]
+    metadata = _metadata_from_frame(decoder, frames_tchw[:1])  # frames_tchw[:1]: [1,C,H,W]
+    return frames_tchw, metadata
+
+
+def decode_frames_cthw_uint8(
+    source: VideoSource,
+    indices: list[int],
+    *,
+    num_threads: int = 1,
+    seek_mode: str = "exact",
+    device: str = "cpu",
+) -> tuple[torch.Tensor, VideoMetadata]:
+    frames_tchw, metadata = decode_frames_tchw_uint8(
+        source, indices, num_threads=num_threads, seek_mode=seek_mode, device=device
+    )
+    frames_cthw = frames_tchw.permute(1, 0, 2, 3).contiguous()  # [C,T,H,W]
+    return frames_cthw, metadata
+
+
+def decode_frames_nhwc_uint8(
+    source: VideoSource,
+    indices: list[int],
+    *,
+    num_threads: int = 1,
+    seek_mode: str = "exact",
+    device: str = "cpu",
+) -> tuple[np.ndarray, VideoMetadata]:
+    frames_tchw, metadata = decode_frames_tchw_uint8(
+        source, indices, num_threads=num_threads, seek_mode=seek_mode, device=device
+    )
+    frames_nhwc = frames_tchw.permute(0, 2, 3, 1).contiguous().numpy()  # [T,H,W,C]
+    return frames_nhwc, metadata
+
+
+def decode_frame_nhwc_uint8(
+    source: VideoSource,
+    index: int,
+    *,
+    num_threads: int = 1,
+    seek_mode: str = "exact",
+    device: str = "cpu",
+) -> tuple[np.ndarray, VideoMetadata]:
+    frames_nhwc, metadata = decode_frames_nhwc_uint8(
+        source, [index], num_threads=num_threads, seek_mode=seek_mode, device=device
+    )
+    return frames_nhwc[0], metadata  # frames_nhwc[0]: [H,W,C]