From b3221765a7023b9b16fa23d1ab21ed647c41f84a Mon Sep 17 00:00:00 2001 From: yangyangt Date: Thu, 25 Jun 2026 21:24:36 -0700 Subject: [PATCH 1/2] Quick fix release redundant code --- .../augmentors/text_transforms_for_image.py | 2 -- .../evaluation/reconstruction_metrics.py | 3 --- cosmos_framework/model/vfm/omni_mot_model.py | 25 +++++++------------ 3 files changed, 9 insertions(+), 21 deletions(-) diff --git a/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py b/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py index 8b094ae..d38fae4 100644 --- a/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py +++ b/cosmos_framework/data/vfm/augmentors/text_transforms_for_image.py @@ -9,8 +9,6 @@ from cosmos_framework.data.imaginaire.webdataset.augmentors.augmentor import Augmentor from cosmos_framework.utils import log -# COSMOS-RELEASE-END-IGNORE - # For the qwen captions, we have 3 variants: short, medium, long # In addition, for synthetic data, we create prompt embeddings as well. # There is quite a bit of entropy in the way prompt data is saved. diff --git a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py index 3537164..eca542a 100644 --- a/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py +++ b/cosmos_framework/model/tokenizer/evaluation/reconstruction_metrics.py @@ -486,9 +486,6 @@ def reset(self) -> None: if self._fid_metric is not None: self._fid_metric.reset() - -# COSMOS-RELEASE-END-IGNORE - __all__ = [ "TokenizerMetric", "PSNRMetric", diff --git a/cosmos_framework/model/vfm/omni_mot_model.py b/cosmos_framework/model/vfm/omni_mot_model.py index f725a18..2c8bf8a 100644 --- a/cosmos_framework/model/vfm/omni_mot_model.py +++ b/cosmos_framework/model/vfm/omni_mot_model.py @@ -2432,6 +2432,15 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): # Peers needed CFG so we ran the uncond forward to keep # FSDP allgather aligned; locally we still return cond. return cond_v + + if not needs_cfg: + # This rank doesn't actually need CFG (guidance==1.0 or sigma + # outside guidance_interval). Return cond_v directly so the + # output is bit-identical to the original no-CFG path; the + # uncond_v forward was only run to keep the FSDP allgather + # sequence aligned with peers. + return cond_v + v_pred = [u_i + guidance * (c_i - u_i) for c_i, u_i in zip(cond_v, uncond_v)] if normalize_cfg: v_pred = [ @@ -2450,22 +2459,6 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): # Same alignment story as above for the postprocess branch. return cond_v - if not needs_cfg: - # This rank doesn't actually need CFG (guidance==1.0 or sigma - # outside guidance_interval). Return cond_v directly so the - # output is bit-identical to the original no-CFG path; the - # uncond_v forward was only run to keep the FSDP allgather - # sequence aligned with peers. - return cond_v - - if not needs_cfg: - # This rank doesn't actually need CFG (guidance==1.0 or sigma - # outside guidance_interval). Return cond_v directly so the - # output is bit-identical to the original no-CFG path; the - # uncond_v forward was only run to keep the FSDP allgather - # sequence aligned with peers. - return cond_v - v_pred = [u_i + guidance * (c_i - u_i) for c_i, u_i in zip(cond_v, uncond_v)] if normalize_cfg: From f6e8c4d1f1a1dca90bb7011307bf54acb739bdd8 Mon Sep 17 00:00:00 2001 From: yangyangt Date: Thu, 25 Jun 2026 23:30:32 -0700 Subject: [PATCH 2/2] fix no-used branch --- cosmos_framework/model/vfm/omni_mot_model.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/cosmos_framework/model/vfm/omni_mot_model.py b/cosmos_framework/model/vfm/omni_mot_model.py index 2c8bf8a..ea2f1b8 100644 --- a/cosmos_framework/model/vfm/omni_mot_model.py +++ b/cosmos_framework/model/vfm/omni_mot_model.py @@ -2432,14 +2432,6 @@ def _single_velocity_fn(tokens: list[list[int]], skip_text_tokens: bool): # Peers needed CFG so we ran the uncond forward to keep # FSDP allgather aligned; locally we still return cond. return cond_v - - if not needs_cfg: - # This rank doesn't actually need CFG (guidance==1.0 or sigma - # outside guidance_interval). Return cond_v directly so the - # output is bit-identical to the original no-CFG path; the - # uncond_v forward was only run to keep the FSDP allgather - # sequence aligned with peers. - return cond_v v_pred = [u_i + guidance * (c_i - u_i) for c_i, u_i in zip(cond_v, uncond_v)] if normalize_cfg: