diff --git a/.agents/skills/veomni-seed-omni/references/data-flow.md b/.agents/skills/veomni-seed-omni/references/data-flow.md index f07d641a25..7c3c8bcdd1 100644 --- a/.agents/skills/veomni-seed-omni/references/data-flow.md +++ b/.agents/skills/veomni-seed-omni/references/data-flow.md @@ -29,20 +29,24 @@ FSDP-anchor images with `mean() * 0`. ## Preprocessor Contract -`ModulePreprocessorBase` lives in `veomni/models/seed_omni/modules/module_processing_base.py`. +`ModulePreprocessorBase` lives in `veomni/models/seed_omni/modules/module_processing_base.py` +(see `docs/seed_omni/mixins/module_preprocessor.md`). Rules: +- Implement `preprocess_conversations(conversation_list, inference=False, **kwargs)` + to mutate `ConversationItem.value` and `meta` in place (or override `__call__` + when writing packed batch tensors). - Keep preprocessors picklable, weight-free, and CPU-only (using tokenizer, image/video/audio processor, config values, and special token IDs) so they can run inside DataLoader workers. -- Mutate `batch["conversation_list"]` in place (or write packed tensors onto - `batch` for packed graphs). - Shared by training and inference: - Training: run by `SeedOmniCollator` inside DataLoader workers. - - Inference: run once by `OmniInferencer._preprocess_request` before the FSM. + - Inference: run once by `OmniInferencer` via `OmniProcessor` before the FSM. - Use `inference=True` for inference-only behavior (vision modules skip FSDP - dummy injection; text encoders append generation prompts when needed). + dummy injection; text encoders append generation prompts when needed); when + `inference=False` and a modality is absent from the micro-batch, attach a + single FSDP anchor row via `ModulePreprocessorBase.append_batch_anchor`. ## Item Routing diff --git a/docs/index.md b/docs/index.md index acd70cfadf..e1c89d7fab 100644 --- a/docs/index.md +++ b/docs/index.md @@ -108,6 +108,7 @@ seed_omni/usage/adding_a_model.md seed_omni/usage/hf_layout.md seed_omni/distributed/per_module_parallelism.md seed_omni/distributed/sequence_parallel.md +seed_omni/mixins/module_preprocessor.md seed_omni/mixins/metric_meter.md seed_omni/mixins/offline_encoding.md ``` diff --git a/docs/seed_omni/design/architecture.md b/docs/seed_omni/design/architecture.md index 227694f728..c48b47d144 100644 --- a/docs/seed_omni/design/architecture.md +++ b/docs/seed_omni/design/architecture.md @@ -114,10 +114,15 @@ Native hooks in `modeling.py`, used during inference with no VeOmni graph: | `reset_local_inference_state()` / `reset_global_inference_state()` | inference | clear per-request / per-generation FSM state | | `finalize(*, ctx)` | inference | flush buffered output when `max_new_tokens` is hit | -Optional capabilities are mixins on the accelerated class, each with its own -page: [Metric Meter](../mixins/metric_meter.md) (per-module tokens and -theoretical FLOPs; `OmniModel` has no single `model_type` to estimate FLOPs on) -and [Offline Encoding](../mixins/offline_encoding.md) (cache a module's encoder +CPU input preparation is a separate picklable worker class, not a mixin: a +native class points `preprocessor_class` at a +[Module CPU Preprocessor](../mixins/module_preprocessor.md) +(`ModulePreprocessorBase`, run by `OmniProcessor` in DataLoader workers for +worker-side CPU input preparation and FSDP batch anchors). Optional capabilities +on the model itself are mixins on the accelerated class, each with its own page: +[Metric Meter](../mixins/metric_meter.md) (per-module tokens and theoretical +FLOPs; `OmniModel` has no single `model_type` to estimate FLOPs on) and +[Offline Encoding](../mixins/offline_encoding.md) (cache a module's encoder output once and train from the cache). ### 2.2 `ConversationItem`: the data carrier (`utils/conversation.py`) diff --git a/docs/seed_omni/mixins/module_preprocessor.md b/docs/seed_omni/mixins/module_preprocessor.md new file mode 100644 index 0000000000..a2ea006bfe --- /dev/null +++ b/docs/seed_omni/mixins/module_preprocessor.md @@ -0,0 +1,133 @@ +# Module CPU Preprocessors (`ModulePreprocessorBase`) + +`ModulePreprocessorBase` is the per-module CPU worker contract for SeedOmni. Each +module that needs CPU-side input preparation (chat-template application, +tokenization, image normalization, or FSDP dummy-anchor injection) subclasses +[`ModulePreprocessorBase`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L55) +in its `processing.py` and points its native model class's `preprocessor_class` +at that subclass. + +- Base class: [`ModulePreprocessorBase`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L55) +- Asset binding helper: [`bind_module_assets`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L145) +- Composed orchestrator: [`OmniProcessor`](../../../veomni/models/seed_omni/processing_omni.py#L81) +- Collator caller: [`SeedOmniCollator`](../../../veomni/data/seed_omni/collator.py#L26) + +## Why preprocessing is per-module and CPU-only + +The global data transform ([`process_seedomni_example`](../../../veomni/data/seed_omni/seedomni_transform.py#L183)) +decodes raw media and produces a model-agnostic `conversation_list` of +[`ConversationItem`](../../../veomni/models/seed_omni/utils/conversation.py#L78) +objects. Only each module knows which items it consumes, how its chat template or +image processor transforms them, and what dummy placeholder shape its FSDP +wrapper needs when a micro-batch has no real items for that modality. + +Running that preparation inside [`ModulePreprocessorBase`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L55) +keeps heavy CPU work (tokenization, image resize/normalize) inside DataLoader +worker processes during training while sharing the exact same preprocessing chain +with single-request inference via [`OmniProcessor`](../../../veomni/models/seed_omni/processing_omni.py#L81). + +## Opting in + +Define a `ModulePreprocessorBase` subclass in `modules///processing.py` +and attach it as `preprocessor_class` on the native model class in `modeling.py`: + +```python +from typing import Any + +import torch + +from veomni.models.seed_omni.modules.module_processing_base import ModulePreprocessorBase +from veomni.models.seed_omni.utils.conversation import ( + _IMG_TAG_KEY, + ConversationItem, + iter_desired_items, +) + + +class VisionPreprocessor(ModulePreprocessorBase): + def __init__(self, image_processor: Any, dtype: torch.dtype = torch.bfloat16) -> None: + self._image_processor = image_processor + self.dtype = dtype + self._dummy_pixel_values: torch.Tensor | None = None + + @classmethod + def from_pretrained( + cls, module_path: str, *, config_overrides: dict[str, Any] | None = None, **kwargs: Any + ) -> "VisionPreprocessor": + del kwargs + config = VisionConfig.from_pretrained(module_path, model_config=config_overrides) + image_processor = ... # build CPU-safe processor from module_path / config + return cls(image_processor=image_processor) + + def bind_dummy_inputs(self, config: Any, dtype: Any = None) -> None: + self._dummy_pixel_values = torch.zeros(3, config.image_size, config.image_size, dtype=dtype or self.dtype) + + def preprocess_conversations( + self, conversation_list: list[list[ConversationItem]], inference: bool = False, **kwargs: Any + ) -> None: + del kwargs + user_images = [ + it + for it in iter_desired_items(conversation_list, types=["image"], roles=["user"]) + if not it.is_dummy + ] + for item in user_images: + item.value = self._image_processor(item.value) + if not user_images and not inference and self._dummy_pixel_values is not None: + self.append_batch_anchor( + conversation_list, + ConversationItem( + type="image", + value=self._dummy_pixel_values.clone(), + role="user", + is_dummy=True, + meta={_IMG_TAG_KEY: "und"}, + ), + ) +``` + +## The functions + +### What a module implements or overrides + +| Function | Module's job | Called by | Effect | +|----------|--------------|-----------|--------| +| [`preprocess_conversations(conversation_list, inference=False, **kwargs)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L89) | **Must implement** (unless overriding `__call__`). Mutate matching `ConversationItem.value` / `meta` in place across `conversation_list` (`list[list[ConversationItem]]`). When `inference=False` and no real item matches this module, append a dummy placeholder via [`append_batch_anchor`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L128). | [`ModulePreprocessorBase.__call__`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L85) | Prepares CPU tensors for the module's `pre_forward` / `generate`. Not implementing it raises `NotImplementedError` when the preprocessor runs. | +| [`from_pretrained(module_path, *, config_overrides=None, **kwargs)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L99) | **May override.** Build and return the CPU preprocessor from the module's checkpoint subfolder without instantiating the `nn.Module`. Default returns `None`. | [`OmniProcessor.from_config`](../../../veomni/models/seed_omni/processing_omni.py#L112) and [`bind_module_assets`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L145) | Returning `None` means the module contributes no CPU preprocessor (for example a pure AR backbone). | +| [`bind_dummy_inputs(config, dtype=None)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L119) | **May override.** Pre-compute dummy tensor(s) from the module's resolved `config` and `dtype` for the `inference=False` branch. Default is a no-op. | [`OmniProcessor.bind_dummy_inputs`](../../../veomni/models/seed_omni/processing_omni.py#L91) during [`OmniTrainer._build_collate_fn`](../../../veomni/trainer/omni/omni_trainer.py#L306) | Attaches FSDP-anchor dummy shapes before DataLoader workers fork/spawn. | +| [`__call__(batch, inference=False, **kwargs)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L85) | **May override** for packed preprocessors that write extra batch-level tensors onto `batch` in addition to or instead of mutating `batch["conversation_list"]`. Default delegates `batch["conversation_list"]` to `preprocess_conversations`. | [`OmniProcessor.preprocess_batch`](../../../veomni/models/seed_omni/processing_omni.py#L231) | Runs this module's CPU preparation step in `config.module_names` order. | + +### What the framework provides (do not override) + +| Function | Called by | Effect | +|----------|-----------|--------| +| [`ModulePreprocessorBase.append_batch_anchor(conversation_list, item)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L128) | Module's `preprocess_conversations` when `inference=False` and no real input for this module exists in the micro-batch | Appends `item` (`is_dummy=True`) to `conversation_list[0]` — a single anchor row per micro-batch rather than one per sample, keeping FSDP2 collectives aligned without running `N` dummy forward passes. | +| [`bind_module_assets(model, *, checkpoint_path=None, preprocessor=None, config_overrides=None)`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L145) | [`PretrainedOmniModule.from_pretrained`](../../../veomni/models/seed_omni/modules/module_modeling_base.py#L32) and [`ModuleRuntime._build_model_assets`](../../../veomni/models/seed_omni/accelerated/omni_module/omni_module_runtime.py#L354) | Copies CPU asset attributes (`_processor`, `_image_processor`, `_video_processor`, `_tokenizer`, `_chat_template`) from the preprocessor onto `model` if not already set, so `generate()` and HF checkpoint export can access them. | + +## Call flow + +- **Training micro-batch**: + - [`OmniTrainer._build_collate_fn`](../../../veomni/trainer/omni/omni_trainer.py#L306) builds [`OmniProcessor.from_config`](../../../veomni/models/seed_omni/processing_omni.py#L112), calls [`OmniProcessor.bind_dummy_inputs`](../../../veomni/models/seed_omni/processing_omni.py#L91), and wraps it in [`SeedOmniCollator`](../../../veomni/data/seed_omni/collator.py#L26). + - [`SeedOmniCollator.__call__`](../../../veomni/data/seed_omni/collator.py#L74) verifies key alignment across samples, groups `conversation_list`, and calls [`OmniProcessor.preprocess_batch(..., inference=False)`](../../../veomni/models/seed_omni/processing_omni.py#L231). + - Each module's [`ModulePreprocessorBase.__call__`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L85) calls [`preprocess_conversations`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L89), mutating items in place and appending at most one batch anchor via [`append_batch_anchor`](../../../veomni/models/seed_omni/modules/module_processing_base.py#L128). +- **Inference request**: + - [`OmniInferencer._build_model`](../../../veomni/trainer/omni/omni_inferencer.py#L120) builds [`OmniProcessor.from_config`](../../../veomni/models/seed_omni/processing_omni.py#L112). + - [`OmniInferencer._run`](../../../veomni/trainer/omni/omni_inferencer.py#L237) calls [`OmniProcessor.__call__(..., inference=True)`](../../../veomni/models/seed_omni/processing_omni.py#L170) → [`OmniProcessor.preprocess`](../../../veomni/models/seed_omni/processing_omni.py#L219) → [`OmniProcessor.preprocess_batch(..., inference=True)`](../../../veomni/models/seed_omni/processing_omni.py#L231). + +## Rules an implementation must follow + +1. **Keep preprocessors picklable, weight-free, and CPU-only.** DataLoader workers fork or spawn with a pickled copy of `OmniProcessor`; never store an `nn.Module` or allocate CUDA tensors inside a preprocessor. +2. **Skip existing dummy rows when selecting real inputs.** Filter with `not item.is_dummy` so a dummy item injected earlier does not get re-processed as real data. +3. **Use `append_batch_anchor` for missing-modality FSDP anchors.** Only inject dummy items when `not inference`, and attach a single micro-batch anchor rather than one per sample. +4. **Honour `config_overrides` in `from_pretrained`.** Apply any `model_config` overrides passed via `config_overrides` on top of the on-disk `config.json` so the worker preprocessor matches the live model's configuration. + +## Tests + +- [`tests/seed_omni/data/test_preprocessor.py`](../../../tests/seed_omni/data/test_preprocessor.py): + - [strict key alignment across collated samples](../../../tests/seed_omni/data/test_preprocessor.py#L52), + - [mixed-resolution images preserved unstacked and normalized per item](../../../tests/seed_omni/data/test_preprocessor.py#L62), and + - [`append_batch_anchor` attaches a single dummy anchor row when a modality is absent](../../../tests/seed_omni/data/test_preprocessor.py#L101). +- [`tests/seed_omni/model/test_processor.py`](../../../tests/seed_omni/model/test_processor.py): + - [preprocessors run in module order](../../../tests/seed_omni/model/test_processor.py#L36), + - [`model_config` and `processor_config` overrides reach `from_pretrained`](../../../tests/seed_omni/model/test_processor.py#L162), and + - [`bind_module_assets` copies missing assets idempotently](../../../tests/seed_omni/model/test_processor.py#L230). diff --git a/tests/data/utils.py b/tests/data/utils.py index d32247e31c..a10e4ecfec 100644 --- a/tests/data/utils.py +++ b/tests/data/utils.py @@ -31,6 +31,9 @@ def setup_test_distributed(args): """Initialize a minimal distributed runtime for data tests. Returns the device and the ``ParallelState``. + + Registers under ``"base"`` because callers stand in for ``BaseTrainer._setup``, + and ``BaseTrainer`` scopes its build under ``use_parallel_state("base")``. """ device_type = get_device_type() if device_type != "cpu": @@ -60,6 +63,7 @@ def setup_test_distributed(args): extra_parallel_placement_innermost=args.model.accelerator.extra_parallel_placement_innermost, extra_parallel_names=args.model.accelerator.extra_parallel_names, dp_mode=args.model.accelerator.fsdp_config.fsdp_mode, + name="base", ) helper.set_seed(args.train.seed, args.train.enable_full_determinism) return device, parallel_state diff --git a/tests/seed_omni/data/__init__.py b/tests/seed_omni/data/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/seed_omni/data/test_audio_ingestion.py b/tests/seed_omni/data/test_audio_ingestion.py new file mode 100644 index 0000000000..d3da48981d --- /dev/null +++ b/tests/seed_omni/data/test_audio_ingestion.py @@ -0,0 +1,359 @@ +"""Audio into the SeedOmni carrier, and the rate that has to survive the trip. + +The data layer loads a clip at whatever rate the file holds and declares it; +each module then converts to its own. That split exists because the two modules +reading a conversation disagree — a Qwen3-Omni thinker's tower runs at 16 kHz, +the codec that turns assistant speech into talker targets at 24 kHz — so any +rate the data layer picked would be wrong for one of them, and the second +conversion cannot restore what the first removed. + +The rate is also what places a clip on TMRoPE's shared clock, so a dropped or +wrong rate is a silent alignment bug rather than an audio-quality one. +""" + +from __future__ import annotations + +import io +import shutil +import subprocess +import wave +from unittest import mock + +import numpy as np +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.data.seed_omni.preprocess import SEED_OMNI_PREPROCESSOR_REGISTRY, conv_preprocess +from veomni.data.seed_omni.seedomni_transform import process_seedomni_example +from veomni.data.seed_omni.utils import audio as audio_utils +from veomni.data.seed_omni.utils.audio import AUDIO_METADATA_KEY, AudioMetadata, load_audio + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +def _wav_bytes(seconds: float, rate: int, channels: int = 1) -> bytes: + """A PCM wav declaring ``rate``, so the loader has a real header to read.""" + frames = int(seconds * rate) + tone = np.sin(2 * np.pi * 220 * np.arange(frames) / rate) + samples = (tone * 16384).astype(" 1: + samples = np.repeat(samples[:, None], channels, axis=1).reshape(-1) + buffer = io.BytesIO() + with wave.open(buffer, "wb") as handle: + handle.setnchannels(channels) + handle.setsampwidth(2) + handle.setframerate(rate) + handle.writeframes(samples.tobytes()) + return buffer.getvalue() + + +def _sample(audios: list[bytes], turns: list[tuple[str, str]]) -> dict: + return { + "source_name": "voice_assistant", + "conversations": [{"from": speaker, "value": value} for speaker, value in turns], + "audios": audios, + } + + +def test_a_clip_arrives_at_its_own_rate_not_a_normalised_one(): + """22.05 kHz is neither module's rate, and it must reach them as 22.05 kHz.""" + items = process_seedomni_example(_sample([_wav_bytes(1.5, 22050)], [("human", "hi"), ("gpt", "hello")]))[0][ + "conversation_list" + ] + audio = next(item for item in items if item.type == "audio") + + assert audio.meta[AUDIO_METADATA_KEY].sampling_rate == 22050 + assert audio.value.dtype == np.float32 + assert audio.value.shape == (33075,) + + +def test_the_declared_rate_and_the_sample_count_agree_on_the_duration(): + """The one invariant every downstream clock reads: samples / rate == seconds.""" + items = process_seedomni_example(_sample([_wav_bytes(2.0, 44100)], [("human", "hi"), ("gpt", "yes")]))[0][ + "conversation_list" + ] + audio = next(item for item in items if item.type == "audio") + + metadata = audio.meta[AUDIO_METADATA_KEY] + assert audio.value.size / metadata.sampling_rate == pytest.approx(2.0, abs=1e-3) + # The item states the duration itself, so a consumer reading the clock does + # not have to still be holding the waveform to work it out. + assert metadata.duration_seconds == pytest.approx(2.0, abs=1e-3) + + +def test_a_stereo_file_is_downmixed_where_its_layout_is_known(): + """The modules take mono and refuse to guess a mix, so this layer owes them one.""" + samples, rate = load_audio(_wav_bytes(1.0, 16000, channels=2)) + + assert rate == 16000 + assert samples.ndim == 1 + assert samples.size == 16000 + + +def test_a_one_frame_stereo_file_is_not_read_as_two_samples(): + """The decoded layout is known, so it is not guessed. + + ``sf.read`` gives ``(frames, channels)``. A guess that the channel axis is + the shorter one is right for any real clip but inverts here, turning two + channels into two samples. + """ + samples, rate = load_audio(_wav_bytes(1 / 16000, 16000, channels=2)) + + assert samples.shape == (1,) + assert rate == 16000 + + +@pytest.mark.parametrize("shape", [(1000, 2), (2, 1000)], ids=["frames-first", "channels-first"]) +def test_a_bare_stereo_array_is_downmixed_whichever_way_round_it_is(shape): + """The channel axis is guessed here, and a wrong guess is silent. + + A bare array carries no layout, so the shorter axis is taken to be the + channels. Inverting that collapses the clip to one value per channel — in + *both* orientations, so a single-orientation test cannot catch it, and the + empty-clip guard does not either since two samples is not zero. + """ + samples, rate = load_audio(np.zeros(shape, dtype=np.float32), audio_sampling_rate=16000) + + assert samples.shape == (1000,) + assert rate == 16000 + + +def test_an_empty_clip_is_refused_rather_than_averaged_into_nan(): + """Averaging an empty axis yields NaN with only a warning. + + Those NaN samples surface much later as a NaN loss, with nothing pointing + back to the file that produced them. + """ + with pytest.raises(ValueError, match="carries no samples"): + load_audio(_wav_bytes(0.0, 16000, channels=2)) + + +@pytest.mark.parametrize("as_bytes", [False, True], ids=["path", "bytes"]) +def test_a_container_libsndfile_cannot_read_falls_back_to_audioread(tmp_path, as_bytes): + """m4a / aac / wma need a second backend, and bytes is the likely form. + + A parquet corpus stores clips inline, so an m4a source arrives as bytes — + the backend takes a filename, hence the spill. The fallback must keep the + native rate: librosa's own default would resample to 22.05 kHz. + + ffmpeg is needed here to *author* the file; decoding may go through it or + gstreamer, so a box with only the latter still exercises the path. + """ + ffmpeg = shutil.which("ffmpeg") + if ffmpeg is None: + pytest.skip("ffmpeg not available to author a non-libsndfile container") + path = tmp_path / "clip.m4a" + subprocess.run( + [ + ffmpeg, + "-f", + "lavfi", + "-i", + "sine=frequency=440:duration=1.5:sample_rate=32000", + "-ac", + "1", + "-y", + str(path), + "-loglevel", + "error", + ], + check=True, + ) + + samples, rate = load_audio(path.read_bytes() if as_bytes else str(path)) + + assert rate == 32000 + assert samples.size / rate == pytest.approx(1.5, abs=0.05) + + +def test_a_bare_array_must_be_told_its_rate(): + """An array carries no header, and guessing one would misplace it on the clock.""" + with pytest.raises(ValueError, match="declares no rate"): + load_audio(np.zeros(1000, dtype=np.float32)) + + samples, rate = load_audio(np.zeros(1000, dtype=np.float32), audio_sampling_rate=8000) + assert (samples.size, rate) == (1000, 8000) + + +def test_a_repeated_ref_is_decoded_once_but_not_shared(): + """Decoded once for speed, copied for safety. + + The two rows holding one clip can be the user's and the assistant's, which + different modules read, so sharing one array would let the first in-place + write anyone adds corrupt the other module's input. + """ + clip = _wav_bytes(0.5, 16000) + decodes = 0 + real_load = audio_utils.load_audio + + def counting_load(*args, **kwargs): + nonlocal decodes + decodes += 1 + return real_load(*args, **kwargs) + + with mock.patch.object(audio_utils, "load_audio", counting_load): + loaded = audio_utils.fetch_audios([clip, clip]) + + assert decodes == 1 + assert loaded[0][0] is not loaded[1][0] + assert np.array_equal(loaded[0][0], loaded[1][0]) + assert loaded[0][1][AUDIO_METADATA_KEY] == loaded[1][1][AUDIO_METADATA_KEY] + + +def _rate_meta(rate: int) -> dict: + """An emitted audio item's meta, which is what ``save_audio`` is handed.""" + return {AUDIO_METADATA_KEY: AudioMetadata(sampling_rate=rate)} + + +def test_a_generated_waveform_is_written_however_the_vocoder_shaped_it(tmp_path): + """The shape and dtype the codec actually emits: batched, and bf16 under AMP. + + Both conversions the writer makes are load-bearing: soundfile takes only + 1-D or 2-D, so an unsqueezed clip raises rather than writing, and + ``Tensor.numpy()`` refuses bfloat16 outright. + """ + rate = 24000 + frames = rate // 10 + tone = np.sin(2 * np.pi * 220 * np.arange(frames) / rate).astype(np.float32) + path = tmp_path / "generated_audio_0.wav" + + audio_utils.save_audio(str(path), torch.from_numpy(tone).to(torch.bfloat16)[None, None, :], _rate_meta(rate)) + + written, written_rate = load_audio(str(path)) + assert written_rate == rate + assert written.shape == tone.shape + # bfloat16 keeps 8 mantissa bits, so the tolerance is set by the dtype the + # vocoder ran in, not by the wav. + np.testing.assert_allclose(written, tone, atol=1e-2) + + +def test_a_generated_clip_with_no_rate_is_refused_rather_than_written_at_a_guess(tmp_path): + """A default rate writes a file that plays at the wrong speed, and nothing + downstream can tell that apart from a model that generated it wrong.""" + with pytest.raises(ValueError, match=AUDIO_METADATA_KEY): + audio_utils.save_audio(str(tmp_path / "generated_audio_0.wav"), torch.zeros(2400), {}) + + +def test_a_multi_channel_waveform_is_refused_by_name(tmp_path): + """Anything still 2-D after the squeeze is a real batch or true stereo, and + neither names a single mono utterance. + + Left to soundfile this is not silent — it reads the 2400 axis as channels + and libsndfile rejects the format — but the error names neither the item nor + its shape, which is the whole diagnostic. + """ + with pytest.raises(ValueError, match="not a single"): + audio_utils.save_audio(str(tmp_path / "generated_audio_0.wav"), torch.zeros(2, 2400), _rate_meta(24000)) + + +def test_an_empty_generated_clip_is_refused_rather_than_written_as_a_bare_header(tmp_path): + """Symmetry with the load side, which refuses an empty clip by name. + + ``squeeze`` leaves a zero-length clip 1-D, so it passes the mono check; + without this guard soundfile writes a valid 44-byte header and the caller + logs a successful write, leaving a playable-looking file with nothing in it. + """ + path = tmp_path / "generated_audio_0.wav" + with pytest.raises(ValueError, match="no samples"): + audio_utils.save_audio(str(path), torch.zeros(1, 1, 0), _rate_meta(24000)) + assert not path.exists() + + +def test_a_one_sample_clip_is_mono_not_a_batching_mistake(tmp_path): + """``squeeze`` takes a single sample down to a scalar, which is still mono. + + Degenerate, but reporting it as "not a single mono waveform. Emit one item + per utterance" would send the reader looking for a batching bug. + """ + path = tmp_path / "generated_audio_0.wav" + audio_utils.save_audio(str(path), torch.ones(1, 1, 1), _rate_meta(24000)) + + samples, rate = load_audio(str(path)) + assert (samples.size, rate) == (1, 24000) + + +def test_more_clips_than_spoken_turns_is_refused_rather_than_silently_dropped(): + with pytest.raises(ValueError, match="spoken turn"): + conv_preprocess( + "voice_assistant", + [{"from": "human", "value": "hi"}, {"from": "gpt", "value": "hello"}], + {"audios": [_wav_bytes(0.5, 16000), _wav_bytes(0.5, 16000)]}, + ) + + +def test_fewer_clips_than_spoken_turns_is_refused_too(): + """Both directions, and both by a named error. + + A bare ``next()`` running out would raise ``StopIteration``, which PEP 479 + turns into a ``RuntimeError`` once it crosses a generator frame — and the + dataset stack is full of them, so the type would depend on the call site. + """ + with pytest.raises(ValueError, match="paired by position"): + conv_preprocess( + "voice_assistant", + [{"from": "human", "value": "hi"}, {"from": "human", "value": "again"}], + {"audios": [_wav_bytes(0.5, 16000)]}, + ) + + +@pytest.fixture +def _declared_rate_source(): + """A source that asserts a rate on the item meta, contradicting the file.""" + SEED_OMNI_PREPROCESSOR_REGISTRY["_declared_rate"] = lambda conversations, example, **kwargs: ( + [["user", ("audio", None, {AUDIO_METADATA_KEY: AudioMetadata(sampling_rate=16000)})]], + {"audio": list(example["audios"])}, + ) + yield "_declared_rate" + del SEED_OMNI_PREPROCESSOR_REGISTRY["_declared_rate"] + + +@pytest.fixture +def _unknown_modality_source(): + """A source declaring refs for a modality that has no fetcher yet.""" + SEED_OMNI_PREPROCESSOR_REGISTRY["_unknown_modality"] = lambda conversations, example, **kwargs: ( + [["user", ("text", "hi")]], + {"action": [object()]}, + ) + yield "_unknown_modality" + del SEED_OMNI_PREPROCESSOR_REGISTRY["_unknown_modality"] + + +def test_a_declared_rate_that_contradicts_the_file_is_refused(_declared_rate_source): + """The header is the only authority on the rate, so a preprocessor does not + get to declare one. + + Preferring the declared value would be a silent misplacement on TMRoPE's + clock rather than a clip that sounds wrong, and letting the decoded one + overwrite it silently would hide the preprocessor's false belief. + """ + with pytest.raises(ValueError, match=f"declares meta key\\(s\\) \\['{AUDIO_METADATA_KEY}'\\]"): + process_seedomni_example({"source_name": _declared_rate_source, "audios": [_wav_bytes(0.5, 22050)]}) + + +def test_refs_for_a_modality_with_no_fetcher_are_refused_by_name(_unknown_modality_source): + """A modality key the transform cannot decode would leave its turns + unpaired, surfacing much later as a count mismatch blaming the wrong thing. + + This is the seam a new modality (action / 3d / camera) lands on: registering + a fetcher in ``MEDIA_FETCHERS`` is what makes the key decodable, and until + then the refs are refused rather than dropped. + """ + with pytest.raises(ValueError, match="'action'.*has no fetcher"): + process_seedomni_example({"source_name": _unknown_modality_source, "conversations": [], "audios": []}) + + +def test_the_voice_assistant_layout_is_speech_in_text_out(): + """The clip is the user's, so it is an encoder input — not a talker target. + + ``role`` is the only thing separating the two uses of ``type="audio"``, and + the transcript is dropped: fed as text the model would answer from it and + the tower would carry no signal at all. + """ + items = process_seedomni_example( + _sample([_wav_bytes(1.0, 22050)], [("human", "what is your name"), ("gpt", "I am Omni")]) + )[0]["conversation_list"] + + assert [(item.role, item.type) for item in items] == [("user", "audio"), ("assistant", "text")] + assert "what is your name" not in str([item.value for item in items if item.type == "text"]) diff --git a/tests/seed_omni/data/test_img_tag.py b/tests/seed_omni/data/test_img_tag.py new file mode 100644 index 0000000000..65936c4b1d --- /dev/null +++ b/tests/seed_omni/data/test_img_tag.py @@ -0,0 +1,657 @@ +"""Per-image ``_img_tag`` (und/gen/edit) data-layer support — Plan 1. + +Covers the model-agnostic data-layer work only: + +* :func:`iter_desired_items` — generic item selection with optional direct + matching over configured ``meta`` values. +* The optional ``(type, value, meta)`` 3-tuple meta channel in + ``_build_conversation_list`` (2-tuple stays backward compatible) — pure + sequential image/video pairing; the preprocessor owns the ref-list layout. +* Existing preprocessors tag their image entries (imagenet1k=gen, + sharegpt4v=und, llava_video=und); text-only tulu stays untagged. Each + preprocessor returns ``(constructed, media_refs)`` where ``media_refs`` is + keyed by item type and each list's length matches that modality's flattened + entry count. +* The ``seed_edit_p23_multi_turn`` preprocessor emits the source/target/copy + edit chain with correct roles and tags, and returns an expanded + ``image_refs`` list (copy-image refs duplicated) so + ``_build_conversation_list`` pairs them by pure sequential order. +""" + +from __future__ import annotations + +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.data.seed_omni.preprocess import conv_preprocess +from veomni.data.seed_omni.seedomni_transform import _build_conversation_list +from veomni.models.seed_omni.utils.conversation import ( + _IMG_TAG_KEY, + ConversationItem, + iter_desired_items, +) + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +# ── helpers ──────────────────────────────────────────────────────────────────── + + +def _img(tag: str | None, role: str = "user", is_dummy: bool = False) -> ConversationItem: + meta = {} if tag is None else {_IMG_TAG_KEY: tag} + return ConversationItem( + type="image", + value=torch.zeros(3, 2, 2, dtype=torch.uint8), + role=role, + is_dummy=is_dummy, + meta=meta, + ) + + +def _txt(role: str = "user") -> ConversationItem: + return ConversationItem(type="text", value="hello", role=role) + + +def _fake_images(n: int) -> list[torch.Tensor]: + # Distinct fill values so sequential pairing is checkable by content. + return [torch.full((3, 2, 2), i, dtype=torch.uint8) for i in range(n)] + + +def _media(**by_type) -> dict[str, list]: + """Decoded media in the shape ``_build_conversation_list`` takes. + + Keyed by item type, each entry a ``(payload, meta)`` pair — what the + fetchers in ``MEDIA_FETCHERS`` return. Payloads are passed bare here and + paired with an empty meta, since these tests are about the layout rather + than about what any modality states. + """ + return {type_: [(payload, {}) for payload in payloads] for type_, payloads in by_type.items()} + + +def _fake_tensors_for_refs(refs) -> list[torch.Tensor]: + # One distinct uint8 tensor per unique ref (by identity), so duplicate refs + # (edit copy-images) decode to the same content — mirroring ``fetch_images`` + # on a repeated ref without actually decoding. + cache: dict[int, torch.Tensor] = {} + out: list[torch.Tensor] = [] + for ref in refs: + key = id(ref) + if key not in cache: + cache[key] = torch.full((3, 2, 2), len(cache), dtype=torch.uint8) + out.append(cache[key]) + return out + + +def _build(source: str, conversations, example: dict): + """Run a preprocessor end-to-end through ``_build_conversation_list``. + + ``example`` carries raw image/video refs (opaque objects are fine — they are + not decoded here); the helper builds distinct fake tensors keyed on ref + identity (one per unique ref) of the lengths the preprocessor reports in its + ``media_refs`` dict, so the sequential pairing in + ``_build_conversation_list`` is exercised and duplicate copy-image refs + resolve to identical content. + """ + constructed, media_refs = conv_preprocess(source, conversations, example) + image_refs = media_refs.get("image", []) + video_refs = media_refs.get("video", []) + media = _media( + image=_fake_tensors_for_refs(image_refs), + # Videos are opaque payload stand-ins — nothing here reads a frame. + video=[object() for _ in video_refs], + ) + items = _build_conversation_list(constructed, media) + return constructed, image_refs, video_refs, items + + +# ── iter_desired_items meta selector ─────────────────────────────────────────── + + +def test_iter_desired_items_no_meta_filter_returns_all_images(): + batch = [[_img("und"), _img("gen"), _img("edit")]] + items = list(iter_desired_items(batch, types=["image"])) + assert [it.meta.get(_IMG_TAG_KEY) for it in items] == ["und", "gen", "edit"] + + +def test_iter_desired_items_filters_by_each_meta_value(): + batch = [[_img("und"), _img("gen"), _img("edit"), _img("und"), _img("gen")]] + assert [ + it.meta.get(_IMG_TAG_KEY) for it in iter_desired_items(batch, types=["image"], meta={_IMG_TAG_KEY: ["und"]}) + ] == ["und", "und"] + assert [ + it.meta.get(_IMG_TAG_KEY) for it in iter_desired_items(batch, types=["image"], meta={_IMG_TAG_KEY: ["gen"]}) + ] == ["gen", "gen"] + assert [ + it.meta.get(_IMG_TAG_KEY) for it in iter_desired_items(batch, types=["image"], meta={_IMG_TAG_KEY: ["edit"]}) + ] == ["edit"] + + +def test_iter_desired_items_meta_value_set_matches_multiple(): + batch = [[_img("und"), _img("gen"), _img("edit")]] + got = [ + it.meta.get(_IMG_TAG_KEY) + for it in iter_desired_items(batch, types=["image"], meta={_IMG_TAG_KEY: ["und", "edit"]}) + ] + assert got == ["und", "edit"] + + +def test_iter_desired_items_type_filter_excludes_text_items(): + batch = [[_img("und"), _txt("user"), _txt("assistant")]] + items = list(iter_desired_items(batch, types=["image"])) + assert len(items) == 1 + assert items[0].type == "image" + + +def test_iter_desired_items_meta_composes_with_roles(): + batch = [ + [ + _img("und", role="user"), + _img("gen", role="assistant"), + _img("edit", role="user"), + _img("gen", role="assistant", is_dummy=True), + ] + ] + by_role = [it.meta.get(_IMG_TAG_KEY) for it in iter_desired_items(batch, types=["image"], roles=["user"])] + assert by_role == ["und", "edit"] + by_both = [ + (it.meta.get(_IMG_TAG_KEY), it.is_dummy) + for it in iter_desired_items( + batch, + types=["image"], + roles=["assistant"], + meta={_IMG_TAG_KEY: ["gen"]}, + ) + ] + assert by_both == [("gen", False), ("gen", True)] + with_tag_and_role = [ + it.meta.get(_IMG_TAG_KEY) + for it in iter_desired_items( + batch, + types=["image"], + roles=["user"], + meta={_IMG_TAG_KEY: ["edit"]}, + ) + ] + assert with_tag_and_role == ["edit"] + + +def test_iter_desired_items_meta_preserves_micro_batch_order_and_reverse(): + a, b, c = _img("und"), _img("gen"), _img("edit") + batch = [[a, b], [c]] + assert list(iter_desired_items(batch, types=["image"])) == [a, b, c] + # reverse_item reverses within each sample. + assert list(iter_desired_items(batch, types=["image"], reverse_item=True)) == [b, a, c] + + +def test_iter_desired_items_missing_meta_key_is_excluded_when_filter_set(): + # An image with no _img_tag is dropped when a tag filter is active, but kept + # when meta is None (selector must not crash on a missing meta key). + untagged = ConversationItem(type="image", value=torch.zeros(3, 2, 2, dtype=torch.uint8), role="user") + batch = [[untagged, _img("gen")]] + assert [ + it.meta.get(_IMG_TAG_KEY) for it in iter_desired_items(batch, types=["image"], meta={_IMG_TAG_KEY: ["gen"]}) + ] == ["gen"] + assert len(list(iter_desired_items(batch, types=["image"]))) == 2 + + +# ── 3-tuple meta channel in _build_conversation_list ─────────────────────────── + + +def test_two_tuple_still_builds_with_empty_meta(): + constructed = [["user", ("text", "hi")], ["assistant", ("image", None)]] + items = _build_conversation_list(constructed, _media(image=_fake_images(1))) + assert items[0].type == "text" and items[0].value == "hi" + assert items[1].type == "image" and items[1].meta == {} + # Fill value 0 — sequential pairing of the first (only) image tensor. + assert int(items[1].value.float().mean()) == 0 + + +def test_a_text_only_sample_needs_no_media_argument(): + """What a text-only preprocessor's ``{}`` refs dict turns into: a build with + no media at all, rather than three empty lists it had to name.""" + items = _build_conversation_list([["user", ("text", "hi")]]) + assert [item.type for item in items] == ["text"] + + +def test_three_tuple_meta_merges_into_item_meta(): + constructed = [["assistant", ("image", None, {_IMG_TAG_KEY: "gen", "extra": 7})]] + items = _build_conversation_list(constructed, _media(image=_fake_images(1))) + assert items[0].meta == {_IMG_TAG_KEY: "gen", "extra": 7} + + +def test_three_tuple_meta_does_not_alias_preprocessor_dict(): + src_meta = {_IMG_TAG_KEY: "edit"} + constructed = [["assistant", ("image", None, src_meta)]] + items = _build_conversation_list(constructed, _media(image=[_fake_images(1)[0]])) + items[0].meta["mutated"] = True + assert "mutated" not in src_meta # item meta is a copy, not the source dict + + +@pytest.mark.parametrize("type_", ["audio", "video"]) +def test_a_preprocessor_cannot_declare_a_key_the_fetcher_writes(type_): + """The decoded file is the only authority on its own timeline, for every + modality alike. Letting ``update`` silently overwrite the declaration would + hide a preprocessor that believes something the file contradicts.""" + key = f"{type_}_metadata" + constructed = [["user", (type_, None, {key: "declared"})]] + media = {type_: [(object(), {key: "decoded"})]} + + with pytest.raises(ValueError, match=f"declares meta key\\(s\\) \\['{key}'\\]"): + _build_conversation_list(constructed, media) + + +def test_unused_media_is_refused_by_modality_and_count(): + # Pure sequential: 1 image turn but 2 tensors supplied → the leftover check + # must catch the count mismatch and name which modality drifted. + constructed = [["user", ("image", None)]] + with pytest.raises(ValueError, match="1 unused image media"): + _build_conversation_list(constructed, _media(image=_fake_images(2))) + + +def test_a_media_turn_with_nothing_left_to_pair_is_refused_by_modality(): + """2 image turns but only 1 tensor. Named ValueError rather than a leaked + StopIteration, which dataset iteration would read as a finished epoch.""" + constructed = [["user", ("image", None)], ["assistant", ("image", None)]] + with pytest.raises(ValueError, match="more image turn"): + _build_conversation_list(constructed, _media(image=_fake_images(1))) + + +# ── existing preprocessors tag image entries ─────────────────────────────────── + + +def test_imagenet1k_tags_assistant_image_gen(): + example = {"images": [object()]} + _, _, _, items = _build("imagenet1k", "cat", example) + types = [(it.type, it.role) for it in items] + assert types == [("text", "user"), ("image", "assistant")] + assert items[0].meta.get(_IMG_TAG_KEY) is None # text entry untagged + assert items[1].meta.get(_IMG_TAG_KEY) == "gen" + + +def test_sharegpt4v_tags_user_image_und(): + conv = [ + {"from": "human", "value": "describe this"}, + {"from": "gpt", "value": "a cat"}, + ] + example = {"images": [object()]} + _, _, _, items = _build("sharegpt4v_sft", conv, example) + img = next(it for it in items if it.type == "image") + assert img.role == "user" + assert img.meta.get(_IMG_TAG_KEY) == "und" + # Text entries (user instruction + assistant reply) stay untagged. + for it in items: + if it.type == "text": + assert _IMG_TAG_KEY not in it.meta + + +def test_llava_video_tags_user_video_und(): + conv = [ + {"from": "human", "value": "describe the clip"}, + {"from": "gpt", "value": "a running dog"}, + ] + example = {"videos": [object()]} + _, _, _, items = _build("llava_video", conv, example) + video = next(it for it in items if it.type == "video") + assert video.role == "user" + assert video.meta.get(_IMG_TAG_KEY) == "und" + for it in items: + if it.type == "text": + assert _IMG_TAG_KEY not in it.meta + + +def test_tulu_text_items_have_no_img_tag(): + conv = {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]} + example: dict = {} + _, _, _, items = _build("tulu-3-sft-mixture", conv, example) + assert all(it.type == "text" for it in items) + for it in items: + assert _IMG_TAG_KEY not in it.meta + assert it.meta == {} + + +def test_tulu_prepends_leading_system_message_to_first_user(): + conv = { + "messages": [ + {"role": "system", "content": "Answer like a pirate."}, + {"role": "user", "content": "Where is the treasure?"}, + {"role": "assistant", "content": "Arr, beyond the reef."}, + ] + } + + constructed, _, _, items = _build("tulu-3-sft-mixture", conv, {}) + + expected_user = "\nAnswer like a pirate.\n\nWhere is the treasure?" + assert constructed == [ + ["user", ("text", expected_user)], + ["assistant", ("text", "Arr, beyond the reef.")], + ] + assert [(item.role, item.value) for item in items] == [ + ("user", expected_user), + ("assistant", "Arr, beyond the reef."), + ] + assert conv["messages"][0] == {"role": "system", "content": "Answer like a pirate."} + assert conv["messages"][1] == {"role": "user", "content": "Where is the treasure?"} + + +@pytest.mark.parametrize( + "messages", + [ + [ + {"role": "user", "content": "hello"}, + {"role": "system", "content": "Answer briefly."}, + {"role": "assistant", "content": "hi"}, + ], + [ + {"role": "system", "content": "Answer briefly."}, + {"role": "system", "content": "Answer precisely."}, + {"role": "user", "content": "hello"}, + ], + ], +) +def test_tulu_rejects_unsupported_system_message_layout(messages): + with pytest.raises(ValueError, match="at most one system message, at the beginning"): + conv_preprocess("tulu-3-sft-mixture", {"messages": messages}, {}) + + +def test_tulu_rejects_system_message_without_user(): + messages = [ + {"role": "system", "content": "Answer briefly."}, + {"role": "assistant", "content": "hello"}, + ] + with pytest.raises(ValueError, match="must contain a user message"): + conv_preprocess("tulu-3-sft-mixture", {"messages": messages}, {}) + + +def test_preprocessors_return_matching_refs(): + # Non-edit preprocessors echo the sample's media refs unchanged so each + # fetcher sees a ref list whose length matches that modality's flattened + # entry count (no expansion needed). A modality the source has none of is + # an absent key, not an empty list — that is what lets a new modality be + # one new key rather than a wider tuple everywhere. + img_ref = object() + vid_ref = object() + + _, refs = conv_preprocess("imagenet1k", "cat", {"images": [img_ref]}) + assert refs == {"image": [img_ref]} + + conv = [{"from": "human", "value": "describe"}, {"from": "gpt", "value": "a cat"}] + _, refs = conv_preprocess("sharegpt4v_sft", conv, {"images": [img_ref]}) + assert refs == {"image": [img_ref]} + + _, refs = conv_preprocess("llava_video", conv, {"videos": [vid_ref]}) + assert refs == {"video": [vid_ref]} + + tulu = {"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "yo"}]} + _, refs = conv_preprocess("tulu-3-sft-mixture", tulu, {}) + assert refs == {} + + +# ── seed_edit_p23_multi_turn preprocessor ────────────────────────────────────── + + +def _edit_conv(n_targets: int) -> list[dict]: + # human source first, then (human instr, gpt target) pairs. n_targets gpt + # images ⇒ total tokens = 1 + n_targets. + conv = [{"from": "human", "value": "edit it to be red"}] + for i in range(n_targets): + conv.append({"from": "gpt", "value": ""}) + if i < n_targets - 1: + conv.append({"from": "human", "value": f"now make it color {i + 2}"}) + return conv + + +def test_seed_edit_multi_turn_chain_structure_three_targets(): + n_targets = 3 + conv = _edit_conv(n_targets) + refs = [f"img{i}" for i in range(1 + n_targets)] # 4 distinct string refs + example = {"images": refs} + constructed, refs_by_type = conv_preprocess("seed_edit_p23_multi_turn", conv, example) + image_refs = refs_by_type["image"] + + # Expected flattened chain: S0, instr1, T1, S1(copy), instr2, T2, S2(copy), + # instr3, T3 (final, no copy). + expected = [ + ("image", "user", "edit"), + ("text", "user", None), + ("image", "assistant", "gen"), + ("image", "assistant", "edit"), + ("text", "user", None), + ("image", "assistant", "gen"), + ("image", "assistant", "edit"), + ("text", "user", None), + ("image", "assistant", "gen"), + ] + actual = [ + (e[0], turn[0], e[2].get(_IMG_TAG_KEY) if len(e) == 3 else None) for turn in constructed for e in turn[1:] + ] + assert actual == expected + + # image_refs length matches the flattened image entry count, with copies + # duplicating the preceding target's ref: [img0, img1, img1, img2, img2, img3]. + n_image_entries = sum(1 for turn in constructed for e in turn[1:] if e[0] == "image") + assert len(image_refs) == n_image_entries == 6 + assert image_refs == [refs[0], refs[1], refs[1], refs[2], refs[2], refs[3]] + # An image-only source declares no other modality at all. + assert set(refs_by_type) == {"image"} + + +def test_seed_edit_multi_turn_image_refs_expansion_matches_primaries_and_copies(): + # Direct check of the ref-list expansion the preprocessor now owns: each + # primary gets example["images"][idx]; each copy re-appends the + # preceding target's ref (same object) without advancing the primary index. + n_targets = 3 + conv = _edit_conv(n_targets) + refs = [object() for _ in range(1 + n_targets)] # 4 distinct opaque refs + example = {"images": refs} + _, refs_by_type = conv_preprocess("seed_edit_p23_multi_turn", conv, example) + image_refs = refs_by_type["image"] + + assert image_refs == [refs[0], refs[1], refs[1], refs[2], refs[2], refs[3]] + # Copies are the very same object as the preceding target's ref. + assert image_refs[2] is refs[1] + assert image_refs[4] is refs[2] + # example["images"] is never mutated. + assert example["images"] == refs + assert len(example["images"]) == 1 + n_targets + + +def test_seed_edit_multi_turn_sequential_pairing_resolves_copies(): + # End-to-end through _build_conversation_list: pure sequential pairing of the + # expanded ref list. Copy items get the same content as their preceding + # target because the helper keys fake tensors on ref identity. + n_targets = 3 + conv = _edit_conv(n_targets) + example = {"images": [object() for _ in range(1 + n_targets)]} + _, _, _, items = _build("seed_edit_p23_multi_turn", conv, example) + img_items = [it for it in items if it.type == "image"] + + assert len(img_items) == 6 # 4 primaries + 2 copies + # source0 → fill 0; target_i → fill i; copy of target_i reuses the same ref. + assert int(img_items[0].value.float().mean()) == 0 # source0 + assert int(img_items[1].value.float().mean()) == 1 # target1 + assert torch.equal(img_items[2].value, img_items[1].value) # source1 = copy of target1 + assert int(img_items[3].value.float().mean()) == 2 # target2 + assert torch.equal(img_items[4].value, img_items[3].value) # source2 = copy of target2 + assert int(img_items[5].value.float().mean()) == 3 # target3 (final) + assert img_items[-1].meta.get(_IMG_TAG_KEY) == "gen" + + # Text items carry no _img_tag. + for it in items: + if it.type == "text": + assert _IMG_TAG_KEY not in it.meta + + +def test_seed_edit_multi_turn_last_target_has_no_copy_single_target(): + # Minimal non-trivial chain: source + exactly one target (the last), so no + # copy should be emitted at all. + conv = [ + {"from": "human", "value": "edit it"}, + {"from": "gpt", "value": ""}, + ] + example = {"images": [object(), object()]} + constructed, refs_by_type = conv_preprocess("seed_edit_p23_multi_turn", conv, example) + image_refs = refs_by_type["image"] + flat = [(e[0], turn[0], e[2].get(_IMG_TAG_KEY) if len(e) == 3 else None) for turn in constructed for e in turn[1:]] + assert flat == [ + ("image", "user", "edit"), + ("text", "user", None), + ("image", "assistant", "gen"), + ] + # No copy emitted for the final (only) target → image_refs == the two primaries. + assert len(image_refs) == 2 + assert image_refs == [example["images"][0], example["images"][1]] + assert set(refs_by_type) == {"image"} + + _, _, _, items = _build("seed_edit_p23_multi_turn", conv, example) + img_items = [it for it in items if it.type == "image"] + assert len(img_items) == 2 # no copy + assert img_items[-1].meta.get(_IMG_TAG_KEY) == "gen" + + +# ── fetch_images ref dedupe ──────────────────────────────────────────────────── + + +def test_fetch_images_dedupes_repeated_bytes_ref_and_clones(): + import io + + from PIL import Image + + from veomni.data.seed_omni.utils.image import fetch_images + + buf = io.BytesIO() + Image.new("RGB", (4, 4), (7, 7, 7)).save(buf, format="PNG") + ref = buf.getvalue() # same bytes object appended 3x + out = fetch_images([ref, ref, ref]) + assert len(out) == 3 + assert all(tuple(t.shape) == (3, 4, 4) for t in out) + # Same decoded content, but each is an independent tensor (clone on reuse). + assert torch.equal(out[0], out[1]) and torch.equal(out[1], out[2]) + assert out[0] is not out[1] and out[1] is not out[2] + out[0][0, 0, 0] = 255 + assert int(out[1][0, 0, 0]) == 7 # clone is independent of the first + + +def test_fetch_images_decodes_repeated_ref_once_and_returns_independent_clones(monkeypatch): + from PIL import Image + + from veomni.data.seed_omni.utils import image as image_utils + + calls = 0 + + def fake_load_image(ref): + nonlocal calls + calls += 1 + assert ref == "same-image" + return Image.new("RGB", (2, 2), (11, 11, 11)) + + monkeypatch.setattr(image_utils, "load_image", fake_load_image) + + out = image_utils.fetch_images(["same-image", "same-image", "same-image"]) + + assert calls == 1 + assert len(out) == 3 + assert all(torch.equal(out[0], item) for item in out[1:]) + out[0][0, 0, 0] = 255 + assert int(out[1][0, 0, 0]) == 11 + + +def test_fetch_images_unique_refs_each_decoded(): + import io + + from PIL import Image + + from veomni.data.seed_omni.utils.image import fetch_images + + buf1 = io.BytesIO() + Image.new("RGB", (4, 4), (7, 7, 7)).save(buf1, format="PNG") + ref1 = buf1.getvalue() + buf2 = io.BytesIO() + Image.new("RGB", (4, 4), (9, 9, 9)).save(buf2, format="PNG") + ref2 = buf2.getvalue() + out = fetch_images([ref1, ref2]) + assert int(out[0].float().mean()) == 7 + assert int(out[1].float().mean()) == 9 + + +def test_fetch_images_dedupes_repeated_pil_ref_by_identity(): + from PIL import Image + + from veomni.data.seed_omni.utils.image import fetch_images + + pil = Image.new("RGB", (4, 4), (5, 5, 5)) + out = fetch_images([pil, pil]) # same PIL object twice → identity-dedupe + assert len(out) == 2 + assert torch.equal(out[0], out[1]) + assert out[0] is not out[1] # clone on reuse + + +def test_fetch_videos_decodes_repeated_ref_once_and_returns_independent_copies(monkeypatch): + import numpy as np + + from veomni.data.seed_omni.utils import video as video_utils + + calls = 0 + + def fake_load_video(ref, **kwargs): + nonlocal calls + calls += 1 + assert ref == "same-video" + return ( + video_utils.VideoInputs( + video=torch.full((2, 3, 2, 2), 13, dtype=torch.uint8), + audio=np.ones((4,), dtype=np.float32), + ), + { + video_utils.VIDEO_METADATA_KEY: video_utils.VideoMetadata( + fps=30.0, total_num_frames=31, frames_indices=[0, 15], requested_fps=kwargs["fps"] + ), + video_utils.AUDIO_METADATA_KEY: video_utils.AudioMetadata(sampling_rate=16000, num_samples=4), + }, + ) + + monkeypatch.setattr(video_utils, "load_video", fake_load_video) + + out = video_utils.fetch_videos(["same-video", "same-video"], fps=3.0, use_audio_in_video=True) + + assert calls == 1 + assert len(out) == 2 + (first, first_meta), (second, second_meta) = out + assert first is not second + assert torch.equal(first.video, second.video) + assert first.audio is not second.audio + # The metadata is copied too: it lands on a per-item meta dict that modules + # write into, so a shared instance would let the first writer reach the other. + assert first_meta[video_utils.VIDEO_METADATA_KEY] is not second_meta[video_utils.VIDEO_METADATA_KEY] + assert first_meta[video_utils.VIDEO_METADATA_KEY] == second_meta[video_utils.VIDEO_METADATA_KEY] + + first.video[0, 0, 0, 0] = 255 + first.audio[0] = 7.0 + assert int(second.video[0, 0, 0, 0]) == 13 + assert float(second.audio[0]) == 1.0 + + +def test_multimodal_load_image_from_path_sets_http_timeout(monkeypatch) -> None: + import io + from types import SimpleNamespace + + from PIL import Image + + from veomni.data.multimodal import image_utils as mm_image_utils + + buf = io.BytesIO() + Image.new("RGB", (2, 2), color=(10, 20, 30)).save(buf, format="PNG") + png_bytes = buf.getvalue() + seen: dict = {} + + def fake_get(url, **kwargs): + seen["url"] = url + seen["kwargs"] = kwargs + return SimpleNamespace(content=png_bytes, raise_for_status=lambda: None) + + monkeypatch.setattr(mm_image_utils.requests, "get", fake_get) + + img = mm_image_utils.load_image_from_path("https://example.com/test.png") + assert img.size == (2, 2) + assert seen == {"url": "https://example.com/test.png", "kwargs": {"timeout": (5, 30)}} diff --git a/tests/seed_omni/data/test_preprocessor.py b/tests/seed_omni/data/test_preprocessor.py new file mode 100644 index 0000000000..9989c915dc --- /dev/null +++ b/tests/seed_omni/data/test_preprocessor.py @@ -0,0 +1,116 @@ +"""Tests for ``SeedOmniCollator`` and ``ModulePreprocessorBase`` contracts (#1063).""" + +from __future__ import annotations + +import pickle + +import pytest +import torch + +from scripts.ci.accelerator.launch_plan import LaunchSpec +from veomni.data.seed_omni.collator import SeedOmniCollator +from veomni.models.seed_omni.modules.module_processing_base import ModulePreprocessorBase +from veomni.models.seed_omni.processing_omni import OmniProcessor +from veomni.models.seed_omni.utils.conversation import ( + _IMG_TAG_KEY, + ConversationItem, + iter_desired_items, +) + + +pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) + + +class _DummyImageNormalizePreprocessor(ModulePreprocessorBase): + """Picklable test preprocessor that normalizes user images and appends a batch anchor.""" + + def __init__(self, dtype: torch.dtype = torch.bfloat16) -> None: + self.dtype = dtype + + def preprocess_conversations( + self, conversation_list: list[list[ConversationItem]], inference: bool = False, **kwargs + ) -> None: + del kwargs + user_images = [ + it for it in iter_desired_items(conversation_list, types=["image"], roles=["user"]) if not it.is_dummy + ] + for item in user_images: + item.value = item.value.to(dtype=self.dtype) / 255.0 + if not user_images and not inference: + self.append_batch_anchor( + conversation_list, + ConversationItem( + type="image", + value=torch.zeros(3, 4, 4, dtype=self.dtype), + role="user", + is_dummy=True, + meta={_IMG_TAG_KEY: "und"}, + ), + ) + + +def test_collator_strict_key_alignment_rejects_mismatched_sample_keys(): + collator = SeedOmniCollator(processor=OmniProcessor({})) + features = [ + {"conversation_list": [ConversationItem(type="text", value="a", role="user")], "ds_idx": 0}, + {"conversation_list": [ConversationItem(type="text", value="b", role="user")]}, + ] + with pytest.raises(ValueError, match="sample 1 has keys"): + collator(features) + + +def test_collator_preserves_mixed_resolution_images_without_stacking(): + """Mixed-resolution images survive collation unstacked and normalized per item.""" + pre = _DummyImageNormalizePreprocessor(dtype=torch.bfloat16) + reloaded_pre = pickle.loads(pickle.dumps(pre)) + collator = SeedOmniCollator(processor=OmniProcessor({"vision": reloaded_pre})) + + small_img = torch.randint(0, 255, (3, 14, 14), dtype=torch.uint8) + wide_img = torch.randint(0, 255, (3, 28, 42), dtype=torch.uint8) + tall_img = torch.randint(0, 255, (3, 56, 20), dtype=torch.uint8) + + features = [ + { + "conversation_list": [ + ConversationItem(type="image", value=small_img, role="user"), + ConversationItem(type="image", value=wide_img, role="user"), + ConversationItem(type="text", value="compare these", role="user"), + ], + "ds_idx": 0, + }, + { + "conversation_list": [ + ConversationItem(type="image", value=tall_img, role="user"), + ConversationItem(type="text", value="describe this", role="user"), + ], + "ds_idx": 1, + }, + ] + + batch = collator(features) + + assert set(batch.keys()) == {"conversation_list", "ds_idx"} + assert batch["ds_idx"] == [0, 1] + s0_imgs = [it.value for it in batch["conversation_list"][0] if it.type == "image"] + s1_imgs = [it.value for it in batch["conversation_list"][1] if it.type == "image"] + assert [tuple(t.shape) for t in s0_imgs] == [(3, 14, 14), (3, 28, 42)] + assert [tuple(t.shape) for t in s1_imgs] == [(3, 56, 20)] + assert all(t.dtype == torch.bfloat16 for t in [*s0_imgs, *s1_imgs]) + + +def test_module_preprocessor_base_appends_single_batch_anchor_when_modality_absent(): + pre = _DummyImageNormalizePreprocessor(dtype=torch.bfloat16) + collator = SeedOmniCollator(processor=OmniProcessor({"vision": pre})) + features = [ + {"conversation_list": [ConversationItem(type="text", value="hi", role="user")]}, + {"conversation_list": [ConversationItem(type="text", value="yo", role="user")]}, + ] + + batch = collator(features) + + dummies = [ + it for it in iter_desired_items(batch["conversation_list"], types=["image"], roles=["user"]) if it.is_dummy + ] + assert len(dummies) == 1 + assert dummies[0] is batch["conversation_list"][0][-1] + assert dummies[0].value.shape == (3, 4, 4) diff --git a/tests/seed_omni/model/test_conversation.py b/tests/seed_omni/model/test_conversation.py index c40ba78ab8..203029e00c 100644 --- a/tests/seed_omni/model/test_conversation.py +++ b/tests/seed_omni/model/test_conversation.py @@ -1,15 +1,195 @@ -"""``ConversationItem`` routing: items are selected by data tags, never by module.""" +"""Unit tests for SeedOmni conversation-list helpers; items are selected by data tags, never by module.""" + +from __future__ import annotations import pytest import torch from scripts.ci.accelerator.launch_plan import LaunchSpec -from veomni.models.seed_omni.utils.conversation import _IMG_TAG_KEY, ConversationItem, iter_desired_items +from veomni.data.seed_omni.utils.media_metadata import AudioMetadata +from veomni.models.seed_omni.utils.conversation import ( + _IMG_TAG_KEY, + ConversationItem, + build_conversation, + collect_desired_values, + get_tail_output_item, + iter_desired_items, + maybe_merge_outputs, + seal_outputs, +) pytestmark = pytest.mark.launch_plan(provider=LaunchSpec(device="cpu")) +def test_build_conversation_text_only_yields_user_part(): + parts = build_conversation(prompt="hello") + assert [(p.type, p.role, p.value) for p in parts] == [("text", "user", "hello")] + + +def test_build_conversation_with_images_places_them_first(): + img_a, img_b = object(), object() + parts = build_conversation(prompt="describe", media={"image": [img_a, img_b]}) + assert [(p.type, p.role) for p in parts] == [ + ("image", "user"), + ("image", "user"), + ("text", "user"), + ] + assert parts[0].value is img_a + assert parts[1].value is img_b + assert parts[2].value == "describe" + + +def test_build_conversation_with_audios_places_them_before_the_prompt(): + wav_a, wav_b = object(), object() + parts = build_conversation(prompt="transcribe", media={"audio": [wav_a, wav_b]}) + assert [(p.type, p.role) for p in parts] == [ + ("audio", "user"), + ("audio", "user"), + ("text", "user"), + ] + assert parts[0].value is wav_a + assert parts[1].value is wav_b + + +def test_build_conversation_keeps_the_mapping_order_and_ends_on_the_prompt(): + img, wav, clip = object(), object(), object() + parts = build_conversation(prompt="what is in these", media={"image": [img], "audio": [wav], "video": [clip]}) + assert [p.type for p in parts] == ["image", "audio", "video", "text"] + + +def test_build_conversation_carries_a_modality_it_was_never_taught(): + """A new modality is an entry in the fetcher table, not an edit to this function. + + Pinned because the alternative — a parameter per modality — is what made + ``video`` silently unsupported here while the data layer already decoded it. + """ + payload = object() + parts = build_conversation(prompt="act", media={"action": [payload]}) + assert [p.type for p in parts] == ["action", "text"] + assert parts[0].value is payload + + +def test_build_conversation_merges_fetcher_meta_onto_the_item(): + """The decode's facts (a clip's rate) have to survive into the item. + + Without this the inference path builds an audio item with no sampling rate — + a fact its consumers require and cannot recover from the samples. + """ + wav = object() + meta = {"audio_metadata": AudioMetadata(sampling_rate=24_000, num_samples=48_000)} + parts = build_conversation(prompt="transcribe", media={"audio": [(wav, meta)]}) + assert parts[0].value is wav + assert parts[0].meta["audio_metadata"].sampling_rate == 24_000 + + +def test_build_conversation_copies_the_meta_it_is_handed(): + """Two items decoded from one ref must not share a meta dict. + + Modules write into ``item.meta`` during forward, so a shared dict lets the + first writer reach the other item. + """ + meta = {"audio_metadata": AudioMetadata(sampling_rate=16_000, num_samples=1600)} + parts = build_conversation(prompt="hi", media={"audio": [(object(), meta), (object(), meta)]}) + assert parts[0].meta is not parts[1].meta + parts[0].meta["written_by_a_module"] = True + assert "written_by_a_module" not in parts[1].meta + + +def test_build_conversation_does_not_mistake_a_tuple_payload_for_a_pair(): + payload = (1, 2) + parts = build_conversation(prompt="hi", media={"image": [payload]}) + assert parts[0].value == payload + assert parts[0].meta == {} + + +def test_audio_items_are_filterable_like_any_other_modality(): + conversation_list = [ + [ + ConversationItem(type="audio", value=torch.zeros(16000), role="user"), + ConversationItem(type="text", value="hi", role="user"), + ], + [ConversationItem(type="audio", value=torch.zeros(8000), role="assistant")], + ] + user_audio = list(iter_desired_items(conversation_list, types=["audio"], roles=["user"])) + assert [tuple(item.value.shape) for item in user_audio] == [(16000,)] + all_audio = list(iter_desired_items(conversation_list, types=["audio"])) + assert len(all_audio) == 2 + + +def test_value_repr_shows_shape_for_array_like_waveforms(): + class FakeArray: + shape = (24000,) + dtype = "float32" + + item = ConversationItem(type="audio", value=FakeArray(), role="user") + assert " value: [FakeArray](24000,)\n" in repr(item) + + +def test_conversation_item_meta_defaults_empty(): + p = ConversationItem(type="text", value="", role="user") + assert p.meta == {} + + +def test_output_merge_and_seal(): + a = ConversationItem(type="output", value=torch.zeros(1, 2, 4), meta={"phase": "text"}) + b = ConversationItem(type="output", value=torch.ones(1, 1, 4), meta={"phase": "text"}) + parts = [a, b] + assert maybe_merge_outputs(parts) + assert len(parts) == 1 + assert parts[0].value.shape == (1, 3, 4) + seal_outputs(parts, new_type="text") + assert parts[0].type == "text" + assert "phase" in parts[0].meta + + +def test_iter_desired_items_allows_multiple_images_per_sample(): + img_a, img_b, img_c = object(), object(), object() + batch = [ + [ + ConversationItem(type="image", value=img_a, role="user"), + ConversationItem(type="image", value=img_b, role="user"), + ], + [ConversationItem(type="image", value=img_c, role="user")], + ] + items = list(iter_desired_items(batch, types=["image"], roles=["user"])) + assert [item.value for item in items] == [img_a, img_b, img_c] + assert collect_desired_values(batch, types=["image"], roles=["user"]) == [img_a, img_b, img_c] + + +def test_iter_desired_items_filters_meta_keys(): + batch = [ + [ + ConversationItem(type="output", value=torch.zeros(1), role="assistant", meta={}), + ConversationItem( + type="output", + value=torch.ones(1), + role="assistant", + meta={"flow_velocity_target": torch.zeros(1)}, + ), + ] + ] + + items = list(iter_desired_items(batch, types=["output"], meta_keys=["flow_velocity_target"])) + + assert len(items) == 1 + assert torch.equal(items[0].value, torch.ones(1)) + + +def test_get_tail_output_item_returns_latest_matching_output(): + first = ConversationItem(type="output", value=torch.zeros(1), role="assistant", meta={"a": 1}) + middle = ConversationItem(type="text", value="done", role="assistant") + second = ConversationItem(type="output", value=torch.ones(1), role="user", meta={"b": 1}) + third = ConversationItem(type="output", value=torch.full((1,), 2.0), role="assistant", meta={"a": 1}) + conversation = [first, middle, second, third] + + assert get_tail_output_item(conversation) is third + assert get_tail_output_item(conversation, meta_keys=["a"]) is third + assert get_tail_output_item(conversation, meta_keys=["b"]) is second + assert get_tail_output_item(conversation, roles=["user"]) is second + assert get_tail_output_item(conversation, meta_keys=["missing"]) is None + + def test_a_dummy_is_selected_by_the_same_filter_as_the_items_it_stands_in_for(): """An encoder reads real and placeholder rows with one filter and tells them apart by ``is_dummy``; a placeholder tagged for another encoder is not taken.""" diff --git a/veomni/data/dataset.py b/veomni/data/dataset.py index 34e2211c74..334fb06103 100755 --- a/veomni/data/dataset.py +++ b/veomni/data/dataset.py @@ -1858,11 +1858,12 @@ def build_weighted_multisource_dataset( train_path=source, seed=seed, transform=transform, + source_name=source_names[idx] if source_names is not None else None, split_by_node=split_by_node, shuffle=shuffle, dataset_repeat=bool(kwargs.get("dataset_repeat", False)), ) - for source in sources + for idx, source in enumerate(sources) ] return WeightedMultiSourceDataset( diff --git a/veomni/data/multimodal/image_utils.py b/veomni/data/multimodal/image_utils.py index 9ef93b6e95..aae5c00b41 100644 --- a/veomni/data/multimodal/image_utils.py +++ b/veomni/data/multimodal/image_utils.py @@ -15,6 +15,7 @@ import io import math +import os from io import BytesIO from typing import ByteString, List, Union @@ -88,9 +89,15 @@ def smart_resize( def load_image_from_path(image: str, **kwargs): if image.startswith("http://") or image.startswith("https://"): - response = requests.get(image, stream=True) + response = requests.get(image, timeout=(5, 30)) + # ``raise_for_status`` surfaces a 404 / 5xx with the URL + status code + # instead of letting PIL raise a downstream ``UnidentifiedImageError`` + # on the HTML error body, which obscures the real failure. + response.raise_for_status() image_obj = Image.open(BytesIO(response.content)) else: + if not os.path.exists(image): + raise FileNotFoundError(f"Image path does not exist: {image}") image_obj = Image.open(image) return image_obj.convert("RGB") @@ -117,8 +124,11 @@ def fetch_images(images: List[ImageInput], **kwargs): def save_image_tensors_to_file(image_tensors: torch.Tensor, output_path: str): - image_tensors = image_tensors * 255.0 - image_tensors = image_tensors.clamp(0, 255) + # `.round()` before `.to(torch.uint8)` is load-bearing — uint8 cast TRUNCATES + # toward zero (so 127.5 → 127, not 128). Without the explicit round the + # saved PNG drifts ±1 LSB from the float pixel intensity, which fails + # bit-exact parity against numpy / PIL reference encoders downstream. + image_tensors = (image_tensors * 255.0).clamp(0, 255).round() image_tensors = image_tensors.cpu().to(torch.uint8).numpy() image = Image.fromarray(image_tensors) image.save(output_path) diff --git a/veomni/data/seed_omni/seedomni_transform.py b/veomni/data/seed_omni/seedomni_transform.py index 5bb4e10a44..46442638c9 100644 --- a/veomni/data/seed_omni/seedomni_transform.py +++ b/veomni/data/seed_omni/seedomni_transform.py @@ -223,7 +223,7 @@ def process_seedomni_example( """ # Non-destructive read — datasets often share dict references and a # subsequent ``__getitem__`` would otherwise see the key gone. - source = example.get("source_name", kwargs.get("source_name")) + source = kwargs.get("source_name") or example.get("source") or example.get("source_name") if source is None: raise KeyError( "process_seedomni_example: sample is missing 'source_name' (and no fallback " diff --git a/veomni/data/seed_omni/utils/video.py b/veomni/data/seed_omni/utils/video.py index d3554729f1..c571109d3f 100644 --- a/veomni/data/seed_omni/utils/video.py +++ b/veomni/data/seed_omni/utils/video.py @@ -59,7 +59,9 @@ class VideoInputs: * ``video`` — kept-frame tensor ``(T, C, H, W)`` uint8. * ``audio`` — optional in-video waveform (``None`` when there is no audio - track or extraction was disabled). No module consumes it yet. + track or extraction was disabled). No module consumes it yet; the intended + wrapping and time-interleave are recorded in + ``docs/seed_omni/design/media.md``. The two timelines are *not* here. They ride on the item's ``ConversationItem.meta`` under :data:`VIDEO_METADATA_KEY` and diff --git a/veomni/models/seed_omni/modules/module_processing_base.py b/veomni/models/seed_omni/modules/module_processing_base.py index ceeb30c8dd..7c9d808b93 100644 --- a/veomni/models/seed_omni/modules/module_processing_base.py +++ b/veomni/models/seed_omni/modules/module_processing_base.py @@ -71,7 +71,12 @@ class ModulePreprocessorBase: tensors (no ``device=``). The main process's thin ``pre_forward`` does the single ``.to(device)``. * **In-place mutation.** ``__call__`` receives the collator ``batch`` dict - and mutates it in place. Subclasses decide which keys they read and write. + (must contain ``conversation_list`` as ``list[list[ConversationItem]]``). + The default path mutates those items' ``value`` / ``meta`` in place, and + the thin ``pre_forward`` / ``generate`` re-selects them with the same + ``type`` / ``role`` / ``meta`` filter to read the heavy work back. A packed preprocessor may override + ``__call__`` and write tensors onto the same dict instead of walking + items. * **Shared by training + inference.** Training runs it inside a collator (DataLoader worker); inference runs it once over the request before the FSM. The ``inference`` flag flips train/infer-only behaviour. @@ -79,9 +84,15 @@ class ModulePreprocessorBase: def __call__(self, batch: dict[str, Any], inference: bool = False, **kwargs: Any) -> None: """Run CPU prep on a collated ``batch`` dict (mutates it in place).""" + self.preprocess_conversations(batch["conversation_list"], inference=inference, **kwargs) + + def preprocess_conversations( + self, conversation_list: list[list[Any]], inference: bool = False, **kwargs: Any + ) -> None: raise NotImplementedError( - f"{type(self).__name__} must implement __call__(batch, inference=False, **kwargs) " - "and mutate the batch in place." + f"{type(self).__name__} must implement " + "preprocess_conversations(conversation_list, inference=False, **kwargs) " + "and mutate it in place." ) @classmethod @@ -113,6 +124,23 @@ def bind_dummy_inputs(self, config: Any, dtype: Any = None) -> None: del config, dtype return None + @staticmethod + def append_batch_anchor(conversation_list: list[list[Any]], item: Any) -> None: + """Attach ``item`` as this micro-batch's single FSDP anchor row. + + The anchor exists so a module whose real inputs are absent this step still + runs forward and backward, keeping its FSDP2 collectives in step with the + other ranks. One row does that: the module's graph hooks batch every row + their filter selects and flag the batch all-dummy, and every + downstream consumer filters dummies out. Appending one per *sample* + instead would put a full-size dummy image per sample through the tower — + on a batch with no real images for this module that was the single largest + term in the step (a 48-sample Janus T2I micro-batch spent more time in the + all-dummy SigLIP pass than in the language model). + """ + if conversation_list: + conversation_list[0].append(item) + def bind_module_assets( model: Any, diff --git a/veomni/models/seed_omni/utils/__init__.py b/veomni/models/seed_omni/utils/__init__.py index dd931bdd3d..17dda1cdf6 100644 --- a/veomni/models/seed_omni/utils/__init__.py +++ b/veomni/models/seed_omni/utils/__init__.py @@ -12,12 +12,26 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SeedOmni utilities: HF → split-checkpoint conversion.""" +"""SeedOmni utilities: conversation carrier + HF → split-checkpoint conversion.""" +from .conversation import ( + ConversationItem, + build_conversation, + collect_desired_values, + iter_desired_items, + maybe_merge_outputs, + seal_outputs, +) from .convert_registry import OMNI_CONVERT_REGISTRY, convert_checkpoint __all__ = [ + "ConversationItem", + "build_conversation", + "maybe_merge_outputs", + "seal_outputs", + "iter_desired_items", + "collect_desired_values", "OMNI_CONVERT_REGISTRY", "convert_checkpoint", ] diff --git a/veomni/models/seed_omni/utils/conversation.py b/veomni/models/seed_omni/utils/conversation.py index 2a342862f0..2eb74c35d4 100644 --- a/veomni/models/seed_omni/utils/conversation.py +++ b/veomni/models/seed_omni/utils/conversation.py @@ -1,4 +1,4 @@ -"""``ConversationItem`` — the single carrier object of SeedOmni V2. +"""``ConversationItem`` — the single carrier object of SeedOmni. The whole pipeline (training and inference) operates on one batched ``conversation_list`` (``list[list[ConversationItem]]`` for training, a flat @@ -19,8 +19,9 @@ ``"audio"`` is a standalone sound item (speech in, or speech out once sealed). Sound carried *inside* a video clip is not an ``"audio"`` item — it rides on the video item's own ``value`` (``VideoInputs.audio``) so that one clip stays - one item and the backbone can interleave both streams on a shared timeline. - Either way its rate is stated the same way, in ``meta["audio_metadata"]``. + one item and the backbone can interleave both streams on a shared timeline; + see ``docs/seed_omni/design/media.md``. Either way its rate is stated the + same way, in ``meta["audio_metadata"]``. * ``value`` — polymorphic: raw content (``str`` / PIL image / pixel tensor / ``(samples,)`` waveform) before encoding, an ``(L, D)`` / ``(1, L, D)`` embedding tensor after.