Release: sync from internal （clean run)

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)

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(),

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)

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],
 )
 
 

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,38 +421,42 @@ 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
             data_dict["selected_caption_type"] = "editing_instruction"
 
             # 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:

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