diff --git a/.file_mapping.json b/.file_mapping.json index 9cbaf068..908fb162 100644 --- a/.file_mapping.json +++ b/.file_mapping.json @@ -1,7 +1,7 @@ { - "_source_commit": "4d9b6cfd3731fdbbca883184937d40a64d2dca52-dirty", - "_dest_commit": "c23e51f2f157ae3e51cfcd86ebfb5464850894f2", - "_generated_at": "2026-09-20T05:50:07Z", + "_source_commit": "d30cb3b32430c5241644a9c0eb54a4c89fb4d1b7-dirty", + "_dest_commit": "96303bb0bdd14d9efa18f24d8ef98c7a0bfb8412", + "_generated_at": "2026-09-21T05:52:31Z", "files": { "imaginaire/__init__.py": "cosmos_framework/__init__.py", "imaginaire/attention/__init__.py": "cosmos_framework/model/attention/__init__.py", @@ -292,20 +292,6 @@ "projects/cosmos3/cosmos3/datasets/augmentors/cropping.py": "cosmos_framework/data/generator/augmentors/cropping.py", "projects/cosmos3/cosmos3/datasets/augmentors/duration_fps_text_timestamps.py": "cosmos_framework/data/generator/augmentors/duration_fps_text_timestamps.py", "projects/cosmos3/cosmos3/datasets/augmentors/duration_fps_text_timestamps_test.py": "cosmos_framework/data/generator/augmentors/duration_fps_text_timestamps_test.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/__init__.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/__init__.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/augmentor.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/augmentor.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/augmentor_test.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/augmentor_test.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/bench.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/bench.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/codec.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/codec.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/contact_sheet.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/contact_sheet.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/contact_sheet_test.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/contact_sheet_test.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/degrade.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/degrade.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/degrade_test.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/degrade_test.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/diffjpeg.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/diffjpeg.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/kernels.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/kernels.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/ops.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/ops.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/packing_test.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/packing_test.py", - "projects/cosmos3/cosmos3/datasets/augmentors/hr_lr_degradation/profiles.py": "cosmos_framework/data/generator/augmentors/hr_lr_degradation/profiles.py", "projects/cosmos3/cosmos3/datasets/augmentors/idle_frames_text_info.py": "cosmos_framework/data/generator/augmentors/idle_frames_text_info.py", "projects/cosmos3/cosmos3/datasets/augmentors/image_editing_transform.py": "cosmos_framework/data/generator/augmentors/image_editing_transform.py", "projects/cosmos3/cosmos3/datasets/augmentors/image_editing_transform_test.py": "cosmos_framework/data/generator/augmentors/image_editing_transform_test.py", diff --git a/cosmos_framework/data/generator/augmentor_provider.py b/cosmos_framework/data/generator/augmentor_provider.py index 7fc757b6..d2d6b6e4 100644 --- a/cosmos_framework/data/generator/augmentor_provider.py +++ b/cosmos_framework/data/generator/augmentor_provider.py @@ -1543,112 +1543,3 @@ def _insert_before(augmentors: dict, anchor_keys: tuple[str, ...], new_key: str, return _insert_relative(augmentors, anchor, new_key, new_value, after=False) -def _insert_low_res_stage(augmentors: dict, add_low_res) -> dict: - """Place ``AddLowRes`` so the LR is derived from exactly the HR frame the model will see. - - - Reflection-padding path (causal VAE): LR is made from the *unpadded* frame, before ``reflection_padding``; - ``SRToTrainingFormat`` pads LR separately to half the HR bucket, so LR and HR stay aligned at the top-left. - - Crop path (non-causal / UniAE, ``crop_to_multiple``): LR is made *after* the centre crop. Making it before - would derive LR from pixels the HR no longer contains (spatial misalignment) and, when the crop changes the - size, a larger LR than the target that ``SRToTrainingFormat`` cannot pad down. - """ - if "reflection_padding" in augmentors: - return _insert_relative(augmentors, "reflection_padding", "add_low_res", add_low_res, after=False) - if "crop_to_multiple" in augmentors: - return _insert_relative(augmentors, "crop_to_multiple", "add_low_res", add_low_res, after=True) - raise KeyError("Pipeline has neither reflection_padding nor crop_to_multiple; cannot place add_low_res") - - -@augmentor_register("video_basic_augmentor_v3_json_caption_sr") -def get_video_augmentor_v3_json_caption_sr( - resolution: str, - sr_scale: float = 2.0, - sr_profiles: dict[str, float] | str = "p1_first_order", - sr_seed_salt: str = "", - sr_chunk_frames: int = 8, - sr_device: str = "cpu", - sr_jpeg_backend: str = "auto", - sr_poisson_mode: str = "auto", - sr_share_vision_temporal_positions: bool = False, - **kwargs: object, -) -> dict[str, object]: - """``video_basic_augmentor_v3_json_caption`` plus an on-the-fly HR-to-LR conditioning stream. - - Adds ``AddLowRes`` (writes ``video_lr`` at ``1/sr_scale`` of the HR frame, uint8) right before - reflection padding, and ``SRToTrainingFormat`` as the last stage, which pads LR to half the HR - bucket and packs ``video = [lr, hr]`` with per-item ``image_size`` and a two-item SequencePlan. - All other stages (caption, chunked decode, sequence plan, sound) are inherited unchanged. - """ - from cosmos_framework.data.generator.augmentors.hr_lr_degradation import augmentor as sr_augmentor - - augmentors = get_video_augmentor_v3_json_caption(resolution=resolution, **kwargs) - add_low_res = L(sr_augmentor.AddLowRes)( - input_keys=["video"], - output_keys=["video_lr"], - args={ - "scale": sr_scale, - "profiles": sr_profiles, - "seed_salt": sr_seed_salt, - "modality": "video", - "chunk_frames": sr_chunk_frames, - "device": sr_device, - "jpeg_backend": sr_jpeg_backend, - "poisson_mode": sr_poisson_mode, - }, - ) - augmentors = _insert_low_res_stage(augmentors, add_low_res) - augmentors["sr_to_training_format"] = L(sr_augmentor.SRToTrainingFormat)( - input_keys=["video", "video_lr"], - args={ - "media_key": "video", - "lr_key": "video_lr", - "scale": sr_scale, - "share_vision_temporal_positions": sr_share_vision_temporal_positions, - "dataset_name": "video_sr", - }, - ) - return augmentors - - -@augmentor_register("image_basic_augmentor_with_tokenization_sr") -def image_basic_augmentor_with_tokenization_sr( - resolution: str, - sr_scale: float = 2.0, - sr_profiles: dict[str, float] | str = "p1_first_order", - sr_seed_salt: str = "", - sr_device: str = "cpu", - sr_jpeg_backend: str = "auto", - sr_poisson_mode: str = "auto", - **kwargs: object, -) -> dict[str, object]: - """``image_basic_augmentor_with_tokenization`` plus an on-the-fly HR-to-LR conditioning image. - - ``AddLowRes`` runs before ``reflection_padding`` (on the resized uint8 image), the LR copy gets - its own ``Normalize`` so both items reach the model as float in [-1, 1], and - ``SRToTrainingFormat`` packs ``images = [lr, hr]`` with per-item ``image_size``. - """ - from cosmos_framework.data.generator.augmentors.hr_lr_degradation import augmentor as sr_augmentor - - augmentors = image_basic_augmentor_with_tokenization(resolution=resolution, **kwargs) - add_low_res = L(sr_augmentor.AddLowRes)( - input_keys=["images"], - output_keys=["images_lr"], - args={ - "scale": sr_scale, - "profiles": sr_profiles, - "seed_salt": sr_seed_salt, - "modality": "image", - "chunk_frames": 1, - "device": sr_device, - "jpeg_backend": sr_jpeg_backend, - "poisson_mode": sr_poisson_mode, - }, - ) - augmentors = _insert_before(augmentors, ("reflection_padding",), "add_low_res", add_low_res) - normalize_lr = L(normalize.Normalize)(input_keys=["images_lr"], args={"mean": 0.5, "std": 0.5}) - augmentors = _insert_before(augmentors, ("text_transform",), "normalize_lr", normalize_lr) - augmentors["sr_to_training_format"] = L(sr_augmentor.SRToTrainingFormat)( - input_keys=["images", "images_lr"], - args={"media_key": "images", "lr_key": "images_lr", "scale": sr_scale, "dataset_name": "image_sr"}, - ) - return augmentors