`
### 16. Algorithm-Specific Training Args
-`TrainingArguments` has algorithm-specific subclasses (`GRPOTrainingArguments`, `DPPOTrainingArguments`, `DPOTrainingArguments`, `DGPOTrainingArguments`, `NFTTrainingArguments`, `AWMTrainingArguments`, `CRDTrainingArguments`, `DiffusionOPDTrainingArguments`). The correct subclass is resolved by `get_training_args_class()` (registry in `hparams/training_args/_registry.py`). Adding a new algorithm requires adding a corresponding subclass and updating the resolver.
+`TrainingArguments` has algorithm-specific subclasses (`SFTTrainingArguments`,
+`OfflineDPOTrainingArguments`, `GRPOTrainingArguments`, `DPPOTrainingArguments`,
+`DPOTrainingArguments`, `DGPOTrainingArguments`, `NFTTrainingArguments`, `AWMTrainingArguments`,
+`CRDTrainingArguments`, `DiffusionOPDTrainingArguments`, and the multi-role distillation classes).
+The correct subclass is resolved by `get_training_args_class()`; adding an algorithm requires a
+corresponding subclass and registry entry.
### 17. YAML Config Structure
Config keys must exactly match Pydantic field names. Typos fail silently with default values. See `examples/` for canonical config templates; structure defined in `hparams/args.py`.
@@ -132,6 +165,13 @@ Config keys must exactly match Pydantic field names. Typos fail silently with de
### 18. All-Rank Synchronization Points
`accelerator.wait_for_everyone()` must be called at critical synchronization points (after preprocessing, before/after evaluation, checkpoint saving). Missing barriers cause deadlocks or race conditions.
+### 18a. Exact Resume Must Cover Replayed Evaluation
+Exact-state identity MUST lock evaluation cadence, sampling configuration, ordered realized eval
+loaders, per-dataset overrides, and eval rewards. Online checkpoints are saved before evaluation
+and replay that evaluation after resume, so treating eval as an operational control permits global
+device RNG drift. Exact-state save MUST fail before adapter or filesystem mutation on a device
+whose RNG Accelerate cannot serialize; MPS users must use model-only checkpoints.
+
### 19. FSDP CPU Efficient Loading
Distributed loading is owned by `ModelLoadCoordinator` and its `BackendLoadRuntime`. TARGET roots may use rank-zero/meta FSDP2 loading only when the adapter declares that capability. AUXILIARY and REWARD resources are materialized as full per-rank replicas; FSDP auxiliary roots receive a cached sampled-fingerprint check, while reward loading is isolated from target-only loading state. Trainer code must not manipulate FSDP loading environment variables or broadcast component weights directly.
diff --git a/.agents/knowledge/philosophy.md b/.agents/knowledge/philosophy.md
index 0cdc4398c..7744d9fd7 100644
--- a/.agents/knowledge/philosophy.md
+++ b/.agents/knowledge/philosophy.md
@@ -1,7 +1,8 @@
# Design Philosophy
-Flow-Factory is a simple, extensible RL fine-tuning framework for diffusion/flow-matching models.
-Models, algorithms, and rewards are decoupled via registries and base-class contracts.
+Flow-Factory is a simple, extensible online and offline fine-tuning framework for
+diffusion/flow-matching models. Models, algorithms, data acquisition, and rewards are decoupled via
+registries and typed contracts.
Both rollout and training run under Hugging Face Accelerate (DDP / DeepSpeed ZeRO-1-2 / FSDP
backends). The single most important invariant is **train-inference consistency**.
diff --git a/.agents/knowledge/topics/adapter_conventions.md b/.agents/knowledge/topics/adapter_conventions.md
index fc55add0c..e0c465dec 100644
--- a/.agents/knowledge/topics/adapter_conventions.md
+++ b/.agents/knowledge/topics/adapter_conventions.md
@@ -26,6 +26,35 @@ All adapters that support CFG must follow a consistent two-stage pattern. Guidan
`flux/flux2_klein.py` — `encode_prompt()` and `_forward()`.
+### Offline flow-matching guidance
+
+Finite-data SFT and offline DPO must not reuse `train.guidance_scale`: that field controls
+generation, and on a conventional CFG adapter it would turn the conditional velocity into a
+conditional/unconditional composite whenever negative embeddings are cached. These trainers
+expand the immutable `adapter.offline_training_forward_overrides` mapping into every SFT policy
+and offline-DPO policy/reference forward. The mapping is layered after both configured training
+arguments and dataset conditions, so adapter-owned model semantics always win.
+
+The base mapping sets `guidance_scale=1.0`, the conventional CFG-off point. An adapter replaces
+the complete mapping when its forward has different semantics or additional guidance branches:
+
+- Z-Image uses `guidance_scale=0.0` because its CFG threshold is `> 0.0`, with normalization and
+ truncation fixed to their neutral settings.
+- Guidance-distilled FLUX.1, FLUX.1-Kontext, and FLUX.2 use the official Diffusers training
+ condition `guidance_scale=3.5`; this is a learned model embedding, not classifier-free guidance.
+- The currently supported Flux2-Klein forward always passes `guidance=None` into its transformer,
+ so its `guidance_scale` remains conventional two-pass CFG and inherits the neutral `1.0`.
+- Wan T2V neutralizes both transformer stages with `guidance_scale=guidance_scale_2=1.0`.
+- SenseNova neutralizes text and image guidance together and disables CFG normalization.
+- Bagel replaces the base mapping with its actual `cfg_text_scale` / `cfg_img_scale` arguments;
+ it must not inherit an irrelevant `guidance_scale` key through its permissive `**kwargs`.
+
+Wan I2V and LTX2 remain behind explicit offline output-codec blockers. Wan I2V must mirror the two
+neutral Wan transformer scales before it is enabled. LTX2 must set video/audio CFG scales and
+modality scales to `1.0`, CFG rescale and STG scales to `0.0`, and neutralize its STG block
+selection with `spatio_temporal_guidance_blocks=None`. The mapping is adapter-owned model
+conditioning, never a sampling or algorithm knob.
+
### Models with model-specific CFG extensions
| Model | Extension | Notes |
@@ -177,6 +206,26 @@ LTX2 packs `[video|audio]` into one `(B, Seq, C)` sequence, so it resolves as PA
13. **SenseNova ragged I2I is per-sample, not NaViT-packed** — SenseNova-U1 1.0/1.5 accepts ordered, variable-size and variable-count reference images. Each sample's references become one variable-length NEO-Unify prefix and remain PIL across preprocessing, rollout and replay. A framework batch may contain several such samples, but `SenseNovaAdapter.inference()` / `forward()` iterate them and call `SenseNovaDenoiser` with B=1; unlike Bagel, independent samples are not concatenated into a packed attention sequence. The native model's `batch_size>1` path only expands one shared prompt/reference KV cache to generate multiple noises for the same condition and is not a ragged multi-sample batch API.
+14. **Offline forward overrides are model conditioning, not sampling CFG** — SFT and offline DPO expand the adapter's complete immutable `offline_training_forward_overrides` mapping into every policy and reference forward, after batch conditions and configured sampling arguments. Conventional CFG adapters use their CFG-off point; guidance-distilled adapters use the value expected by their learned guidance embedder; multi-branch adapters neutralize every active branch under its real forward argument names. Replace the base mapping rather than adding unrelated keys, and never infer it from cached negative embeddings or expose it as an algorithm knob.
+
+## Fix Records
+
+### Sampling CFG leaked into finite-data velocity matching
+
+- **Date**: 2026-08-28
+- **Symptom**: SFT and offline DPO could optimize a CFG-composite velocity when
+ `train.guidance_scale > 1.0`, while their target remained the conditional flow-matching
+ velocity.
+- **Root Cause**: The shared forward helper copied the generation-oriented training arguments
+ into offline forwards without an adapter-owned model-conditioning override.
+- **Fix**: Added the immutable `BaseAdapter.offline_training_forward_overrides` mapping, declared
+ complete model-specific neutral or guidance-distilled mappings (including Wan T2V, SenseNova,
+ and Bagel multi-branch CFG), and expanded it into every offline policy/reference forward after
+ configured and batch arguments.
+- **Lesson**: Generation controls and finite-data model conditioning may share a low-level
+ argument name but must have separate semantic owners.
+- **Related Constraint**: #7
+
## Cross-refs
- UP: `architecture.md` "Adapter Pattern", `constraints.md` #5 #11-12
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 00f56f916..110ffcfdb 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -86,6 +86,110 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Lesson**: Reflection is only a default. Dynamic preprocessors need an explicit cache contract.
- **Related Constraint**: N/A
+### Preserve contract-locked terminology during architecture rewrites
+- **Date**: 2026-08-28
+- **Symptom**: The SenseNova documentation regression test failed after the architecture table retained the correct behavior but dropped the contract phrase `ordered variable-count references`.
+- **Root Cause**: A broad documentation rewrite paraphrased a tested semantic distinction without checking the existing documentation contract.
+- **Fix**: `.agents/knowledge/architecture.md` now restores the exact phrase while retaining the grouped `images` and non-NaViT execution details; the documentation suite verifies both distinctions.
+- **Lesson**: Search documentation tests before rewriting architecture terminology, especially where wording distinguishes adapters with superficially similar multi-reference inputs.
+- **Related Constraint**: N/A
+
+### Distributed exact-state phases need synchronized failure and one publisher
+- **Date**: 2026-08-28
+- **Symptom**: One rank could fail exact-resume RNG/hash preflight while a peer entered `Accelerator.load_state()` and hung; concurrent saves could also pass destination preflight together and write the same staging directory.
+- **Root Cause**: Rank-local filesystem work was followed by raw barriers or core mutation without first gathering errors, and publication ownership was checked non-atomically before artifact creation.
+- **Fix**: `trainers/abc.py` now gathers all-rank errors after barrier-free path resolution, runtime preflight, and core load; commits runtime progress only after every core load succeeds; atomically elects the global publisher before core save; and keeps every filesystem claim until all publishers install their final directory. `trainers/common/runtime_identity.py` hashes FSDP wrap/state-dict topology and the full DeepSpeed batch/accumulation plan, while `runtime_state.py` validates per-device RNG topology and state installability. Runtime manifests are written only after Accelerator artifacts finish. Targeted multi-rank failure simulations cover preflight, load, and publisher races.
+- **Lesson**: A raw barrier is unsafe after rank-local I/O that can raise. Structure distributed checkpoints as preflight → synchronized error gather → backend mutation → synchronized error gather → manifest → atomic publication, acquire the publication claim before any process writes shared staging, and retain it until every filesystem reports success. Exact resume identity must cover backend topology, not only parameter names and shapes.
+- **Related Constraint**: #18
+
+### Validate nested media cardinality before flattening
+- **Date**: 2026-08-28
+- **Symptom**: Valid offline Flux1-Kontext batches shaped as `List[List[PIL]]` raised `TypeError` while checking whether a sample carried multiple condition images.
+- **Root Cause**: `_standardize_image_input()` flattened the nested batch before running its per-sample cardinality check, so the check called `len()` on each PIL image.
+- **Fix**: Flux1-Kontext now checks and warns on the original nested batch before selecting the first image, with a regression covering two offline single-image rows.
+- **Lesson**: Perform shape and cardinality validation at the boundary where that structure still exists; flattening destroys the evidence needed to validate it safely.
+- **Related Constraint**: N/A
+
+### Optional conditions require row-preserving empty sentinels
+- **Date**: 2026-08-28
+- **Symptom**: Legal batches mixing omitted and present negative prompts or condition images could reach tokenizers as `None`, lose their outer batch interpretation, or fail in an image encoder on an empty list.
+- **Root Cause**: Optionality was validated per record, but preprocessing and collation lacked homogeneous representations for a missing value inside a mixed batch.
+- **Fix**: Mixed optional negative prompts project missing values as empty strings; `is_multi_image_batch()` recognizes empty per-sample lists; Flux2/Klein emit aligned `None` latent/ID slots; and Bagel preserves empty condition-image slots. Arrow and offline-collator regressions cover the complete path.
+- **Lesson**: A batch-level optional field still needs one explicit slot per row. Normalize at the preprocessing boundary while retaining the original record identity for provenance.
+- **Related Constraint**: N/A
+
+### Preprocessing cache identity includes precision policy
+- **Date**: 2026-08-28
+- **Symptom**: Changing `component_load_dtypes` or `frozen_parameters_dtype` could reuse condition embeddings computed under a different component precision policy.
+- **Root Cause**: The offline cache fingerprint named the model but omitted the precision configuration that controls preprocessing component loading and storage.
+- **Fix**: Offline condition-cache extras now include a sorted canonical JSON representation of both dtype policies, with regressions for load-policy changes, frozen-policy changes, and mapping-order stability.
+- **Lesson**: Cache identity must include every configuration value that can change preprocessing numerics, even when that value is enforced during model loading rather than passed to the preprocessing function.
+- **Related Constraint**: #20
+
+### Batched Arrow schemas must cover later optional values
+- **Date**: 2026-08-28
+- **Symptom**: A valid optional-image dataset ordered as two prompt-only rows followed by two image-conditioned rows failed in the second `Dataset.map` chunk while casting image bytes or latent tensors to the first chunk's empty/string schema.
+- **Root Cause**: HuggingFace fixed the writer schema from the all-empty first output chunk, while Flux2 also omitted its image output keys whenever that individual chunk contained no images.
+- **Fix**: `data_utils/dataset.py` now scans real map-sized slices, dropping resolved columns from the bounded scan until it finds each later typed representative chunk; it probes those chunks only when a source column transitions from empty in the first chunk to typed later, restores Python/NumPy/torch/MPS and explicit-generator RNG state after that schema-only probe, and passes the resulting explicit `Features` through the `datasets==3.3.2`-compatible map surface. Flux2 now emits aligned image columns whenever the source image field exists; Flux2 and Flux2-Klein regressions verify identical empty-chunk output structure, a spy rejects whole-column reads, and a four-row offline cache regression covers the real Arrow path.
+- **Lesson**: A batched map's output key set and feature types are dataset-level contracts, not properties of whichever values happen to appear in the current chunk. Optional adapters must emit empty row slots consistently, and data writers must derive schemas from representative typed values before committing the first Arrow batch. A schema probe is intentionally narrow and restores RNG, but preprocessors still own their normal deterministic/cache-safe behavior; arbitrary adapter-owned mutable state is not rollback-safe.
+- **Related Constraint**: N/A
+
+### Exact resume must lock future execution and preserve acquisition boundaries
+- **Date**: 2026-08-28
+- **Symptom**: An exact checkpoint could pass preflight after objective, seed, scheduler, ordered training-data, or replayed evaluation semantics changed; a remote-rank runtime-child failure could let peers return from resume; an online resume retried its immutable source `checkpoint-N`; offline save-before-eval captured RNG that did not match the next uninterrupted epoch; and MPS could publish an exact checkpoint that its loader would always reject.
+- **Root Cause**: Resume identity stopped at physical model/optimizer/backend structure and treated RNG-consuming evaluation as operational, runtime-child commit was outside the synchronized distributed phase, one generic checkpoint/evaluation order ignored the different replay boundaries of generated versus finite-dataset acquisition, and save preflight did not reject a device RNG unsupported by Accelerate.
+- **Fix**: Runtime identity now includes trainer-extensible execution and rank-free data-contract digests derived from resolved objective/forward settings, realized training loader provenance/order/geometry, and the ordered realized evaluation path (cadence, arguments, dataset overrides, rewards, and prepared loaders). Runtime-child commit and attachment use a synchronized all-rank error phase. The first online boundary skips a save only when its resolved real path equals the exact-resume source, while still evaluating; offline boundaries evaluate before saving so exact checkpoints capture post-evaluation RNG, model-only saves retain the same observable order without claiming RNG restoration, and MPS exact save fails before adapter or filesystem mutation.
+- **Lesson**: Exact resume compatibility covers every computation and ordered data stream that can affect future state, including evaluation replay, not only the state container. Every rank-local resume mutation needs a synchronized failure boundary, checkpoint placement must match whether acquisition resumes before a rollout or after a completed data epoch, and save must not publish a state the matching load path cannot restore.
+- **Related Constraint**: #18
+
+### Multi-source schedule seeds must be process-independent
+- **Date**: 2026-08-28
+- **Symptom**: The same multi-source counts, configured seed, and epoch could produce a different source order on separate ranks or after restart.
+- **Root Cause**: `WeightedSourceBatchScheduler` seeded its generator with Python's salted `hash()` over a tuple containing a string, so the result depended on each process's `PYTHONHASHSEED`.
+- **Fix**: `data_utils/multi_source.py` now derives a domain-separated unsigned 64-bit seed from SHA-256; a subprocess regression compares schedules under distinct `PYTHONHASHSEED` values.
+- **Lesson**: Never use Python object hashes as distributed or persistent RNG seeds. Derive seeds from an explicitly versioned, stable byte representation.
+- **Related Constraint**: #9
+
+### Exact resume must include non-param-group optimizer semantics
+- **Date**: 2026-08-28
+- **Symptom**: An exact checkpoint could pass compatibility preflight after a role's gradient clipping threshold or update frequency changed.
+- **Root Cause**: Runtime identity covered realized optimizer parameter groups, but `max_grad_norm` and `update_frequency` are consumed by role optimization outside those groups.
+- **Fix**: `trainers/common/runtime_identity.py` now hashes resolved per-role optimizer arguments, including algorithm-provided defaults; runtime-identity regressions verify clipping and cadence drift change the execution digest while operational controls remain mutable.
+- **Lesson**: Exact-resume identity must cover every value that controls whether and how an optimizer update occurs, not only values serialized in optimizer parameter groups.
+- **Related Constraint**: #18
+
+### Distillation rollout cursors derive from persisted progress
+- **Date**: 2026-08-28
+- **Symptom**: DMD2, TDM, and TDM-R1 exact resumes restarted prompt acquisition from dataloader epoch zero even though the checkpoint recorded completed rollout iterations.
+- **Root Cause**: The trainers retained a live Python iterator and local dataloader epoch counter, while exact runtime state persisted only `TrainingProgress`; infinite grouped samplers never raised `StopIteration`, so the local epoch counter did not describe their real position either.
+- **Fix**: `trainers/distillation/distillation_runtime.py` now reconstructs the global consumed-batch cursor as `rollout_iteration * gradient_accumulation_steps`, uses the realized finite loader length before sampler/config fallbacks, maps the cursor to sampler epoch and intra-epoch offset, and restores Python/NumPy/torch CPU/CUDA/MPS plus explicit loader-generator RNG around iterator reconstruction and skips. Regressions compare uninterrupted and resumed real `DataLoader`, infinite grouped-sampler, and finite multi-source sequences.
+- **Lesson**: Do not serialize Python iterators or maintain a second checkpoint authority. At legal checkpoint boundaries, derive replayable loader position from persisted progress plus identity-locked realized geometry, and treat iterator construction and replayed skips as RNG-consuming side effects that must be neutralized.
+- **Related Constraint**: #18
+
+### Offline media bytes belong to the exact data identity
+- **Date**: 2026-08-28
+- **Symptom**: Replacing target, chosen, or rejected media in place left the same path-based record digest, so exact-resume preflight accepted a run whose next on-the-fly VAE inputs had changed. Input media replacement could likewise reuse stale condition embeddings.
+- **Root Cause**: Offline identities included normalized media type, path, and rate metadata but not file content.
+- **Fix**: `data_utils/offline_dataset.py` streams each unique normalized media path through SHA-256 once per source build. Input digests participate in condition IDs and cache fingerprints; supervision digests participate in full record IDs. No media, decoded pixels, or VAE latents are copied or cached.
+- **Lesson**: A path is provenance, not immutable content. Exact future-data contracts and preprocessing caches must identify external file bytes when those bytes are decoded on demand.
+- **Related Constraint**: #18
+
+### Runtime identity excludes transport-only global source IDs
+- **Date**: 2026-08-28
+- **Symptom**: Inserting or reordering an eval-only dataset renumbered training sources and caused exact-resume rejection even though the ordered training data and reward mathematics were unchanged.
+- **Root Cause**: Data identity hashed the full global name-to-ID registry and offline numeric `source_id`, which are transport metadata assigned across both train and eval entries.
+- **Fix**: `trainers/common/runtime_identity.py` now locks ordered realized training source names and loader schemas while excluding global numeric IDs. Structural and real-`Arguments` regressions verify eval-only changes preserve both execution and data digests, while training-source order/count changes remain incompatible.
+- **Lesson**: Compatibility identities should include semantic names and order, not remappable integer handles whose only purpose is runtime transport.
+- **Related Constraint**: #18
+
+### Offline epoch semantics follow the realized official sampler
+- **Date**: 2026-08-28
+- **Symptom**: Loader documentation claimed that every global sample appears exactly once and that an offline epoch is never padded, while the intentionally selected official `DistributedSampler(drop_last=False)` repeats tail indices when dataset size is not divisible by world size.
+- **Root Cause**: The design correctly defined an epoch as a complete rank-local dataloader traversal, but documentation conflated that with global sample uniqueness and with the separate prohibition on inventing batches to close a gradient-accumulation window.
+- **Fix**: Loader, workflow, dataset, sampler, and constraint documentation now preserve PyTorch's standard tail-equalization semantics and state that the framework adds no batches merely for accumulation. Existing uneven-geometry tests continue to lock official sampler behavior.
+- **Lesson**: When delegating sharding to an official sampler, define epoch semantics over its realized finite loader. Distinguish sampler-level repeated indices from optimizer-level synthetic padding.
+- **Related Constraint**: #9
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/.agents/knowledge/topics/samplers.md b/.agents/knowledge/topics/samplers.md
index ecfcd7165..c53b0e546 100644
--- a/.agents/knowledge/topics/samplers.md
+++ b/.agents/knowledge/topics/samplers.md
@@ -6,7 +6,10 @@
## Overview
-Flow-Factory uses **K-Repeat Sampling** (Stage 2 of the pipeline) to generate `K` copies of each unique prompt for group-wise advantage estimation. Three sampler implementations exist, differing in **how repeated samples are distributed across ranks**.
+Generation acquisition uses **K-Repeat Sampling** to create `K` copies of each unique prompt for
+group-wise advantage estimation. Three framework samplers differ in how repeated samples are
+distributed across ranks. Dataset acquisition is a separate path and uses PyTorch's official
+`DistributedSampler` without K-repeat geometry.
| Property | DistributedKRepeatSampler | GroupContiguousSampler | GroupDistributedSampler |
|----------|--------------------------|----------------------|------------------------|
@@ -16,6 +19,16 @@ Flow-Factory uses **K-Repeat Sampling** (Stage 2 of the pipeline) to generate `K
| **Auto-adjustment** | GCD-based rounding | LCM-based rounding (stricter) | O(√B) divisor search (`_align_for_group_distributed`) |
| **Use case** | Fallback when geometric constraints for group_contiguous are unsatisfied | Default when constraints are met (minimal communication) | DGPO — rank-identical prompt contract for local `torch.unique` |
+### Offline Dataset Sampler
+
+SFT and offline DPO always use `torch.utils.data.DistributedSampler`, call
+`set_epoch(data_epoch)`, and exhaust the finite loader once per data epoch. Every source weight is
+`1`; `gradient_accumulation_steps` is explicit; and each rank's batch count must be divisible by
+it. The loader is not passed to `Accelerator.prepare()`. With the official sampler's default
+`drop_last=False`, a non-divisible global tail is repeated deterministically to equalize rank
+lengths; the complete resulting loader traversal, rather than global sample uniqueness, defines
+the epoch.
+
---
## How Each Sampler Works
@@ -303,7 +316,7 @@ These are caught at sampler construction time. The auto-adjustment in `_align_ba
## Impact on Other Components
-- **Constraint #9 in [`../constraints.md`](../constraints.md)**: The dataloader is NOT prepared via `accelerator.prepare()` — both samplers handle distribution themselves.
+- **Constraint #9 in [`../constraints.md`](../constraints.md)**: No train dataloader is prepared via `accelerator.prepare()`; its selected grouped or official sampler owns distribution.
- **RewardProcessor**: When GroupContiguousSampler is active, groupwise rewards can be computed locally per rank. When DistributedKRepeatSampler is active, the RewardProcessor must gather group members across ranks.
- **AdvantageProcessor**: Automatically skips `accelerator.gather()` calls when `sampler_type == "group_contiguous"` (all group members already local); uses `all_reduce(count, sum, sum_sq)` for global_std (3 scalars). When `sampler_type == "distributed_k_repeat"`, packs all rewards + unique_ids into a single tensor for one `accelerator.gather()` call.
@@ -320,5 +333,5 @@ data:
## Cross-refs
- `constraints.md` #9, #9a (accelerator prepare scope, sampler geometric constraints)
-- `architecture.md` "Six-Stage Training Pipeline" (Stage 2: K-Repeat Sampling)
+- `architecture.md` "Execution Pipelines" (generation acquisition)
- `architecture.md` "Advantage Computation" (communication path depends on sampler type)
diff --git a/.agents/skills/ff-new-algorithm/SKILL.md b/.agents/skills/ff-new-algorithm/SKILL.md
index 83ff6d7f6..ecb46f37f 100644
--- a/.agents/skills/ff-new-algorithm/SKILL.md
+++ b/.agents/skills/ff-new-algorithm/SKILL.md
@@ -1,18 +1,21 @@
---
name: ff-new-algorithm
-description: "Complete workflow for adding a new RL training algorithm. Covers paradigm selection, TrainingArguments subclass, trainer implementation, registry, example config, and verification. Trigger: 'add algorithm', 'new trainer', 'new training method', 'implement algorithm'."
+description: "Complete workflow for adding an online or offline training algorithm. Covers execution-contract and paradigm selection, TrainingArguments subclass, trainer implementation, registry, example config, and verification. Trigger: 'add algorithm', 'new trainer', 'new training method', 'implement algorithm'."
---
-# New RL Algorithm Integration
+# New Training Algorithm Integration
> **Authoritative reference**: `guidance/algorithms.md`
## Prerequisites
Determine your algorithm's characteristics:
+- **Acquisition**: Generated rollouts or a finite dataset? (`generation` / `dataset`)
+- **Feedback**: Runtime reward/advantage or none? (`runtime_reward` / `none`)
- **Paradigm**: Coupled or Decoupled? (`constraints.md` #7)
- **Dynamics**: Which SDE/ODE formulation? (`Flow-SDE`, `Dance-SDE`, `CPS`, `ODE`)
-- **Advantage**: How are advantages computed from rewards? (Most algorithms can delegate to `AdvantageProcessor`)
+- **Supervision**: Prompt-only generation, demonstrations, preference pairs, or a new typed record?
+- **Advantage**: If feedback is enabled, how are advantages computed? (Most reward-based algorithms can delegate to `AdvantageProcessor`)
- **Loss**: What is the policy optimization objective?
## Phase 1: Design
@@ -20,11 +23,25 @@ Determine your algorithm's characteristics:
1. **Study existing implementations**:
- Coupled example: `trainers/rl/grpo.py` (GRPO)
- Decoupled example: `trainers/rl/nft.py` (DiffusionNFT) or `trainers/rl/awm.py` (AWM)
+ - Finite demonstration example: `trainers/offline/sft.py` (SFT)
+ - Finite preference example: `trainers/offline/offline_dpo.py` (offline DPO)
2. **Identify what's shared vs unique** (`constraints.md` #11):
- - Shared: the epoch loop (`BaseTrainer.start`), data loading, reward computation,
- `AdvantageProcessor`, `prepare_feedback`, `compute_advantages`, adapter interface, checkpoint logic
+ - Shared: the cycle loop (`BaseTrainer.start`), acquisition dispatch, progress counters,
+ adapter interface, checkpoint/eval boundaries, role optimization, and exact-resume identity
+ - Conditional: runtime rewards, `AdvantageProcessor`, `prepare_feedback`, and
+ `compute_advantages` exist only when the feedback contract requests them
- Unique: the loss function and the algorithm-specific hyperparameters. Never restate the loop
- - Per-epoch hook order: `sample()` → `prepare_feedback()` → `optimize()` (see `guidance/workflow.md`)
+ - Generation hook order: `sample()` → optional `prepare_feedback()` → `optimize()`
+ - Dataset hook order: official finite loader traversal → `optimize_batch(batch)`
+3. **Declare one immutable execution contract**:
+ - Online RL: `ONLINE_EXECUTION_CONTRACT` (`generation + runtime_reward`)
+ - Generation without rewards: `ONLINE_NO_FEEDBACK_EXECUTION_CONTRACT`
+ - Finite offline training: `OFFLINE_EXECUTION_CONTRACT` (`dataset + none`)
+
+Keep execution semantics orthogonal to `PipelineIOContract`. The algorithm owns how examples are
+acquired and optimized; the adapter owns accepted input/output media, geometry, and output-state
+encoding. A new offline record shape belongs in the typed data layer, never in model-specific loss
+branches.
## Phase 2: Configuration
@@ -95,25 +112,30 @@ from .training_args import MyAlgoTrainingArguments
### Step 3 — Create Trainer Class
+For a generated online algorithm:
+
```python
-# src/flow_factory/trainers/rl/my_algo.py
+# src/flow_factory/trainers/rl/my_online_algo.py
+from ...contracts import ONLINE_EXECUTION_CONTRACT
from ..abc import BaseTrainer
from ..registry import register_trainer
-@register_trainer('my_algo')
-class MyAlgoTrainer(BaseTrainer):
- """My custom RL algorithm trainer."""
+@register_trainer("my-online-algo")
+class MyOnlineAlgoTrainer(BaseTrainer):
+ """My generated-acquisition algorithm."""
+
+ execution_contract = ONLINE_EXECUTION_CONTRACT
- # Do NOT define start(). BaseTrainer.start() owns the epoch loop: reseed, checkpoint on
- # save_freq, evaluate on eval_freq, _run_training_step(), ema_step, _after_optimizer_step.
- # evaluate(), prepare_feedback() and compute_advantages() are likewise CONCRETE base
- # methods. optimize() is the only abstract one.
+ # Do NOT define start(). BaseTrainer.start() owns the acquisition loop: reseed,
+ # periodic boundaries, acquisition dispatch, EMA, and _after_acquisition_cycle().
+ # evaluate(), prepare_feedback(), and compute_advantages() are concrete base methods.
+ # A generation trainer implements sample() and optimize(samples).
#
# Vary behavior through hooks instead of restating the loop:
# sampling_context() - wrap the rollout (e.g. install a snapshot's weights)
# _run_training_step() - replace the sample -> feedback -> optimize middle
# _after_gradient_step() - run right after each optimizer step
- # _after_optimizer_step() - run once per epoch, after the EMA step
+ # _after_acquisition_cycle() - run once per rollout iteration or data epoch
# _declare_model_variants() - declare several trainable copies (see component_variants.md)
def sample(self):
@@ -129,8 +151,38 @@ class MyAlgoTrainer(BaseTrainer):
pass
```
-> **Note**: `AdvantageProcessor` is auto-instantiated in `BaseTrainer._init_reward_model()`.
-> Reward-based trainers delegate via `self.advantage_processor.compute_advantages()` — see `architecture.md` "Advantage Computation". (Pure-distillation trainers like `diffusion-opd` skip rewards/advantages with a no-op `prepare_feedback()`.)
+For a finite offline algorithm:
+
+```python
+# src/flow_factory/trainers/offline/my_offline_algo.py
+from ...contracts import OFFLINE_EXECUTION_CONTRACT
+from ..abc import BaseTrainer
+from ..registry import register_trainer
+
+@register_trainer("my-offline-algo")
+class MyOfflineAlgoTrainer(BaseTrainer):
+ """My finite-dataset algorithm."""
+
+ paradigm = "decoupled"
+ execution_contract = OFFLINE_EXECUTION_CONTRACT
+
+ def _build_train_dataloader(self):
+ """Build a finite official-DistributedSampler loader for one typed schema."""
+ # Reuse build_offline_train_dataloader when the supervision type matches;
+ # otherwise extend the typed schema/collator first.
+ ...
+
+ def optimize_batch(self, batch):
+ """Apply one gradient-accumulation microstep from a dataset batch."""
+ # Decode output media in the dataset and encode it on demand through
+ # adapter.encode_output_state(); never add target VAE latents to the cache.
+ ...
+```
+
+> **Note**: `AdvantageProcessor` is relevant only to `runtime_reward` feedback.
+> Reward-based trainers delegate via `self.advantage_processor.compute_advantages()` — see
+> `architecture.md` "Advantage Computation". `none` feedback bypasses rewards structurally; do not
+> emulate that by overriding reward methods with incidental no-ops.
### Step 4 — Register in Trainer Registry
@@ -188,9 +240,12 @@ optimizers:
- [ ] `MyAlgoTrainingArguments` correctly parsed from YAML
- [ ] `get_training_args_class('my_algo')` returns correct subclass
- [ ] `get_trainer_class('my_algo')` loads `MyAlgoTrainer`
-- [ ] Training runs end-to-end for ≥2 epochs without errors
+- [ ] `execution_contract` matches the argument class and implemented optimization hook
+- [ ] Training runs end-to-end for ≥2 acquisition cycles without errors
+- [ ] Dataset acquisition defines one epoch as one complete finite dataloader traversal
- [ ] Loss values are numerically reasonable (not NaN, decreasing)
-- [ ] Rewards improve over training
+- [ ] Rewards improve over training when feedback is `runtime_reward`
+- [ ] Offline supervision media is encoded on the fly and excluded from preprocessing caches
- [ ] Checkpoint save/load works correctly
- [ ] Works with at least two different model adapters
- [ ] Coupled algorithms only use SDE dynamics
@@ -206,3 +261,6 @@ optimizers:
6. **Reimplementing advantage gather/scatter** — use `self.advantage_processor.compute_advantages()` instead; it handles both sampler topologies automatically
7. **Extending `GRPOTrainer` unnecessarily** — unless your algorithm extends GRPO's PPO-clipped loss, extend `BaseTrainer` directly (as NFT and AWM do)
8. **Optimizer-time CFG without `get_preprocess_guidance_scale()`** — if your algorithm calls `adapter.forward(guidance_scale=X)` where X > 1.0 but `training_args.guidance_scale` ≤ 1.0, negative prompts won't be encoded at preprocessing time and CFG silently falls back to no-CFG. Override `get_preprocess_guidance_scale()` in your TrainingArguments subclass to return `max(guidance_scale, your_optimize_cfg)`. See DGPO's `kl_cfg` for a real example.
+9. **Inferring online/offline behavior from batch keys** — declare `ExecutionContract`; keep acquisition and feedback independent from the model I/O schema.
+10. **Using `optimize()` for finite data** — dataset acquisition calls `optimize_batch(batch)` and advances `data_epoch` only after clean loader exhaustion.
+11. **Caching target/chosen/rejected latents** — cache prompt/input conditions only; output media is decoded and encoded on demand through the adapter output codec.
diff --git a/AGENTS.md b/AGENTS.md
index 06cbb0fa3..ce79bf0e8 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -2,9 +2,9 @@
## Project Overview
-Flow-Factory is a unified **online RL fine-tuning framework** for diffusion/flow-matching models. It provides a modular architecture where trainers, model adapters, and reward models are independently extensible via a registry-based plugin system.
+Flow-Factory is a unified **online and offline fine-tuning framework** for diffusion/flow-matching models. It provides a modular architecture where trainers, model adapters, data acquisition, and reward models are independently extensible through typed contracts and registries.
-- **Algorithms**: GRPO, GRPO-Guard, DPPO, DPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1
+- **Algorithms**: SFT, offline DPO, online DPO, GRPO, GRPO-Guard, DPPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1
- **Models**: FLUX.1 (+Kontext), FLUX.2 (+Klein), SD3.5, Qwen-Image (+Edit-Plus), Z-Image, Wan2 (T2V/I2V), LTX2 (T2AV/I2AV), Bagel, SenseNova-U1 (1.0/1.5; T2I + ordered multi-reference I2I)
- **Rewards**: PickScore (+Rank), CLIP, CLAP, ImageBind, OCR, GenEval/GenEval2, HPSv2, VLM-Evaluate, rational-rewards, and custom rewards
- **Python**: >=3.10 | **PyTorch**: >=2.6.0 | **License**: Apache-2.0
@@ -59,8 +59,8 @@ See `.agents/knowledge/architecture.md` "Module Dependency Graph" for full detai
| Document | Purpose |
|----------|---------|
-| `guidance/workflow.md` | 6-stage training pipeline with code examples |
-| `guidance/algorithms.md` | All algorithms (GRPO, GRPO-Guard, DPPO, DPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1) deep dive |
+| `guidance/workflow.md` | Unified generation/dataset acquisition contracts plus the online 6-stage pipeline |
+| `guidance/algorithms.md` | All algorithms (SFT, offline DPO, GRPO, GRPO-Guard, DPPO, online DPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1) deep dive |
| `guidance/rewards.md` | Reward system design, custom model creation |
| `guidance/new_model.md` | Step-by-step model adapter integration |
| `guidance/acceleration.md` | Acceleration plugin layer (compile, attention backend, feature caching) |
@@ -76,7 +76,7 @@ Skills follow the [Agent Skills](https://agentskills.io) open standard. Each ski
| `/ff-review` | Pre-commit code review | Before committing changes |
| `/ff-new-model` | Model adapter integration | Adding support for a new diffusion model |
| `/ff-new-reward` | Reward model integration | Adding a new reward function |
-| `/ff-new-algorithm` | RL algorithm integration | Adding a new training algorithm |
+| `/ff-new-algorithm` | Online/offline algorithm integration | Adding a new training algorithm |
### Quick Decision Guide
diff --git a/README.md b/README.md
index 60f43ec0a..8e0b9ce45 100644
--- a/README.md
+++ b/README.md
@@ -4,7 +4,7 @@
Flow-Factory
- Easy Reinforcement Learning for Diffusion and Flow-Matching Models
+ Unified Online RL and Offline Fine-Tuning for Diffusion and Flow-Matching Models
# 🔥 News
@@ -45,6 +45,7 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
- [Quick Start Example](#quick-start-example)
- [Guidance](#-guidance)
- [Dataset](#-dataset)
+ - [Offline SFT and Preference Data](#offline-sft-and-preference-data)
- [Text-to-Image & Text-to-Video](#text-to-image--text-to-video)
- [Image-to-Image & Image-to-Video](#image-to-image--image-to-video)
- [Reward Model](#-reward-model)
@@ -96,6 +97,12 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
> To support new models, see [Guidance/New Model](guidance/new_model.md).
+> **Offline output support:** SFT and offline DPO currently support `sd3-5`, `flux1`,
+> `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`,
+> `bagel`, `sensenova`, and `wan2_t2v`. Wan I2V, LTX2, and MiniMax H3 fail fast on their
+> currently unresolved output/condition or audio-video semantics. See the
+> [offline model matrix](guidance/datasets.md#offline-model-support).
+
> **MiniMax H3 status:** the T2VA debug and
> [native-quality FSDP2](examples/grpo/lora/minimax_h3_t2va/quality_720p_fsdp2.yaml)
> paths are real-weight
@@ -107,9 +114,11 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
# 💻 Supported Algorithms
-| Algorithm | `trainer_type` | Paper |
+| Algorithm | `trainer_type` | Reference / objective |
|----------------|----------------|-------|
-| DPO | dpo | [Diffusion-DPO](https://arxiv.org/abs/2311.12908) |
+| SFT | sft | Supervised flow matching over V2 demonstrations |
+| Offline DPO | offline-dpo | [Diffusion-DPO](https://arxiv.org/abs/2311.12908) over V2 preference pairs |
+| Online DPO | dpo | [Diffusion-DPO](https://arxiv.org/abs/2311.12908) with generated, reward-ranked pairs |
| GRPO | grpo | [Flow-GRPO](https://arxiv.org/abs/2505.05470) / [Dance-GRPO](https://arxiv.org/abs/2505.07818) |
| DiffusionNFT | nft | [DiffusionNFT](https://arxiv.org/abs/2509.16117) |
| AWM | awm | [Advantage Weighted Matching](https://arxiv.org/abs/2509.25050) |
@@ -186,6 +195,13 @@ Start training with the following simple command:
ff-train examples/grpo/lora/flux1/default.yaml
```
+Offline smoke recipes use strict V2 manifests and require no training reward model:
+
+```bash
+ff-train examples/sft/lora/sd3_5/default.yaml
+ff-train examples/offline_dpo/lora/sd3_5/default.yaml
+```
+
# 📖 Guidance
We provide a set of guidance documents to help you understand the framework and extend it. For a comprehensive understanding of the framework's design and motivation, refer to our [technique report](https://arxiv.org/abs/2602.12529).
@@ -193,7 +209,7 @@ We provide a set of guidance documents to help you understand the framework and
| Document | Description |
|---|---|
| [Workflow](guidance/workflow.md) | End-to-end training pipeline: the overall stages from data preprocessing to policy optimization |
-| [Algorithms](guidance/algorithms.md) | Supported algorithms (GRPO, GRPO-Guard, DPPO, DiffusionNFT, AWM, DPO, DGPO, CRD, DiffusionOPD, DMD2, TDM, TDM-R1) and their configurations |
+| [Algorithms](guidance/algorithms.md) | Supported online RL, SFT, offline DPO, and distillation algorithms and their configurations |
| [Rewards](guidance/rewards.md) | Reward model system: built-in models, custom rewards, and remote reward servers |
| [Datasets](guidance/datasets.md) | Dataset schemas, media paths, and ordered-reference inputs |
| [New Model](guidance/new_model.md) | How to add support for a new Diffusion/Flow-Matching model |
@@ -214,6 +230,21 @@ The unified structure of dataset is:
|----|---| ...
```
+## Offline SFT and Preference Data
+
+SFT and offline DPO use strict JSONL with `schema_version: 2`. Public media objects always use the
+`type` discriminator; `kind` is not accepted in V2:
+
+```jsonl
+{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"targets/poster.png"}]}},"metadata":{}}
+{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"pairs/chosen.png"}]},"rejected":{"media":[{"type":"image","path":"pairs/rejected.png"}]}},"metadata":{}}
+```
+
+Prompt and input-condition encodings are cached. Target, chosen, and rejected media are decoded and
+encoded on the fly; their VAE latents are never stored in the preprocessing cache. One offline
+epoch is one complete dataloader traversal sharded by PyTorch's official `DistributedSampler`. See the
+[dataset guide](guidance/datasets.md#offline-v2-records) for the full schema and cadence rules.
+
## Text-to-Image & Text-to-Video
For text-to-image and text-to-video tasks, the only required input is the **prompt** in plain text format. Use `train.txt` and `test.txt` (optional) with following format:
diff --git a/examples/README.md b/examples/README.md
index 7ae385850..29e5d8e80 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -10,7 +10,7 @@ examples/{algorithm}/{finetune_type}/{model_type}/{variant}.yaml
| Level | Description | Examples |
|-------|-------------|---------|
-| `algorithm` | Training algorithm | `grpo`, `dppo`, `nft`, `awm`, `dgpo`, `dpo`, `crd`, `opd`, `dmd2`, `tdm`, `tdm_r1` |
+| `algorithm` | Training algorithm | `sft`, `offline_dpo`, `grpo`, `dppo`, `nft`, `awm`, `dgpo`, `dpo`, `crd`, `opd`, `dmd2`, `tdm`, `tdm_r1` |
| `finetune_type` | Parameter-efficient or full | `lora`, `full` |
| `model_type` | Model family (underscore-separated) | `flux1`, `sd3_5`, `wan21`, `ltx2` |
| `variant` | Config variant | `default.yaml`, `nocfg.yaml`, `t2v.yaml` |
@@ -24,6 +24,20 @@ examples/{algorithm}/{finetune_type}/{model_type}/{variant}.yaml
ff-train examples/grpo/lora/flux1/default.yaml
```
+## Offline examples
+
+- [`sft` with SD3.5](sft/lora/sd3_5/default.yaml) consumes V2
+ `demonstration` records from [`examples/data/sft_sd3_5`](data/sft_sd3_5/train.jsonl).
+- [`offline-dpo` with SD3.5](offline_dpo/lora/sd3_5/default.yaml) consumes V2
+ `preference` records from
+ [`examples/data/offline_dpo_sd3_5`](data/offline_dpo_sd3_5/train.jsonl).
+
+The two tiny manifests reuse repository images so their paths resolve without a separate dataset
+download. They are configuration and smoke-test fixtures, not quality-training datasets. Offline
+training requires an explicit integer `gradient_accumulation_steps`; the number of rank-local
+dataloader batches must be divisible by it. See the [dataset guide](../guidance/datasets.md#offline-v2-records)
+for the production schema and media requirements.
+
## DMD2 and TDM
- [`dmd2` SD3.5 OCR recipe](dmd2/lora/sd3_5/ocr.yaml) — validated in a
diff --git a/examples/data/offline_dpo_sd3_5/train.jsonl b/examples/data/offline_dpo_sd3_5/train.jsonl
new file mode 100644
index 000000000..f8c47b544
--- /dev/null
+++ b/examples/data/offline_dpo_sd3_5/train.jsonl
@@ -0,0 +1,2 @@
+{"schema_version":2,"input":{"prompt":"Render a clean Flow-Factory emblem.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"../../../assets/logo-no-bg.png"}]},"rejected":{"media":[{"type":"image","path":"../../../assets/wandb_metrics.png"}]}},"metadata":{"fixture":"repository-images"}}
+{"schema_version":2,"input":{"prompt":"Render an experiment dashboard.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"../../../assets/wandb_images.png"}]},"rejected":{"media":[{"type":"image","path":"../../../assets/logo.png"}]}},"metadata":{"fixture":"repository-images"}}
diff --git a/examples/data/sft_sd3_5/train.jsonl b/examples/data/sft_sd3_5/train.jsonl
new file mode 100644
index 000000000..0e4b77fc1
--- /dev/null
+++ b/examples/data/sft_sd3_5/train.jsonl
@@ -0,0 +1,2 @@
+{"schema_version":2,"input":{"prompt":"Render a clean Flow-Factory emblem.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"../../../assets/logo-no-bg.png"}]}},"metadata":{"fixture":"repository-image"}}
+{"schema_version":2,"input":{"prompt":"Render an experiment dashboard.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"../../../assets/wandb_images.png"}]}},"metadata":{"fixture":"repository-image"}}
diff --git a/examples/offline_dpo/lora/sd3_5/default.yaml b/examples/offline_dpo/lora/sd3_5/default.yaml
new file mode 100644
index 000000000..1a1698785
--- /dev/null
+++ b/examples/offline_dpo/lora/sd3_5/default.yaml
@@ -0,0 +1,72 @@
+# Single-process offline-DPO smoke recipe over a V2 preference manifest.
+# Replace the tiny repository-image dataset with real preference pairs for quality runs.
+launcher: "accelerate"
+config_file: null
+num_processes: 1
+main_process_port: 29501
+mixed_precision: "bf16"
+
+data:
+ datasets:
+ - name: offline_preferences
+ dataset_dir: "examples/data/offline_dpo_sd3_5"
+ train:
+ weight: 1 # Offline epochs require unit source weights and full traversal.
+ max_dataset_size: 2
+ enable_preprocess: true
+ preprocessing_batch_size: 2
+ dataloader_num_workers: 0
+ force_reprocess: false
+ cache_dir: "~/.cache/flow_factory/datasets"
+ sampler_type: "auto" # Offline training selects torch DistributedSampler.
+
+model:
+ finetune_type: "lora"
+ lora_rank: 16
+ lora_alpha: 16
+ target_modules: "default"
+ model_name_or_path: "stabilityai/stable-diffusion-3.5-medium"
+ model_type: "sd3-5"
+ resume_path: null
+ resume_type: null
+
+log:
+ run_name: "sd3_5_offline_dpo_smoke"
+ project: "Flow-Factory"
+ logging_backend: "none"
+ save_dir: "saves/"
+ save_freq: 0
+ save_model_only: true
+
+train:
+ trainer_type: "offline-dpo"
+ max_epochs: 1 # One epoch is one complete dataloader traversal.
+ resolution: 512
+ guidance_scale: 1.0
+ per_device_batch_size: 1
+ gradient_accumulation_steps: 1
+ beta: 2000.0
+ weighting_scheme: "logit_normal" # Options: "logit_normal", "uniform"
+ num_train_timesteps: 1 # Shared chosen/rejected Monte Carlo time terms.
+ timestep_range: 0.99
+ time_shift: 1.0
+ logit_mean: 0.0
+ logit_std: 1.0
+ ref_param_device: "cpu"
+ ema_decay: 0.0
+ enable_gradient_checkpointing: false
+ seed: 42
+
+scheduler:
+ dynamics_type: "ODE"
+
+eval:
+ eval_freq: 0
+
+optimizers:
+ - name: default
+ learning_rate: 1.0e-5
+ weight_decay: 1.0e-4
+ betas: [0.9, 0.999]
+ eps: 1.0e-8
+ max_grad_norm: 1.0
diff --git a/examples/sft/lora/sd3_5/default.yaml b/examples/sft/lora/sd3_5/default.yaml
new file mode 100644
index 000000000..2d07fa263
--- /dev/null
+++ b/examples/sft/lora/sd3_5/default.yaml
@@ -0,0 +1,70 @@
+# Single-process SFT smoke recipe over a V2 demonstration manifest.
+# Replace the tiny repository-image dataset with a real training corpus for quality runs.
+launcher: "accelerate"
+config_file: null
+num_processes: 1
+main_process_port: 29500
+mixed_precision: "bf16"
+
+data:
+ datasets:
+ - name: offline_demonstrations
+ dataset_dir: "examples/data/sft_sd3_5"
+ train:
+ weight: 1 # Offline epochs require unit source weights and full traversal.
+ max_dataset_size: 2
+ enable_preprocess: true
+ preprocessing_batch_size: 2
+ dataloader_num_workers: 0
+ force_reprocess: false
+ cache_dir: "~/.cache/flow_factory/datasets"
+ sampler_type: "auto" # Offline training selects torch DistributedSampler.
+
+model:
+ finetune_type: "lora"
+ lora_rank: 16
+ lora_alpha: 16
+ target_modules: "default"
+ model_name_or_path: "stabilityai/stable-diffusion-3.5-medium"
+ model_type: "sd3-5"
+ resume_path: null
+ resume_type: null
+
+log:
+ run_name: "sd3_5_sft_smoke"
+ project: "Flow-Factory"
+ logging_backend: "none"
+ save_dir: "saves/"
+ save_freq: 0
+ save_model_only: true
+
+train:
+ trainer_type: "sft"
+ max_epochs: 1 # One epoch is one complete dataloader traversal.
+ resolution: 512
+ guidance_scale: 1.0
+ per_device_batch_size: 1
+ gradient_accumulation_steps: 1
+ weighting_scheme: "logit_normal" # Options: "logit_normal", "uniform"
+ num_train_timesteps: 1 # Monte Carlo terms averaged inside each microbatch.
+ timestep_range: 0.99
+ time_shift: 1.0
+ logit_mean: 0.0
+ logit_std: 1.0
+ ema_decay: 0.0
+ enable_gradient_checkpointing: false
+ seed: 42
+
+scheduler:
+ dynamics_type: "ODE"
+
+eval:
+ eval_freq: 0
+
+optimizers:
+ - name: default
+ learning_rate: 1.0e-5
+ weight_decay: 1.0e-4
+ betas: [0.9, 0.999]
+ eps: 1.0e-8
+ max_grad_norm: 1.0
diff --git a/guidance/algorithms.md b/guidance/algorithms.md
index a6e99380c..67c2bca26 100644
--- a/guidance/algorithms.md
+++ b/guidance/algorithms.md
@@ -15,6 +15,10 @@
- [DPPO](#dppo)
+- [SFT](#sft)
+
+- [Offline DPO](#offline-dpo)
+
- [DPO](#dpo)
- [DGPO](#dgpo)
@@ -37,12 +41,26 @@
## Overview
-Flow-Factory provides unified implementations of state-of-the-art RL algorithms for flow-matching models. All algorithms share the same model adapter and reward interfaces, enabling direct comparison under controlled conditions.
+Flow-Factory provides unified online RL, distillation, and offline objectives for flow-matching
+models. Algorithms share model and dataset contracts while keeping objective-specific policy
+updates independent.
+
+Two execution dimensions are separate from the mathematical paradigm:
+
+| Dimension | Values | Meaning |
+|---|---|---|
+| Acquisition | `generation`, `dataset` | Generate a rollout collection or fetch every batch from a finite dataloader. |
+| Feedback | `runtime_reward`, `none` | Run reward/advantage processing or optimize without runtime feedback. |
+
+GRPO and online DPO use `generation + runtime_reward`. Generation-based distillation uses
+`generation + none`. SFT and offline DPO use `dataset + none`: their sampling stage is the dataset
+acquisition driver, so `adapter.inference()` and `sample()` are not called for training. This keeps
+online/offline selection out of model adapters and loss functions.
At a high level, the supported algorithms fall into three paradigms:
- **Coupled paradigm (GRPO and variants)**: Training timesteps are coupled with the SDE-based sampling dynamics, requiring tractable log-probability computation for policy gradient optimization.
-- **Decoupled paradigm (DPO, DiffusionNFT, AWM, DGPO, CRD, TDM-R1)**: Training timesteps are decoupled from the actual sampling dynamics, making them inherently solver-agnostic.
+- **Decoupled paradigm (SFT, offline DPO, online DPO, DiffusionNFT, AWM, DGPO, CRD, TDM-R1)**: Training timesteps are decoupled from the actual sampling dynamics, making them inherently solver-agnostic.
- **Distillation paradigm (DiffusionOPD, DMD2, TDM)**: Students match flow-matching targets. DiffusionOPD uses a teacher; DMD2 and TDM keep a fake score on one model bundle and update it before the generator.
DMD2, TDM, and TDM-R1 update fake first. TDM-R1 then updates the surrogate
@@ -190,9 +208,76 @@ train:
Like GRPO, DPPO is **coupled** and must use SDE dynamics (`Flow-SDE`, `Dance-SDE`, `CPS`). `DPPOTrainingArguments` does not inherit `GRPOTrainingArguments` (no `clip_range`) — its field set is intentionally minimal. When `kl_beta > 0`, the KL-vs-reference term is evaluated at `kl_guidance_scale`; this is reflected in `DPPOTrainingArguments.get_preprocess_guidance_scale()` so negative prompts are encoded at preprocessing whenever `kl_guidance_scale > 1.0`. Example configs: `examples/dppo/lora/{flux2_klein_base,sd3_5}/geneval2_{single,multi}.yaml`.
+## SFT
+
+SFT trains directly from V2 `demonstration` records. The target media is decoded from its source
+file and encoded to the model's clean output state on every microbatch. At each independently
+sampled flow time, the trainer noises that clean state, predicts velocity through the ordinary
+adapter `forward()` contract, and minimizes the per-sample flow-matching error. Several
+`num_train_timesteps` terms are averaged inside the microbatch before one backward pass.
+
+```yaml
+train:
+ trainer_type: sft
+ max_epochs: 4
+ per_device_batch_size: 1
+ gradient_accumulation_steps: 4
+ weighting_scheme: logit_normal # logit_normal or uniform
+ num_train_timesteps: 1
+ timestep_range: 0.99 # scalar -> (0, scalar), or [lower, upper]
+ time_shift: 1.0
+ logit_mean: 0.0
+ logit_std: 1.0
+```
+
+`max_epochs` counts complete dataloader traversals. `num_train_timesteps` is a Monte Carlo axis,
+not a hidden gradient-accumulation multiplier. SFT has no reference model and rejects training
+`rewards`; `eval_rewards` remain available for generation-based evaluation.
+
+See the [SFT configuration](../examples/sft/lora/sd3_5/default.yaml) and the
+[V2 demonstration schema](datasets.md#demonstration-supervision).
+
+## Offline DPO
+
+Offline DPO applies the Diffusion-DPO objective [[11]](#ref11) to V2 `preference` records rather
+than generating and reward-ranking samples.
+Chosen and rejected media are encoded on the fly under the same input condition. For each loss
+term they share the primary timestep, exact component-time mapping, and diffusion noise. The
+current policy and a frozen reference each produce chosen/rejected flow-matching errors; the
+shared DPO objective applies `beta` to the policy-versus-reference error delta.
+
+```yaml
+train:
+ trainer_type: offline-dpo
+ max_epochs: 1
+ per_device_batch_size: 1
+ gradient_accumulation_steps: 1
+ beta: 2000.0
+ weighting_scheme: logit_normal
+ num_train_timesteps: 1
+ timestep_range: 0.99
+ time_shift: 1.0
+ logit_mean: 0.0
+ logit_std: 1.0
+ ref_param_device: cpu
+```
+
+The reference is mandatory; `reference_free` is not implemented. LoRA runs use the base model by
+disabling trainable adapters inside the reference scope. Full-parameter runs keep the frozen
+reference snapshot on `ref_param_device`. Training rewards and online pair formation are not part
+of this path.
+
+See the [offline-DPO configuration](../examples/offline_dpo/lora/sd3_5/default.yaml), the
+[V2 preference schema](datasets.md#preference-supervision), and the current
+[offline model matrix](datasets.md#offline-model-support).
+
## DPO
-DPO (Direct Preference Optimization) [[11]](#ref11) is a **decoupled** algorithm that optimises a pairwise preference loss on flow-matching velocity targets. Instead of per-sample policy-gradient ratios, it forms chosen/rejected pairs within each group (based on per-sample advantages), then minimises a Bradley-Terry preference loss over the DSM errors of the two policies (current vs. frozen reference). To use this algorithm, set:
+The existing `dpo` trainer is **online DPO** [[11]](#ref11). It is a decoupled algorithm that generates samples,
+scores them with runtime rewards, and forms chosen/rejected pairs within each prompt group from
+the resulting advantages. It then minimizes a Bradley-Terry preference loss over the DSM errors
+of the current and frozen-reference policies. Use `offline-dpo` instead when the dataset already
+contains preference pairs. To select online DPO, set:
```yaml
train:
diff --git a/guidance/datasets.md b/guidance/datasets.md
index 94b0ae850..8edd8aac9 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -53,8 +53,163 @@ resolved against their corresponding media directory. MiniMax H3 Ref2VA is diffe
relative `references[*].path` and `references[*].audio_path` values are resolved against
`dataset_dir`; absolute paths are accepted unchanged.
+## Offline V2 records
+
+SFT and offline DPO require strict JSONL records with `"schema_version": 2`. In this release, V2 is
+a supervised offline format: every record must carry either demonstration or preference
+supervision. It is not an alternative input format for the existing online generation loader.
+V2 separates the model input from its supervision, so dataset rows remain model- and
+algorithm-neutral:
+
+```text
+record
+├── input # prompt plus optional condition media
+├── supervision # demonstration or preference
+└── metadata # optional JSON provenance; never a model input
+```
+
+Every V2 media object uses `type` as its only public discriminator. The accepted values are
+`image`, `video`, and `audio`:
+
+```json
+{"type":"image","path":"images/source.png"}
+{"type":"video","path":"videos/clip.mp4","fps":24.0}
+{"type":"audio","path":"audios/clip.wav","sample_rate":48000}
+```
+
+Do not write `kind` in a V2 record. Some ordered-reference adapters still consume a validated
+legacy `kind` mapping internally; the V2 condition projection creates that private bridge only at
+the adapter preprocessing boundary. It is not part of the public V2 schema.
+
+### Demonstration supervision
+
+One SFT row has a shared input and one target candidate:
+
+```jsonl
+{"schema_version":2,"input":{"prompt":"Restyle the source as a watercolor.","media":[{"type":"image","path":"conditions/source.png"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"targets/watercolor.png"}]}},"metadata":{"license":"example"}}
+```
+
+Use it with `train.trainer_type: sft`. `target.media` is an ordered sequence because a pipeline
+may eventually emit several modalities. Adapter-level codec availability and explicit blockers
+are checked before model weights are loaded. The offline loader then validates each record's exact
+output sequence, rates, input cardinality, and batch capability before condition preprocessing or
+training; adapter-specific encoded geometry is validated at the output-codec boundary.
+
+### Preference supervision
+
+One offline-DPO row shares the input across a chosen and rejected candidate:
+
+```jsonl
+{"schema_version":2,"input":{"prompt":"A clean typographic poster.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"pairs/chosen.png"}]},"rejected":{"media":[{"type":"image","path":"pairs/rejected.png"}]}},"metadata":{"annotator":"example"}}
+```
+
+Use it with `train.trainer_type: offline-dpo`. A training source must be homogeneous: every row
+must carry the supervision type required by its trainer. Prompt-only rows, mixed demonstration and
+preference rows, unknown keys, and non-V2 records fail during manifest loading.
+
+All V2 media paths are resolved against that source's `dataset_dir`; an absolute path is retained.
+Images and videos have built-in CPU decoders. Video targets require PyAV 18 or newer. There is no
+default audio target decoder yet, which is one reason the current audio-video adapters are blocked
+for offline objectives.
+
+Tiny schema-complete fixtures and configs are available for
+[SFT](../examples/sft/lora/sd3_5/default.yaml) and
+[offline DPO](../examples/offline_dpo/lora/sd3_5/default.yaml).
+
+Evaluation still uses generation acquisition, including when the trainer is SFT or offline DPO.
+Consequently, a split enabled through `data.datasets[*].eval` must use one of the legacy
+prompt/condition formats documented under [Common task formats](#common-task-formats), not a
+supervised V2 record. A single dataset directory may therefore contain a V2 `train.jsonl` and a
+legacy prompt-only `test.jsonl` without conflating their roles.
+
+### Condition cache and target lifecycle
+
+Offline preprocessing intentionally caches only the input side:
+
+```text
+V2 input
+ -> project prompt and condition media
+ -> adapter.preprocess_func under no-grad
+ -> Arrow cache of prompt/condition tensors
+
+V2 target, chosen, or rejected
+ -> decode from the source file in Dataset.__getitem__
+ -> collate decoded CPU media
+ -> adapter.encode_output_state under no-grad on every training microbatch
+ -> clean latent state for the objective
+```
+
+Target, chosen, and rejected payloads, their VAE latents, and supervision metadata are never
+written to the Arrow condition cache. There is no target-VAE preprocessing cache. This avoids a
+second large media-derived dataset on disk and keeps output geometry and posterior semantics owned
+by the adapter. The frozen VAE (or other output codec component) remains available at training
+time and performs the comparatively small on-the-fly encode; evaluation already needs the decoder.
+
+During offline dataset construction, every unique normalized input and supervision media path is
+streamed once through SHA-256, with digests memoized only for that source build. These digests are
+identity metadata: media payloads, decoded pixels, and output latents are neither copied nor cached.
+Replacing an input condition file in place therefore changes its condition identity and invalidates
+the Arrow cache automatically. Replacing target, chosen, or rejected media in place changes the
+full record identity used by exact-resume checks; supervision is still decoded afresh and requires
+no target-cache invalidation step.
+
+### Offline dataloader and epoch semantics
+
+Offline sources are concatenated once and sharded with PyTorch's official
+`torch.utils.data.DistributedSampler`, including a one-process run. The trainer calls
+`sampler.set_epoch(data_epoch)` and advances `data_epoch` only after the finite dataloader is
+exhausted successfully. Therefore one offline epoch has the standard meaning: one complete
+dataloader traversal. An exception or interruption during a partial traversal does not publish a
+completed epoch.
+
+Keep these configuration rules:
+
+```yaml
+data:
+ datasets:
+ - name: demonstrations
+ dataset_dir: "dataset/demonstrations"
+ train: {weight: 1}
+ enable_preprocess: true
+ sampler_type: auto
+train:
+ max_epochs: 4
+ per_device_batch_size: 1
+ gradient_accumulation_steps: 4
+```
+
+- Each offline source must use `train.weight: 1`; replacement weighting would make full traversal
+ stop meaning one data epoch.
+- `gradient_accumulation_steps` must be an explicit positive integer. The number of rank-local
+ batches must divide evenly by it; the framework never adds batches merely to close a partial
+ accumulation window or silently flushes one.
+- PyTorch's official `DistributedSampler` owns cross-rank tail handling. With its default
+ `drop_last=False`, it may repeat tail indices when the global dataset size is not divisible by
+ world size so every rank traverses the same number of samples. This is standard sampler behavior;
+ an offline epoch is the resulting complete dataloader traversal, not a global uniqueness claim.
+- The offline loader is already rank-sharded and is not passed to `Accelerator.prepare()`.
+- `num_train_timesteps` controls independently sampled Monte Carlo loss terms averaged inside a
+ microbatch. It does not multiply gradient accumulation or change epoch length.
+
+### Offline model support
+
+The V2 schema is broader than the codecs currently implemented by adapters. Static capability
+validation fails before heavyweight model loading when a selected pipeline cannot preserve its
+output semantics.
+
+| Offline status | Model types | Notes |
+|---|---|---|
+| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. |
+| Supported | `wan2_t2v` | Video targets require `fps`; the codec resamples to configured frames/rate and samples the Wan VAE posterior on the fly. |
+| Blocked | `wan2_i2v` | Output geometry depends on the first-frame VAE latent/mask, while the current condition cache does not preserve the source pixels needed by that binder. |
+| Blocked | `ltx2_t2av`, `ltx2_i2av` | Lossless audio decode/rate metadata and exact audio-video duration alignment are not unified; I2AV also needs the pinned first-frame active mask. |
+| Blocked | `minimax-h3-t2va`, `minimax-h3-fl2va`, `minimax-h3-ref2va` | The audio-video boundary and official target-video posterior policy are not yet defined for offline targets. |
+
## Common task formats
+The following compact formats remain supported for generation acquisition. They are separate from
+strict V2 offline records and do not carry output supervision.
+
### Text-conditioned generation
Text-to-image, text-to-video, and text-to-audio-video datasets may use plain text:
@@ -201,7 +356,8 @@ directory. See the [FL2VA dataset fixture](../dataset/minimax_h3_fl2va/train.jso
### Ref2VA: `minimax-h3-ref2va`
-Ref2VA uses a non-empty ordered `"references"` array containing image, video, and audio entries:
+The existing online Ref2VA loader uses a legacy non-empty ordered `"references"` array containing
+image, video, and audio entries:
```jsonl
{"prompt":"Create a coherent scene using the references in order.","references":[{"kind":"image","path":"references/style.png"},{"kind":"video","path":"references/motion.mp4","fps":12.0},{"kind":"audio","path":"references/ambience.wav","sample_rate":16000}]}
@@ -223,6 +379,10 @@ Supported entries:
soundtrack or a separate dataset-relative `audio_path`; a video `sample_rate` override requires
`audio_path`. Unknown keys and unsupported `kind` values fail before preprocessing.
+This legacy online manifest is distinct from the strict V2 format above. A V2 record always uses
+`input.media[*].type`; offline condition projection performs any required legacy `kind` conversion
+internally.
+
See the [Ref2VA dataset fixture](../dataset/minimax_h3_ref2va/train.jsonl), its
[local fixture notes](../dataset/minimax_h3_ref2va/README.md), and the
[Ref2VA GRPO configuration](../examples/grpo/lora/minimax_h3_ref2va/default.yaml).
@@ -240,9 +400,10 @@ TXT/JSONL row
-> collate cached fields for rollout
```
-Prompt encoders, VAEs, and processors are preprocessing components. The trainer can offload them
-after cache creation because optimization consumes their cached outputs rather than decoding and
-encoding every source row again.
+Prompt encoders, condition VAEs, and processors are preprocessing components. Online RL can
+offload them after cache creation because optimization consumes cached conditions. Offline SFT and
+offline DPO reload any output-codec components declared by the adapter and encode target,
+chosen, and rejected media on the fly; those output states are never cached.
Ref2VA adds an ordered-reference path:
diff --git a/guidance/new_model.md b/guidance/new_model.md
index 3069bcafa..025686559 100644
--- a/guidance/new_model.md
+++ b/guidance/new_model.md
@@ -12,6 +12,7 @@
- [Step 5: Implement `inference()`](#step-5-implement-inference)
- [Step 6: Implement `forward()`](#step-6-implement-forward)
- [Step 7: Register the Adapter](#step-7-register-the-adapter)
+- [Advanced: Offline Output-State Encoding](#advanced-offline-output-state-encoding)
- [Advanced: Custom `preprocess_func`](#advanced-custom-preprocess_func)
- [Advanced: Pseudo-Pipeline for Non-Diffusers Models](#advanced-pseudo-pipeline-for-non-diffusers-models)
- [Data Format Conventions](#data-format-conventions)
@@ -19,7 +20,7 @@
## Overview
-Flow-Factory uses a **model adapter** pattern that wraps [diffusers](https://github.com/huggingface/diffusers) pipelines into a unified interface for RL training. Each adapter maps a diffusers pipeline to a consistent API that the training loop can call without knowing model-specific details.
+Flow-Factory uses a **model adapter** pattern that wraps [diffusers](https://github.com/huggingface/diffusers) pipelines into a unified interface for online and offline fine-tuning. Each adapter maps a diffusers pipeline to a consistent API that the training loop can call without knowing model-specific details.
The relationship is straightforward:
@@ -114,7 +115,7 @@ class MyModelSample(T2ISample):
| `I2AVSample` | Image-to-audio-video | `ImageConditionSample` subclass |
| `V2VSample` | Video-to-video | `VideoConditionSample` subclass |
-> See [`src/flow_factory/samples/samples.py`](src/flow_factory/samples/samples.py) for all available classes.
+> See [`src/flow_factory/samples/samples.py`](../src/flow_factory/samples/samples.py) for all available classes.
> **Key**: The `_shared_fields` class variable declares fields that are identical across a batch (e.g., `height`, `width`, `latent_index_map`). During `BaseSample.stack()`, shared fields take the first element instead of stacking.
@@ -192,7 +193,7 @@ class MyModelAdapter(BaseAdapter):
| `preprocessing_modules` | `['text_encoders', 'vae']` |
| `inference_modules` | `['transformer', 'vae']` |
-Override only when your model deviates — for example, [WAN-T2V](src/flow_factory/models/wan/wan2_t2v.py) models need `['text_encoders', 'vae', 'image_encoder']` for preprocessing and conditionally include `transformer_2` for inference.
+Override only when your model deviates — for example, [WAN-T2V](../src/flow_factory/models/wan/wan2_t2v.py) models need `['text_encoders', 'vae', 'image_encoder']` for preprocessing and conditionally include `transformer_2` for inference.
> **Tip**: Use `print(dict(self.pipeline.named_children()))` to discover available component names.
@@ -547,6 +548,68 @@ model:
model_name_or_path: "org/my-model-checkpoint"
```
+## Advanced: Offline Output-State Encoding
+
+Online-only adapters can keep the default `build_output_state_codec() -> None`. To support SFT or
+offline DPO, an adapter must additionally declare both sides of its pipeline and provide an
+on-the-fly output codec:
+
+1. Set a class-level `pipeline_io_contract`. It owns input media counts/order/binding, negative
+ prompt policy, the exact ordered output media sequence, rate requirements, geometry source, and
+ batch capability.
+2. Override `build_output_state_codec()` with a declaration-only codec. Its
+ `required_components` names logical runtime components such as `("vae",)`; construction must
+ not load, materialize, move, replace, or cast them.
+3. Return an `EncodedOutputState` containing a detached `LatentState`, output-derived forward and
+ decode contexts, and one exact geometry signature per sample.
+4. Override `_validate_encoded_output_geometry()` so configured, condition-derived, and
+ output-derived dimensions cannot drift silently.
+5. Declare a complete immutable `offline_training_forward_overrides` mapping whenever the base
+ `{"guidance_scale": 1.0}` contract does not describe the adapter. Offline trainers apply this
+ mapping after sampling configuration and cached batch conditions, so it owns loss-time model
+ conditioning. Conventional CFG branches must all be set to their neutral point (for example,
+ both Wan transformer scales or Bagel text/image CFG scales); guidance-distilled models instead
+ declare the explicit guidance-embedding value used by their official training recipe. Replace
+ the complete mapping so permissive `**kwargs` forwards do not receive unrelated base keys.
+
+The dataset remains responsible only for strict V2 parsing and CPU media decoding. The adapter
+owns numerical output semantics. The SFT/offline-DPO trainer calls `encode_output_state()` on every
+microbatch under `torch.no_grad`; target, chosen, and rejected latents are not preprocessing-cache
+columns. The declared output components are loaded through `ModelLoadCoordinator`, never from
+inside the codec.
+
+Condition encoding and target encoding should share role-neutral numerical transforms instead of
+duplicating VAE math. Extract helpers for pixel preprocessing, posterior extraction, latent
+normalization, patchification, IDs, and packing, then make the posterior policy an explicit
+argument:
+
+```python
+def encode_vae_image(adapter, pixels, *, sample_mode, generator=None):
+ posterior = adapter.vae.encode(pixels).latent_dist
+ latent = (
+ posterior.sample(generator=generator)
+ if sample_mode == "sample"
+ else posterior.mode()
+ )
+ return normalize_and_pack(adapter, latent)
+```
+
+The helper is role-neutral; the caller is not. Follow the official Diffusers pipeline for each
+role. Condition paths commonly use posterior `argmax`/`mode` for stable conditioning, while
+training targets use posterior `sample` and forward the caller's generator. Never merge the two
+entry points in a way that silently changes this policy. Tests should compare the shared transform
+against the pinned Diffusers helper and assert both sample/argmax behavior and generator routing.
+
+An output codec is not merely an `encode_image()` alias. It may need output-specific geometry,
+multi-component state order, active masks, rate alignment, or forward context that condition
+encoding does not own. If those semantics are not lossless, set a concrete
+`output_state_codec_unavailable_reason` so offline selection fails before downloading weights.
+
+SenseNova is an example of an important boundary: its existing condition schema uses grouped
+`images` with within-type order. Do not advertise heterogeneous ordered references merely because
+several images are accepted. The public V2 discriminator remains `type`; conversion to a legacy
+adapter-internal `kind` entry, when genuinely required, belongs only in the condition projection.
+
## Advanced: Custom `preprocess_func`
The default `preprocess_func` calls `encode_prompt`, `encode_image`, `encode_video` and `encode_audio` independently. Override it when your model requires **cross-modal preprocessing** — for example, FLUX.2 uses its text encoder to "upsample" (rewrite) prompts based on input images before encoding ([here](https://github.com/X-GenGroup/Flow-Factory/blob/main/src/flow_factory/models/flux/flux2.py#L371)):
@@ -837,9 +900,14 @@ Before submitting a new model adapter, verify:
- [ ] **`encode_audio()`** — Override only if your model consumes audio; handles `MultiAudioBatch` input format (text/image/video-only models inherit the no-op default)
- [ ] **`inference()`** — Accepts both raw and pre-encoded inputs; returns `List[Sample]`
- [ ] **`forward()`** — Single denoising step; ends with `self.scheduler.step()`; returns `SDESchedulerOutput`
+- [ ] **Pipeline I/O contract** — Declares exact input/output media, rate, geometry, and batch semantics before enabling offline training
+- [ ] **Output-state codec (when supported)** — Declaration-only logical component requirements; on-the-fly detached target encoding; exact geometry validation
+- [ ] **Offline forward overrides (when supported)** — Complete immutable adapter mapping; sampling controls never define offline loss semantics; every CFG branch is neutralized or every distilled guidance condition is explicitly pinned
+- [ ] **Role-neutral encoder math** — Condition/output paths reuse transforms but explicitly preserve official posterior `sample` versus `argmax` policy and generator routing
+- [ ] **Explicit offline blocker (when unsupported)** — `output_state_codec_unavailable_reason` names the missing lossless semantic boundary
- [ ] **Sample dataclass** — All fields without batch dimension; `_shared_fields` correctly set; custom field types are consistent (no `Tensor` vs `List[Tensor]` mixing across samples)
- [ ] **Registry entry** — Added to `_MODEL_ADAPTER_REGISTRY`
-- [ ] **Tested** — Runs at least one epoch of GRPO training without errors
+- [ ] **Tested** — Runs at least one rollout cycle for online support and one complete dataloader epoch for any declared offline support
## Component Runtime and Structured Replay
@@ -875,4 +943,4 @@ rather than replacing registered modules.
MiniMax H3 is the reference for workflow-pruned modular components and separate
video/audio trajectories. See `src/flow_factory/models/minimax_h3/` and the
-[MiniMax H3 dataset contracts](datasets.md#minimax-h3-datasets).
\ No newline at end of file
+[MiniMax H3 dataset contracts](datasets.md#minimax-h3-datasets).
diff --git a/guidance/workflow.md b/guidance/workflow.md
index 1dd42de13..7af9c133b 100644
--- a/guidance/workflow.md
+++ b/guidance/workflow.md
@@ -4,6 +4,7 @@
- [Overview](#overview)
- [Stage 1: Data Preprocessing](#stage-1-data-preprocessing)
+- [Offline Dataset Training](#offline-dataset-training)
- [Stage 2: K-Repeat Sampling](#stage-2-k-repeat-sampling)
- [Stage 3: Trajectory Generation](#stage-3-trajectory-generation)
- [Stage 4: Reward Computation](#stage-4-reward-computation)
@@ -13,11 +14,35 @@
## Overview
-Flow-Factory follows an **online RL** training paradigm for diffusion/flow-matching models. Each epoch executes a six-stage pipeline:
+Flow-Factory has one training kernel with orthogonal execution contracts. An algorithm declares
+where examples come from and whether they need runtime feedback:
+
+| Contract axis | Value | Runtime behavior |
+|---|---|---|
+| Acquisition | `generation` | Run adapter inference to create a rollout collection. |
+| Acquisition | `dataset` | Fetch and optimize every batch from a finite dataloader. |
+| Feedback | `runtime_reward` | Compute rewards and advantages before optimization. |
+| Feedback | `none` | Pass acquired examples directly to the objective. |
+
+This produces three currently useful compositions:
+
+```text
+online RL generation + runtime_reward
+generation distillation generation + none
+SFT / offline DPO dataset + none
+```
+
+Acquisition is not a model concern, and feedback is not inferred from a batch shape. The trainer,
+its algorithm-specific arguments, and the shared driver must declare the same immutable contract.
+Pipeline I/O is a separate adapter contract that validates input modalities, ordered references,
+output modality/rates, geometry ownership, and batch capability. The data schema therefore does
+not need model-specific tensor names, while the objective does not need file-format branches.
+
+The familiar online RL path still executes six stages:
```
┌─────────────────────────────────────────────────────────────────────────────────┐
-│ Flow-Factory Training Epoch │
+│ Flow-Factory Online Rollout Cycle │
│ │
│ ┌─────────────┐ ┌──────────┐ ┌──────────────┐ ┌───────────────┐ │
│ │ Data │ │ K-Repeat │ │ Trajectory │ │ Reward │ │
@@ -32,21 +57,24 @@ Flow-Factory follows an **online RL** training paradigm for diffusion/flow-match
└─────────────────────────────────────────────────────────────────────────────────┘
```
-The high-level training loop lives once in `BaseTrainer.start()`; no algorithm restates it:
+The high-level loop lives once in `BaseTrainer.start()`. Its acquisition driver selects either a
+generated collection or one finite dataset traversal:
```python
# src/flow_factory/trainers/abc.py — BaseTrainer.start()
def start(self):
while self.should_continue_training():
- # Reseed, checkpoint on save_freq, evaluate on eval_freq (omitted for brevity)
- self._run_training_step() # Stages 2-6, the algorithm-specific middle
- self.adapter.ema_step(step=self.epoch)
- self._after_optimizer_step()
- self.epoch += 1
+ driver.prepare_cycle(self, progress, seed=train.seed)
+ driver.run_cycle(self, progress)
+ if contract.acquisition is AcquisitionMode.GENERATION:
+ self.adapter.ema_step(step=progress.rollout_iteration)
+ # Runs once per completed rollout iteration or fully exhausted data epoch.
+ self._after_acquisition_cycle()
+ # Advance only after the selected acquisition completes successfully.
+ progress = progress.advance_acquisition(contract.acquisition, completed=True)
```
-The default `_run_training_step` is the familiar three-stage sequence, so a reward-based
-algorithm only implements `optimize()`:
+Generation acquisition retains the familiar sequence:
```python
samples = self.sample() # Stages 2 + 3
@@ -54,19 +82,24 @@ self.prepare_feedback(samples) # Stages 4 + 5 (rewards + advantages)
self.optimize(samples) # Stage 6 (DPO: pair formation + loss here)
```
-> **Note**: Stage 1 (preprocessing) runs *once* before training begins and is cached to disk. Stages 2–6 repeat every epoch. The three methods above map directly to those stages: `sample` → trajectory rollouts; `prepare_feedback` → finalize rewards from the buffer and compute advantages; `optimize` → policy update (DPO additionally forms chosen/rejected pairs at the start of `optimize` before the loss). `optimize()` is the only abstract one; vary the rest through `sampling_context`, `_run_training_step` and `_after_optimizer_step`.
+Dataset acquisition never calls `sample()` or `adapter.inference()` for training. It calls
+`optimize_batch(batch)` for every loader batch and advances the data epoch only after exhaustion.
+This is the concrete form of treating the RL sampling stage as a no-op/fetch operation for an
+offline algorithm, without teaching online trainers about offline batch formats.
## Stage 1: Data Preprocessing
-**Goal**: Encode raw text prompts (and optional images / videos / audio files) into model-ready tensor representations *before* training begins, eliminating redundant computation during the RL loop and enabling components offloading such as **text-encoder**, **image-encoder**, and **audio-encoder** (when applicable).
+**Goal**: Encode raw text prompts (and optional condition images / videos / audio files) into
+model-ready tensor representations before repeated training work, eliminating redundant condition
+encoding and enabling component offloading.
### Input / Output
| | Description |
|---|---|
-| **Input** | Raw dataset: `train.jsonl` or `train.txt` containing prompts, optional image / video / audio paths |
-| **Output** | Cached HuggingFace Dataset on disk with pre-encoded tensors (`prompt_embeds`, `prompt_ids`, `pooled_prompt_embeds`, `image_latents`, etc.) |
+| **Input** | Raw dataset containing prompts and optional condition media paths. Offline objectives use the `input` portion of strict V2 records. |
+| **Output** | Cached HuggingFace Dataset with input tensors (`prompt_embeds`, `prompt_ids`, condition latents, etc.). Output supervision is excluded. |
### How It Works
@@ -105,13 +138,21 @@ def preprocess_func(self, prompt, images, ...):
- **Cache layout**: The merged cache directory looks like `{cache_dir}/{fingerprint}/_parts/rank_{i:05d}_of_{N:05d}/cache-{fingerprint}_shard{i}of{N-1}.arrow`, plus the top-level `state.json` and `dataset_info.json`. While preprocessing is in flight, the same content lives under `{cache_dir}/{fingerprint}.tmp/`, with a `_build_meta.json` sentinel that records `num_shards` so a subsequent run with the same `num_shards` can resume from any per-rank Arrow files that were already written before a crash, while a different `num_shards` triggers a clean wipe.
- **No HF default-cache copy**: Because each `map()` call sets `cache_file_name`, HuggingFace does **not** also write a duplicate `cache-*.arrow` under `~/.cache/huggingface/datasets/...`.
- **Intelligent caching**: A hash fingerprint of `(dataset, split, max_dataset_size, preprocess_func source, preprocess_kwargs, extra_hash_strs)` (the last includes `model_type` and `model_name_or_path`) determines the cache path. Subsequent runs that match the fingerprint take the fast path without any `Dataset.map` invocation.
-- **Component offloading**: Text encoders and VAEs are loaded for preprocessing, then offloaded before the training loop to free VRAM for the denoising model.
+- **Component offloading**: Text and condition encoders can be offloaded after cache creation. An
+ offline adapter's declared output codec components are reloaded through the component lifecycle
+ for on-the-fly target encoding.
+- **No target cache**: SFT targets and offline-DPO chosen/rejected candidates are decoded per
+ dataset access and encoded per training microbatch. Target payloads, latent states, and metadata
+ never enter the Arrow condition cache.
### Configuration
```yaml
data:
- dataset: "path/to/dataset"
+ datasets:
+ - name: example
+ dataset_dir: "path/to/dataset"
+ train: {weight: 1}
enable_preprocess: true # Enable offline preprocessing
force_reprocess: false # Force re-encoding even if cache exists; essential if code is modified without changing config
preprocessing_batch_size: 16 # Batch size for encoding
@@ -119,9 +160,64 @@ data:
preprocess_parallelism: "local" # "local" = per-node parallelism (no shared FS required); "global" = cross-node (shared FS required)
```
+## Offline Dataset Training
+
+SFT and offline DPO use strict V2 JSONL manifests described in the
+[dataset guide](datasets.md#offline-v2-records). Their finite loader is constructed with PyTorch's
+official `DistributedSampler`, even for `num_replicas=1`, and is not prepared by Accelerator:
+
+```text
+sampler.set_epoch(data_epoch)
+for batch in dataloader:
+ condition = cached prompt/input tensors
+ output = freshly decoded target or chosen/rejected media
+ trainer.optimize_batch(batch)
+data_epoch += 1 # only after clean exhaustion
+```
+
+One complete dataloader traversal is one offline epoch. Source weights must be `1`,
+`data.sampler_type` remains `auto`, and `gradient_accumulation_steps` is an explicit integer. The
+rank-local batch count must be divisible by gradient accumulation; the framework does not add
+batches to close a partial accumulation window or flush one at epoch end. The official
+`DistributedSampler` retains its standard `drop_last=False` behavior and may repeat global tail
+indices to equalize rank lengths. Those indices are part of the finite loader traversal.
+
+Progress uses three independent counters:
+
+| Counter | Advances when | Used by |
+|---|---|---|
+| `rollout_iteration` | One generation acquisition completes. | Online rollout cadence and its compatibility `epoch`. |
+| `data_epoch` | One finite offline dataloader traversal completes. | Offline epoch, sampler shuffle epoch, save/eval boundaries. |
+| `optimizer_step` | One optimizer update completes. | Training metrics and optimizer-step state. |
+
+Offline EMA updates on optimizer-step cadence because one data epoch may contain many optimizer
+updates. Online algorithms retain their rollout-cycle EMA cadence. Similarly,
+`num_train_timesteps` means independently sampled flow-matching terms averaged inside one offline
+microbatch; it neither advances `optimizer_step` nor multiplies gradient accumulation.
+
+An exact-state checkpoint (`log.save_model_only: false`) locks the realized training loader plus
+evaluation cadence, sampling arguments, ordered eval loaders, per-dataset overrides, and eval
+reward configuration. Online resume replays evaluation after its pre-rollout checkpoint, and
+evaluation adapters or rewards may consume global RNG before the next training acquisition. MPS
+cannot currently save exact state because Accelerate does not serialize MPS RNG; use
+`log.save_model_only: true` on Apple Silicon.
+
+The target codec is adapter-owned but role-neutral at its numerical core. Condition and output
+encoding reuse the same pixel preprocessing, VAE transform, normalization, and packing helpers.
+Their semantic policies remain explicit: official condition paths commonly use posterior
+`argmax`, while stochastic target training uses posterior `sample` with an optional generator.
+Sharing a transform must never silently erase that role boundary.
+
+For offline DPO, chosen and rejected arms share the primary timestep, component-time mapping, and
+diffusion noise. Both policy arms run before one frozen-reference scope covers both reference
+forwards. SFT has no reference branch.
+
## Stage 2: K-Repeat Sampling
+Stages 2–6 in this guide describe generation acquisition. Dataset acquisition replaces Stages 2–5
+with the finite-loader path above and enters its objective through `optimize_batch()`.
+
**Goal**: Construct batches where each unique prompt appears exactly $K$ times (`group_size`), enabling group-relative advantage computation.
### Input / Output
@@ -350,13 +446,13 @@ train:
## Stage 6: Policy Optimization
-**Goal**: Update the denoising model's parameters using the computed advantages and PPO-style clipped policy gradient.
+**Goal**: Update the denoising model through the selected online or offline objective.
### Input / Output
| | Description |
|---|---|
-| **Input** | `List[BaseSample]` with advantages, trajectories, and log-probs stored |
+| **Input** | Generated `List[BaseSample]` for generation acquisition, or one typed offline batch for dataset acquisition. |
| **Output** | Updated model parameters; logged loss metrics |
### How It Works (GRPO)
@@ -412,7 +508,9 @@ def optimize(self, samples):
| **AWM** | Samples fresh timesteps; weights velocity matching loss by advantage; PPO clipping + EMA-KL regularization |
| **DGPO** | Samples fresh timesteps via `TimeSampler`; applies group-level preference objective with optional PPO clipping and EMA-reference KL |
| **CRD** | Samples fresh timesteps; reward distillation against CFG-guided teacher with adaptive KL; old/sampling model snapshots and centered advantages |
-| **DPO** | Preference loss on chosen/rejected pairs; pairs formed at the start of `optimize` after advantages |
+| **DPO** | Online preference loss on reward-ranked pairs formed at the start of `optimize` after advantages |
+| **SFT** | On-the-fly target encoding followed by independently noised flow-matching loss in `optimize_batch` |
+| **Offline DPO** | On-the-fly chosen/rejected encoding, shared timestep/noise, and policy-vs-reference DPO loss in `optimize_batch` |
### Optimizer Configuration
@@ -559,7 +657,7 @@ unverified; use DDP or FSDP.
### Key Points
-- **Inner epochs**: Samples can be reused for multiple optimization passes (`num_inner_epochs`), amortizing the cost of sampling.
+- **Generation inner epochs**: Generated samples can be reused for multiple optimization passes (`num_inner_epochs`), amortizing rollout cost. Offline `max_epochs` instead counts full loader traversals.
- **Gradient accumulation**: The `accelerator.accumulate()` context handles gradient accumulation across timesteps and micro-batches, with optimizer steps only at sync boundaries.
- **KL regularization**: Optional penalty keeping the policy close to a reference model (or EMA model for AWM), preventing reward hacking.
- **Per-timestep iteration**: GRPO iterates over each stored trajectory timestep, computing loss at each. NFT, AWM, DGPO, and CRD sample fresh timesteps independently of the sampling trajectory.
@@ -593,7 +691,25 @@ Epoch N
└── Optimizer step at sync boundaries
```
-*DPO*: form chosen/rejected pairs at the **start** of `optimize()` (after advantages exist), then run the preference loss; there is no pair formation in `prepare_feedback()`.
+*Online DPO*: form chosen/rejected pairs at the **start** of `optimize()` (after advantages exist), then run the preference loss; there is no pair formation in `prepare_feedback()`.
+
+A complete offline epoch is a different acquisition shape:
+
+```text
+Data epoch N
+├── DistributedSampler.set_epoch(N)
+├── Exhaust every rank-local dataloader batch
+│ ├── Reuse cached prompt/input condition
+│ ├── Decode target or chosen/rejected media from source
+│ ├── Encode output state on the fly
+│ ├── Compute SFT or offline-DPO loss
+│ └── Optimizer step at explicit GAS sync boundaries
+└── Advance data_epoch only after clean exhaustion
+```
+
+There is no rollout, training reward, advantage computation, or online pair formation in this
+path. Save/evaluation boundaries use the completed data epoch; training metrics use the independent
+optimizer-step counter.
## Structured multimodal trajectories
@@ -609,4 +725,4 @@ The component order is adapter-owned, for example `("video", "audio")` in MiniMa
H3. Conditioning is packed and replayed by the adapter; trainers such as GRPO,
GRPO-Guard, DPPO, DiffusionNFT, AWM, DPO, DGPO, CRD, and DiffusionOPD consume the
same state interface. H3 accepts neutral guidance `1.0`; framework-interface compatibility
-does not itself establish real-weight numerical parity.
\ No newline at end of file
+does not itself establish real-weight numerical parity.
From 2ba3d7e4cc3ededf5fb12d2027c5ea063291ab0a Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Fri, 28 Aug 2026 19:32:06 +0800
Subject: [PATCH 23/76] fix(examples): align MiniMax H3 T2VA recipe contract
---
.agents/knowledge/topics/fix_patterns.md | 8 ++++++++
examples/README.md | 13 +++++++------
guidance/datasets.md | 14 +++++++++-----
tests/examples/test_minimax_h3_examples.py | 19 +++++++++++++++++++
4 files changed, 43 insertions(+), 11 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 110ffcfdb..9b6f6737b 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -190,6 +190,14 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Lesson**: When delegating sharding to an official sampler, define epoch semantics over its realized finite loader. Distinguish sampler-level repeated indices from optimizer-level synthetic padding.
- **Related Constraint**: #9
+### Recipe migrations must move tests and documentation with the config
+- **Date**: 2026-08-28
+- **Symptom**: Rebasing onto the precision-aware loading branch changed the MiniMax H3 T2VA default to the shared `vid_prompt` source, added ImageBind routing, and removed its old unvalidated warning, while the executable example test and user guides still required the prior dataset and wording.
+- **Root Cause**: The recipe-only commit updated YAML semantics without treating example assertions, dataset links, dependency notes, and validation-status language as one public workflow contract.
+- **Fix**: The T2VA default test now locks the shared TXT manifests and CLAP/ImageBind routing while retaining the dedicated JSONL fixture check for the validated debug recipe. The example and dataset guides now describe the shared source, ImageBind dependency, and the exact evidence boundary without claiming a completed long run.
+- **Lesson**: An example configuration is executable documentation. Any recipe migration must update its production parse test, linked data provenance, optional dependency instructions, and validation claims in the same integration change.
+- **Related Constraint**: #15
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/examples/README.md b/examples/README.md
index 29e5d8e80..2dea62db8 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -89,9 +89,10 @@ The T2VA `debug.yaml` recipe is real-weight validated with the 61 GB checkpoint
1 GPU and 16 GPUs across two nodes completed CPS rollout, video/audio decode,
CLAP reward, GRPO replay/backward/optimizer step, and LoRA checkpoint save/resume.
Its 64x96 canvas is intentionally a correctness geometry. The quality-oriented T2VA
-default remains an unverified quality starting point. FL2VA and Ref2VA are
-**Schema/API validated only** rather than claims of training stability or reward
-improvement.
+default is now the shared-`vid_prompt`, LoRA-rank-64 baseline aligned with the LTX2
+T2AV recipe and uses both CLAP and ImageBind rewards. It is configuration/API
+validated; no completed long-run reward trend is claimed. FL2VA and Ref2VA are also
+**Schema/API validated only**, rather than claims of training stability or reward improvement.
The T2VA `quality_720p_fsdp2.yaml` recipe is the active native-quality path:
768x1344, 124 frames, 24 denoising steps, LoRA rank 64 / alpha 128, and two
@@ -99,9 +100,9 @@ updates from 48 prompt groups per epoch. Its real-weight FSDP2 initialization,
checkpoint, native-resolution decode, and CLAP evaluation are validated; a
completed long-run reward trend is not yet claimed.
-FL2VA and Ref2VA use Meta ImageBind for audio-video alignment. Install ImageBind
-and PyTorchVideo from their upstream repositories before running those examples;
-ImageBind is licensed CC-BY-NC-SA 4.0 (NonCommercial).
+The aligned T2VA default, FL2VA, and Ref2VA use Meta ImageBind for audio-video
+alignment. Install ImageBind and PyTorchVideo from their upstream repositories before
+running those examples; ImageBind is licensed CC-BY-NC-SA 4.0 (NonCommercial).
```bash
pip install git+https://github.com/facebookresearch/ImageBind.git
diff --git a/guidance/datasets.md b/guidance/datasets.md
index 8edd8aac9..c402f34d7 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -321,8 +321,10 @@ T2VA `debug.yaml` is real-weight validated on 1 and 16 GPUs, including LoRA
checkpoint save/resume. Its 64x96 canvas validates
correctness and memory fit, not visual quality or reward improvement.
`quality_720p_fsdp2.yaml` has real-weight initialization, checkpoint, native-resolution
-decode, and evaluation coverage; no long-run reward trend is claimed. The default,
-FL2VA, and Ref2VA configs remain schema/API-validated starting points.
+decode, and evaluation coverage; no long-run reward trend is claimed. The aligned default uses
+the shared `dataset/vid_prompt` source, LoRA rank 64, and CLAP plus ImageBind rewards; it remains a
+configuration/API-validated baseline without a published long-run trend. FL2VA and Ref2VA remain
+schema/API-validated starting points.
### T2VA: `minimax-h3-t2va`
@@ -332,9 +334,11 @@ T2VA is prompt-only:
{"prompt":"A small paper windmill turns beside a quiet stream with synchronized birdsong."}
```
-Do not include negative prompts, images, or references. Use the
-[T2VA dataset fixture](../dataset/minimax_h3_t2va/train.jsonl) with the
-[T2VA GRPO configuration](../examples/grpo/lora/minimax_h3_t2va/default.yaml).
+Do not include negative prompts, images, or references. The
+[T2VA default GRPO configuration](../examples/grpo/lora/minimax_h3_t2va/default.yaml) uses the
+shared [`vid_prompt` TXT dataset](../dataset/vid_prompt/train.txt). The dedicated
+[T2VA JSONL fixture](../dataset/minimax_h3_t2va/train.jsonl) remains the compact input for the
+real-weight validated `debug.yaml` recipe.
### FL2VA: `minimax-h3-fl2va`
diff --git a/tests/examples/test_minimax_h3_examples.py b/tests/examples/test_minimax_h3_examples.py
index 345786ec1..e56f19d56 100644
--- a/tests/examples/test_minimax_h3_examples.py
+++ b/tests/examples/test_minimax_h3_examples.py
@@ -89,6 +89,15 @@ def test_examples_parse_through_production_config_and_registry(
assert "N + 1 states and exactly N trainable transitions" in yaml_text
assert "B=1" in yaml_text
assert "no CFG" in yaml_text
+ if directory == "minimax_h3_t2va":
+ assert "not been run with the 61 GB checkpoint" not in yaml_text
+ assert [reward.reward_model for reward in config.reward_args] == [
+ "clap",
+ "imagebind",
+ ]
+ assert all(reward.applicable_datasets == ["vid_prompt"] for reward in config.reward_args)
+ else:
+ assert "not been run with the 61 GB checkpoint" in yaml_text
assert "stg_scale" not in yaml_text
assert "modality_scale" not in yaml_text
assert "negative_prompt" not in yaml_text
@@ -120,12 +129,22 @@ def test_t2va_validated_variants_parse(
def test_t2va_manifests_contain_prompt_only() -> None:
+ """Keep the dedicated JSONL fixture used by the validated debug recipe strict."""
for split in ("train", "test"):
rows = _read_jsonl(ROOT / "dataset/minimax_h3_t2va" / f"{split}.jsonl")
assert rows
assert all(set(row) == {"prompt"} and row["prompt"] for row in rows)
+def test_t2va_default_shared_text_manifests_contain_prompts() -> None:
+ """The aligned default recipe uses the shared prompt-only TXT dataset."""
+ for split in ("train", "test"):
+ prompts = (ROOT / "dataset/vid_prompt" / f"{split}.txt").read_text(encoding="utf-8")
+ lines = prompts.splitlines()
+ assert lines
+ assert all(prompt.strip() for prompt in lines)
+
+
def test_fl2va_manifests_preserve_one_or_two_ordered_images() -> None:
dataset_dir = ROOT / "dataset/minimax_h3_fl2va"
for split in ("train", "test"):
From 142760e4e51504d2da0c435dca035571c3b08ba1 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Fri, 28 Aug 2026 22:18:06 +0800
Subject: [PATCH 24/76] feat(minimax-h3): add offline AV target encoding
---
.../knowledge/topics/adapter_conventions.md | 2 +
.agents/knowledge/topics/minimax_h3.md | 68 ++
README.md | 6 +-
guidance/datasets.md | 23 +-
.../data_utils/offline_dataset.py | 50 +
.../data_utils/offline_train_data.py | 8 +-
src/flow_factory/models/minimax_h3/_output.py | 891 ++++++++++++++++++
.../models/minimax_h3/adapters.py | 67 +-
.../models/minimax_h3/workflow.py | 68 +-
src/flow_factory/models/pipeline_contracts.py | 58 +-
tests/data_utils/test_offline_dataset.py | 48 +-
tests/data_utils/test_offline_train_data.py | 43 +
tests/models/minimax_h3/test_output_codec.py | 604 ++++++++++++
tests/models/minimax_h3/test_review_fixes.py | 2 +
.../test_offline_output_capability_matrix.py | 10 +-
.../test_pipeline_contract_constructors.py | 92 ++
16 files changed, 2003 insertions(+), 37 deletions(-)
create mode 100644 src/flow_factory/models/minimax_h3/_output.py
create mode 100644 tests/models/minimax_h3/test_output_codec.py
create mode 100644 tests/models/test_pipeline_contract_constructors.py
diff --git a/.agents/knowledge/topics/adapter_conventions.md b/.agents/knowledge/topics/adapter_conventions.md
index e0c465dec..8f2ce5dd0 100644
--- a/.agents/knowledge/topics/adapter_conventions.md
+++ b/.agents/knowledge/topics/adapter_conventions.md
@@ -48,6 +48,8 @@ the complete mapping when its forward has different semantics or additional guid
- SenseNova neutralizes text and image guidance together and disables CFG normalization.
- Bagel replaces the base mapping with its actual `cfg_text_scale` / `cfg_img_scale` arguments;
it must not inherit an irrelevant `guidance_scale` key through its permissive `**kwargs`.
+- MiniMax H3 T2VA inherits neutral `guidance_scale=1.0`; its strict forward validates that
+ interface value even though the guidance-distilled checkpoint has no CFG branch.
Wan I2V and LTX2 remain behind explicit offline output-codec blockers. Wan I2V must mirror the two
neutral Wan transformer scales before it is enabled. LTX2 must set video/audio CFG scales and
diff --git a/.agents/knowledge/topics/minimax_h3.md b/.agents/knowledge/topics/minimax_h3.md
index a20ab1097..fd41f1962 100644
--- a/.agents/knowledge/topics/minimax_h3.md
+++ b/.agents/knowledge/topics/minimax_h3.md
@@ -53,6 +53,73 @@ and lets that boundary stay strict.
`samples/references.py`.
- PyAV >=18.0.0 decodes video/audio references, including embedded or separate soundtracks.
+## Offline T2VA output contract
+
+`minimax-h3-t2va` supports SFT and offline DPO with one exact ordered output pair:
+video first, then audio. Both `fps` and `sample_rate` are required in V2 supervision.
+Targets are decoded on demand; neither pixels, waveforms, nor VAE latents enter the condition
+cache. The pipeline's single-sample capability also forces condition-cache preprocessing to B=1,
+independently of the global preprocessing batch-size setting.
+
+The codec cross-validates cached T2VA layout and geometry against the current training config. It
+resamples video onto the configured fixed 24-fps grid and canvas, truncates audio on its declared
+source clock before a single conversion to the audio-VAE rate, and aligns stereo audio to the exact
+latent duration, samples and normalizes the video posterior, takes and normalizes the official
+audio posterior mode, then packs structured rows in `("video", "audio")` order. The codec does
+not duplicate input-owned fields in output forward context. Replay nests the flat cached layout
+and derives empty T2VA condition prefixes from the current state, preserving storage dtype and
+device. Exact velocity-only offline forwards return before either component scheduler steps, so
+SFT and offline DPO do not sample unused transitions or perturb scheduler RNG cadence. Every
+encoded row count must match the cached layout before transformer execution.
+
+FL2VA and Ref2VA remain online-only. Their output AV encoding can reuse the T2VA numerical
+codec, but their cached media conditions still need a separately owned, reproducible
+condition-prefix binder. In particular, both offline-DPO arms must consume the same conditioned
+prefix noise; do not generate those prefixes independently inside the chosen/rejected codecs.
+
+## Fix records
+
+### Offline targets preserve configured geometry and logical source clocks
+
+- **Date**: 2026-08-28
+- **Symptom**: An internally valid stale H3 condition cache could select another output canvas or
+ frame count, while an audio `sample_rate` override caused decoder resampling followed by a
+ second codec resample.
+- **Root Cause**: The output codec trusted cached geometry without comparing the current training
+ config, and the generic audio decoder treated source-rate metadata as a target decode rate.
+- **Fix**: The codec now cross-validates cached H/W and the officially aligned frame count against
+ current training arguments. Audio decoding preserves file samples; the codec truncates on the
+ declared source clock before exactly one model-rate conversion.
+- **Lesson**: Cached geometry must be checked against its configured authority, and media rate
+ overrides describe logical source clocks rather than preprocessing requests.
+- **Related Constraint**: #8, #26
+
+### Velocity-only H3 forwards bypass both schedulers
+
+- **Date**: 2026-08-28
+- **Symptom**: SFT and offline DPO requested only velocity but still sampled unused video/audio
+ scheduler transitions, wasting memory and changing RNG cadence.
+- **Root Cause**: The adapter boundary did not route the exact velocity-only request to the
+ existing scheduler-free `forward_h3_state` path.
+- **Fix**: The adapter detects a non-log-probability `("velocity",)` request with no replay next
+ state, returns `MultiModalStepOutput(velocity=...)`, and never enters either scheduler.
+- **Lesson**: Decoupled velocity objectives must stop at model prediction; requesting fewer return
+ fields is not sufficient if the adapter still executes transition side effects.
+- **Related Constraint**: #7
+
+### Single-sample capability governs condition preprocessing
+
+- **Date**: 2026-08-28
+- **Symptom**: A multi-row H3 offline manifest could reach its B=1 preprocessor with the global
+ condition-cache batch size, even though the training loader correctly rejected B>1.
+- **Root Cause**: The offline cache builder forced row-wise preprocessing only for ordered
+ references and did not apply the adapter's general batching capability.
+- **Fix**: The cache builder now derives its effective preprocessing batch size from both ordered
+ reference binding and `BatchCapability.SINGLE_SAMPLE`.
+- **Lesson**: One adapter-owned batching contract must govern preprocessing and model execution;
+ otherwise framework stages can disagree before training starts.
+- **Related Constraint**: #8, #12
+
## Verification boundary
All workflows have pinned API/schema/no-weight verification. T2VA additionally completed
@@ -68,6 +135,7 @@ improvement, convergence, or numerical parity.
- [ ] Rerun the real public-symbol and no-weight component-spec/workflow probes.
- [ ] Run H3 scheduler/runtime/registry/reference tests in the pinned environment.
- [ ] Parse all H3 examples through `Arguments.load_from_yaml`.
+- [ ] Run the T2VA output-codec and common SFT/offline-DPO structured-state tests.
- [ ] Rerun the documented T2VA real-weight smoke before changing support or memory claims.
## Cross-refs
diff --git a/README.md b/README.md
index 8e0b9ce45..0b9954858 100644
--- a/README.md
+++ b/README.md
@@ -99,8 +99,10 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
> **Offline output support:** SFT and offline DPO currently support `sd3-5`, `flux1`,
> `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`,
-> `bagel`, `sensenova`, and `wan2_t2v`. Wan I2V, LTX2, and MiniMax H3 fail fast on their
-> currently unresolved output/condition or audio-video semantics. See the
+> `bagel`, `sensenova`, `wan2_t2v`, and `minimax-h3-t2va`. Wan I2V, LTX2, and the
+> conditioned MiniMax H3 FL2VA/Ref2VA workflows fail fast on their currently unresolved
+> output/condition semantics. MiniMax H3 T2VA targets use an exact ordered video/audio pair,
+> encoded on demand into its structured latent state. See the
> [offline model matrix](guidance/datasets.md#offline-model-support).
> **MiniMax H3 status:** the T2VA debug and
diff --git a/guidance/datasets.md b/guidance/datasets.md
index c402f34d7..ff0d74726 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -108,9 +108,11 @@ must carry the supervision type required by its trainer. Prompt-only rows, mixed
preference rows, unknown keys, and non-V2 records fail during manifest loading.
All V2 media paths are resolved against that source's `dataset_dir`; an absolute path is retained.
-Images and videos have built-in CPU decoders. Video targets require PyAV 18 or newer. There is no
-default audio target decoder yet, which is one reason the current audio-video adapters are blocked
-for offline objectives.
+Images, videos, and audio have built-in CPU decoders. Video targets require PyAV 18 or newer.
+Decoded audio is a detached CPU `float32` waveform shaped `(channels, samples)`. A manifest
+`sample_rate` is a logical source-clock override and does not pre-resample the decoded samples;
+source-clock truncation, channel conversion, the single model-rate conversion, posterior selection,
+and latent packing remain adapter-owned.
Tiny schema-complete fixtures and configs are available for
[SFT](../examples/sft/lora/sd3_5/default.yaml) and
@@ -201,9 +203,20 @@ output semantics.
|---|---|---|
| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. |
| Supported | `wan2_t2v` | Video targets require `fps`; the codec resamples to configured frames/rate and samples the Wan VAE posterior on the fly. |
+| Supported | `minimax-h3-t2va` | Every candidate is an exact ordered `(video, audio)` pair with required `fps` and `sample_rate`. The codec aligns both streams to configured H3 geometry, samples the video posterior, takes the official audio-posterior mode, and packs structured video/audio rows on the fly. Condition preprocessing and training remain B=1. |
| Blocked | `wan2_i2v` | Output geometry depends on the first-frame VAE latent/mask, while the current condition cache does not preserve the source pixels needed by that binder. |
-| Blocked | `ltx2_t2av`, `ltx2_i2av` | Lossless audio decode/rate metadata and exact audio-video duration alignment are not unified; I2AV also needs the pinned first-frame active mask. |
-| Blocked | `minimax-h3-t2va`, `minimax-h3-fl2va`, `minimax-h3-ref2va` | The audio-video boundary and official target-video posterior policy are not yet defined for offline targets. |
+| Blocked | `ltx2_t2av`, `ltx2_i2av` | Their adapter codec still needs exact LTX-specific audio/video duration alignment, latent packing, and decode context; I2AV also needs the pinned first-frame active mask. |
+| Blocked | `minimax-h3-fl2va`, `minimax-h3-ref2va` | Their cached input media still need a shared, reproducible offline condition-prefix binder; offline DPO must reuse the same conditioned prefix noise for both preference arms. |
+
+MiniMax H3 T2VA supervision lists video first and audio second. The target video must cover the
+configured 24-fps duration; it is deterministically sampled onto that frame grid and resized to the
+configured canvas. Audio is converted to stereo at the H3 audio-VAE rate, then trimmed or
+right-padded to the exact aligned latent duration. For example:
+
+```jsonl
+{"schema_version":2,"input":{"prompt":"Ocean waves beneath an aurora.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/aurora.mp4","fps":24.0},{"type":"audio","path":"targets/aurora.wav","sample_rate":32000}]}},"metadata":{}}
+{"schema_version":2,"input":{"prompt":"Ocean waves beneath an aurora.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"video","path":"pairs/chosen.mp4","fps":24.0},{"type":"audio","path":"pairs/chosen.wav","sample_rate":32000}]},"rejected":{"media":[{"type":"video","path":"pairs/rejected.mp4","fps":24.0},{"type":"audio","path":"pairs/rejected.wav","sample_rate":32000}]}},"metadata":{}}
+```
## Common task formats
diff --git a/src/flow_factory/data_utils/offline_dataset.py b/src/flow_factory/data_utils/offline_dataset.py
index 0822179bf..9043a15bc 100644
--- a/src/flow_factory/data_utils/offline_dataset.py
+++ b/src/flow_factory/data_utils/offline_dataset.py
@@ -38,6 +38,8 @@
from pydantic import ValidationError
from torch.utils.data import Dataset
+from ..utils.audio import load_audio
+
try:
import av
except ImportError:
@@ -511,10 +513,57 @@ def decode_video(asset: MediaAsset) -> np.ndarray:
return np.ascontiguousarray(video, dtype=np.uint8)
+def decode_audio(asset: MediaAsset) -> torch.Tensor:
+ """Decode one target audio asset as a detached CPU waveform.
+
+ The returned contiguous ``float32`` tensor has shape ``(channels, samples)``.
+ The manifest ``sample_rate`` is a logical source-clock override, so decoding
+ intentionally preserves the file's samples instead of resampling them. The
+ :class:`DecodedMedia` boundary carries that override to the model codec, which
+ owns source-clock truncation, channel conversion, and the single model-rate
+ resample. This mirrors the source-clock semantics of video ``fps`` overrides.
+ Keeping this function at module scope makes the default decoder safe to pickle
+ under spawn-based DataLoader workers.
+
+ Args:
+ asset: Normalized audio reference with a resolved local path and optional rate override.
+
+ Returns:
+ Detached contiguous CPU waveform shaped ``(channels, samples)``.
+
+ Raises:
+ TypeError: If the audio backend returns a non-tensor or non-floating payload.
+ ValueError: If the decoded waveform is empty, non-finite, or not two-dimensional.
+ """
+ waveform = load_audio(asset.path, sample_rate=None)
+ if not isinstance(waveform, torch.Tensor):
+ raise TypeError(
+ f"failed to decode target audio {asset.path!r}: expected torch.Tensor, "
+ f"received {type(waveform).__name__}"
+ )
+ if not waveform.is_floating_point():
+ raise TypeError(
+ f"failed to decode target audio {asset.path!r}: expected floating waveform, "
+ f"received {waveform.dtype}"
+ )
+ waveform = waveform.detach().to(device="cpu", dtype=torch.float32).contiguous()
+ if waveform.ndim != 2 or waveform.shape[0] < 1 or waveform.shape[1] < 1:
+ raise ValueError(
+ f"failed to decode target audio {asset.path!r}: expected non-empty waveform "
+ f"shaped (channels,samples), received {tuple(waveform.shape)}"
+ )
+ if not torch.isfinite(waveform).all():
+ raise ValueError(
+ f"failed to decode target audio {asset.path!r}: waveform contains non-finite values"
+ )
+ return waveform
+
+
DEFAULT_MEDIA_DECODERS: Mapping[MediaType, MediaDecoder] = MappingProxyType(
{
"image": decode_image,
"video": decode_video,
+ "audio": decode_audio,
}
)
@@ -840,6 +889,7 @@ def _require_record_supervision(
"PreferenceOutputBatch",
"compute_offline_condition_id",
"compute_offline_record_id",
+ "decode_audio",
"decode_image",
"decode_video",
"load_offline_manifest",
diff --git a/src/flow_factory/data_utils/offline_train_data.py b/src/flow_factory/data_utils/offline_train_data.py
index 83f6d9493..4fc16575a 100644
--- a/src/flow_factory/data_utils/offline_train_data.py
+++ b/src/flow_factory/data_utils/offline_train_data.py
@@ -223,6 +223,7 @@ def build_offline_train_dataloader(
preprocess_func=preprocess_func,
preprocess_kwargs=normalized_preprocess_kwargs,
preprocessing_batch_size=data_args.preprocessing_batch_size,
+ batch_capability=pipeline_io_contract.batch_capability,
force_reprocess=data_args.force_reprocess,
extra_hash_strs=[*normalized_extra_hash_strs, f"offline_train_source:{source.name}"],
preprocess_parallelism=data_args.preprocess_parallelism,
@@ -314,6 +315,7 @@ def _build_distributed_condition_cache(
preprocess_func: PreprocessCallable,
preprocess_kwargs: Mapping[str, Any],
preprocessing_batch_size: int,
+ batch_capability: BatchCapability,
force_reprocess: bool,
extra_hash_strs: Sequence[str],
preprocess_parallelism: Literal["global", "local"],
@@ -322,7 +324,11 @@ def _build_distributed_condition_cache(
) -> HFDataset:
"""Build one input-only cache through the existing rank-safe orchestrator."""
ordered_references = _supports_ordered_references(preprocess_func)
- effective_batch_size = 1 if ordered_references else preprocessing_batch_size
+ effective_batch_size = (
+ 1
+ if ordered_references or batch_capability is BatchCapability.SINGLE_SAMPLE
+ else preprocessing_batch_size
+ )
raw_dataset = project_offline_condition_dataset(
records,
source_name=source_name,
diff --git a/src/flow_factory/models/minimax_h3/_output.py b/src/flow_factory/models/minimax_h3/_output.py
new file mode 100644
index 000000000..dff208c3f
--- /dev/null
+++ b/src/flow_factory/models/minimax_h3/_output.py
@@ -0,0 +1,891 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""On-the-fly audiovisual target encoding for MiniMax H3 T2VA."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from dataclasses import dataclass
+from numbers import Real
+from typing import Any, ClassVar, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ...contracts import MediaType
+from ...samples import LatentState
+from ...utils.audio import convert_audio
+from ..configured_image_output import retrieve_vae_latents
+from ..output_state import (
+ DecodedMediaBatch,
+ EncodedOutputState,
+ GeometrySignature,
+ MediaGeometrySignature,
+)
+from ._common import pack_audio_latents, pack_video_latents, validate_target_state
+from .workflow import _normalize_geometry, _normalize_layout
+
+_GEOMETRY_FIELDS = (
+ "height",
+ "width",
+ "num_frames",
+ "num_latent_frames",
+ "latent_height",
+ "latent_width",
+ "num_audio_latents",
+)
+_LAYOUT_FIELDS = (
+ "position_ids",
+ "token_tags",
+ "video_indices",
+ "audio_indices",
+ "text_indices",
+ "num_condition_video_rows",
+ "num_condition_audio_rows",
+)
+_RELEASED_PATCH_SIZE = (1, 2, 2)
+_RELEASED_VIDEO_LATENT_CHANNELS = 24
+_RELEASED_AUDIO_CHANNELS = 2
+_RELEASED_AUDIO_LATENT_CHANNELS = 32
+_RELEASED_AUDIO_SAMPLE_RATE = 32000
+_RELEASED_AUDIO_HOP_LENGTH = 800
+
+
+@dataclass(frozen=True, slots=True)
+class _H3ModelShape:
+ patch_size: Tuple[int, int, int]
+ video_latent_channels: int
+ audio_channels: int
+ audio_latent_channels: int
+
+
+@dataclass(frozen=True, slots=True)
+class MiniMaxH3AVOutputCodec:
+ """Encode one configured T2VA video/audio target into packed H3 rows."""
+
+ adapter: Any
+ required_components: ClassVar[Tuple[str, ...]] = ("vae", "audio_vae")
+
+ def encode_output_state(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> EncodedOutputState:
+ """Encode the exact ``(video, audio)`` target sequence for one sample.
+
+ Diffusers defines H3's conditioning encoder, but not an offline target
+ encoder. Video posterior sampling is therefore a framework policy inferred
+ from latent-diffusion training objectives. In particular, the condition-only
+ float16 rounding step is intentionally not applied to clean targets. Audio
+ targets use the posterior mode because the official H3 pipeline never samples
+ the audio posterior. Both components are normalized and packed exactly like
+ generated target rows.
+
+ Args:
+ media_batch: One decoded video/audio pair.
+ condition: Cached T2VA condition and its authoritative packed layout.
+ generator: Optional generator forwarded only to video posterior sampling.
+
+ Returns:
+ Detached clean video/audio rows plus the output-owned decode context.
+ """
+ if len(media_batch) != 1:
+ raise ValueError(
+ f"MiniMax H3 T2VA output encoding requires B=1, received B={len(media_batch)}"
+ )
+ candidate = media_batch[0]
+ if len(candidate) != 2:
+ raise ValueError(
+ "MiniMax H3 T2VA output codec expected exact (video, audio) media, "
+ f"received {len(candidate)} items"
+ )
+ video_media, audio_media = candidate
+ geometry = resolve_h3_output_geometry(self.adapter, condition)
+ model_shape = _resolve_h3_model_shape(self.adapter)
+
+ video_pixels = prepare_h3_target_video(
+ video_media.payload,
+ source_fps=video_media.fps,
+ target_frames=geometry["num_frames"],
+ target_fps=float(self.adapter.pipeline.fps),
+ height=geometry["height"],
+ width=geometry["width"],
+ )
+ video_latents = encode_h3_target_video(
+ self.adapter,
+ video_pixels,
+ generator=generator,
+ )
+ expected_video_shape = (
+ 1,
+ model_shape.video_latent_channels,
+ geometry["num_latent_frames"],
+ geometry["latent_height"],
+ geometry["latent_width"],
+ )
+ if tuple(video_latents.shape) != expected_video_shape:
+ raise ValueError(
+ "MiniMax H3 target video latent geometry mismatch: "
+ f"expected {expected_video_shape}, received {tuple(video_latents.shape)}"
+ )
+
+ audio_vae = self.adapter.get_component("audio_vae")
+ sample_rate = _positive_int(
+ getattr(audio_vae.config, "sampling_rate", None),
+ "audio_vae.config.sampling_rate",
+ )
+ hop_length = _positive_int(
+ getattr(audio_vae, "hop_length", None),
+ "audio_vae.hop_length",
+ )
+ target_audio_samples = geometry["num_audio_latents"] * hop_length
+ waveform = prepare_h3_target_audio(
+ audio_media.payload,
+ source_sample_rate=audio_media.sample_rate,
+ target_sample_rate=sample_rate,
+ target_samples=target_audio_samples,
+ target_duration_seconds=geometry["num_frames"] / float(self.adapter.pipeline.fps),
+ )
+ audio_latents = encode_h3_target_audio(self.adapter, waveform)
+ expected_audio_shape = (
+ model_shape.audio_channels,
+ model_shape.audio_latent_channels,
+ geometry["num_audio_latents"],
+ )
+ if tuple(audio_latents.shape) != expected_audio_shape:
+ raise ValueError(
+ "MiniMax H3 target audio latent geometry mismatch: "
+ f"expected {expected_audio_shape}, received {tuple(audio_latents.shape)}"
+ )
+
+ clean_state = LatentState(
+ {
+ "video": pack_video_latents(video_latents.to(torch.float32)),
+ "audio": pack_audio_latents(audio_latents.to(torch.float32).unsqueeze(0)),
+ }
+ )
+ validate_target_state(clean_state)
+ _validate_h3_t2va_input_layout(condition, clean_state)
+ frame_rate = float(self.adapter.pipeline.fps)
+ signature = GeometrySignature(
+ media=(
+ MediaGeometrySignature(
+ type=MediaType.VIDEO,
+ height=geometry["height"],
+ width=geometry["width"],
+ frames=geometry["num_frames"],
+ fps=frame_rate,
+ ),
+ MediaGeometrySignature(
+ type=MediaType.AUDIO,
+ samples=target_audio_samples,
+ sample_rate=sample_rate,
+ ),
+ )
+ )
+ return EncodedOutputState(
+ clean_state=clean_state,
+ forward_context={},
+ decode_context={"geometry": geometry},
+ geometry_signatures=(signature,),
+ )
+
+
+def resolve_h3_output_geometry(adapter: Any, condition: Mapping[str, Any]) -> Dict[str, int]:
+ """Resolve and validate configured geometry from the cached T2VA condition.
+
+ Args:
+ adapter: Active MiniMax H3 adapter with configured pipeline components.
+ condition: Cached input condition carrying flat H3 geometry fields.
+
+ Returns:
+ Canonical positive integer geometry validated against the training config.
+ """
+ geometry = _normalize_geometry(condition)
+ missing = tuple(field for field in _GEOMETRY_FIELDS if field not in geometry)
+ if missing:
+ raise ValueError(f"MiniMax H3 offline condition geometry missing fields={missing}")
+ for field in _GEOMETRY_FIELDS:
+ _positive_int(geometry[field], f"condition.{field}")
+
+ pipeline = adapter.pipeline
+ model_shape = _resolve_h3_model_shape(adapter)
+ frame_rate = _positive_real(getattr(pipeline, "fps", None), "pipeline.fps")
+ if frame_rate != 24.0:
+ raise ValueError(
+ f"MiniMax H3 target encoding requires fixed 24 fps, received {frame_rate!r}"
+ )
+ configured_frame_rate = _positive_real(
+ getattr(adapter.training_args, "frame_rate", None),
+ "training_args.frame_rate",
+ )
+ if configured_frame_rate != frame_rate:
+ raise ValueError(
+ "MiniMax H3 configured frame rate must match the model clock: "
+ f"expected {frame_rate}, received {configured_frame_rate}"
+ )
+ configured_height = _positive_int(
+ getattr(adapter.training_args, "height", None),
+ "training_args.height",
+ )
+ configured_width = _positive_int(
+ getattr(adapter.training_args, "width", None),
+ "training_args.width",
+ )
+ if (geometry["height"], geometry["width"]) != (
+ configured_height,
+ configured_width,
+ ):
+ raise ValueError(
+ "MiniMax H3 cached output canvas does not match the current configured geometry: "
+ f"expected {(configured_height, configured_width)}, received "
+ f"{(geometry['height'], geometry['width'])}"
+ )
+
+ frames_per_chunk = _positive_int(
+ getattr(pipeline, "vae_frames_per_chunk", None),
+ "pipeline.vae_frames_per_chunk",
+ )
+ latents_per_chunk = _positive_int(
+ getattr(pipeline, "vae_latents_per_chunk", None),
+ "pipeline.vae_latents_per_chunk",
+ )
+ if latents_per_chunk >= frames_per_chunk:
+ raise ValueError(
+ "MiniMax H3 video VAE chunk geometry requires latents_per_chunk < "
+ f"frames_per_chunk, received {(latents_per_chunk, frames_per_chunk)}"
+ )
+ configured_num_frames = _positive_int(
+ getattr(adapter.training_args, "num_frames", None),
+ "training_args.num_frames",
+ )
+ aligned_configured_num_frames = configured_num_frames + (
+ (latents_per_chunk - configured_num_frames) % frames_per_chunk
+ )
+ if geometry["num_frames"] != aligned_configured_num_frames:
+ raise ValueError(
+ "MiniMax H3 cached output frame count does not match the current configured "
+ "frame count after official VAE alignment: "
+ f"expected {aligned_configured_num_frames} from configured "
+ f"num_frames={configured_num_frames}, received {geometry['num_frames']}"
+ )
+
+ min_duration = _positive_real(
+ getattr(pipeline, "min_duration", None),
+ "pipeline.min_duration",
+ )
+ max_duration = _positive_real(
+ getattr(pipeline, "max_duration", None),
+ "pipeline.max_duration",
+ )
+ if min_duration > max_duration:
+ raise ValueError(
+ "MiniMax H3 pipeline duration bounds are inverted: "
+ f"min_duration={min_duration}, max_duration={max_duration}"
+ )
+ duration = geometry["num_frames"] / frame_rate
+ if not min_duration <= duration <= max_duration:
+ raise ValueError(
+ "MiniMax H3 configured output duration is outside the pipeline contract: "
+ f"expected [{min_duration}, {max_duration}] seconds, received {duration} "
+ f"from {geometry['num_frames']} frames at {frame_rate} fps"
+ )
+
+ spatial_ratio = _positive_int(
+ getattr(pipeline, "vae_spatial_compression_ratio", None),
+ "pipeline.vae_spatial_compression_ratio",
+ )
+ expected_latent_height = geometry["height"] // spatial_ratio
+ expected_latent_width = geometry["width"] // spatial_ratio
+ if geometry["height"] % spatial_ratio or geometry["width"] % spatial_ratio:
+ raise ValueError(
+ "MiniMax H3 output height/width must be divisible by the video VAE spatial ratio "
+ f"{spatial_ratio}, received {(geometry['height'], geometry['width'])}"
+ )
+ if (geometry["latent_height"], geometry["latent_width"]) != (
+ expected_latent_height,
+ expected_latent_width,
+ ):
+ raise ValueError(
+ "MiniMax H3 cached spatial latent geometry mismatch: expected "
+ f"{(expected_latent_height, expected_latent_width)}, received "
+ f"{(geometry['latent_height'], geometry['latent_width'])}"
+ )
+ _, patch_height, patch_width = model_shape.patch_size
+ if geometry["latent_height"] % patch_height or geometry["latent_width"] % patch_width:
+ raise ValueError(
+ "MiniMax H3 cached latent height/width must be divisible by transformer patch "
+ f"{model_shape.patch_size}, received "
+ f"{(geometry['latent_height'], geometry['latent_width'])}"
+ )
+
+ num_frames = geometry["num_frames"]
+ if num_frames % frames_per_chunk != latents_per_chunk:
+ raise ValueError(
+ "MiniMax H3 target num_frames must satisfy the video VAE chunk geometry "
+ f"F % {frames_per_chunk} == {latents_per_chunk}, received {num_frames}"
+ )
+ expected_video_latents = (
+ num_frames - latents_per_chunk
+ ) // frames_per_chunk * latents_per_chunk + 2
+ if geometry["num_latent_frames"] != expected_video_latents:
+ raise ValueError(
+ "MiniMax H3 cached temporal latent geometry mismatch: expected "
+ f"{expected_video_latents}, received {geometry['num_latent_frames']}"
+ )
+ patch_time = model_shape.patch_size[0]
+ if geometry["num_latent_frames"] % patch_time:
+ raise ValueError(
+ "MiniMax H3 cached latent frame count must be divisible by transformer temporal "
+ f"patch {patch_time}, received {geometry['num_latent_frames']}"
+ )
+
+ audio_vae = adapter.get_component("audio_vae")
+ sample_rate = _positive_int(
+ getattr(audio_vae.config, "sampling_rate", None),
+ "audio_vae.config.sampling_rate",
+ )
+ hop_length = _positive_int(
+ getattr(audio_vae, "hop_length", None),
+ "audio_vae.hop_length",
+ )
+ pipeline_sample_rate = _positive_int(
+ getattr(pipeline, "audio_sampling_rate", None),
+ "pipeline.audio_sampling_rate",
+ )
+ if (
+ sample_rate,
+ pipeline_sample_rate,
+ hop_length,
+ ) != (
+ _RELEASED_AUDIO_SAMPLE_RATE,
+ _RELEASED_AUDIO_SAMPLE_RATE,
+ _RELEASED_AUDIO_HOP_LENGTH,
+ ):
+ raise ValueError(
+ "MiniMax H3 output codec requires the released audio clock "
+ f"sample_rate/hop_length={(_RELEASED_AUDIO_SAMPLE_RATE, _RELEASED_AUDIO_HOP_LENGTH)}, "
+ f"received config/pipeline/hop={(sample_rate, pipeline_sample_rate, hop_length)}"
+ )
+ expected_audio_latents = int(round(num_frames / frame_rate * sample_rate / hop_length))
+ if geometry["num_audio_latents"] != expected_audio_latents:
+ raise ValueError(
+ "MiniMax H3 cached audio latent geometry mismatch: expected "
+ f"{expected_audio_latents}, received {geometry['num_audio_latents']}"
+ )
+ return geometry
+
+
+def prepare_h3_target_video(
+ payload: Any,
+ *,
+ source_fps: Any,
+ target_frames: int,
+ target_fps: float,
+ height: int,
+ width: int,
+) -> torch.Tensor:
+ """Resample decoded RGB frames onto the configured H3 video grid.
+
+ Args:
+ payload: Decoded uint8 RGB frames shaped ``(F, H, W, 3)``.
+ source_fps: Logical source frame rate from the media manifest.
+ target_frames: Exact aligned output frame count.
+ target_fps: Fixed model frame rate.
+ height: Configured output height.
+ width: Configured output width.
+
+ Returns:
+ Float32 pixels shaped ``(1, 3, F, H, W)`` in the unit interval.
+ """
+ if not isinstance(payload, np.ndarray):
+ raise TypeError(
+ "MiniMax H3 target video expected a decoded NumPy array, "
+ f"received {type(payload).__name__}"
+ )
+ if payload.dtype != np.uint8 or payload.ndim != 4 or payload.shape[-1] != 3:
+ raise ValueError(
+ "MiniMax H3 target video must be uint8 RGB shaped (F,H,W,3), "
+ f"received dtype={payload.dtype}, shape={tuple(payload.shape)}"
+ )
+ if payload.shape[0] < 1:
+ raise ValueError("MiniMax H3 target video must contain at least one frame")
+ source_fps = _positive_real(source_fps, "target video fps")
+ target_fps = _positive_real(target_fps, "model video fps")
+ frames = payload
+ if source_fps != target_fps:
+ # Match H3's official ffmpeg-style fps filter: each source frame is held
+ # until the rounded slot of the next frame, including the stream endpoint.
+ scale = target_fps / source_fps
+ slots = np.floor(np.arange(frames.shape[0]) * scale + 0.5).astype(np.int64)
+ endpoint = math.floor(frames.shape[0] * scale + 0.5)
+ frames = np.repeat(frames, np.diff(slots, append=endpoint), axis=0)
+ if frames.shape[0] < target_frames:
+ required_duration = target_frames / target_fps
+ available_duration = frames.shape[0] / target_fps
+ raise ValueError(
+ "MiniMax H3 target video is too short for configured temporal geometry: "
+ f"requires {target_frames} frames/{required_duration:.6f}s after rate conversion, "
+ f"has {frames.shape[0]} frames/{available_duration:.6f}s"
+ )
+ frames = np.ascontiguousarray(frames[:target_frames])
+ if frames.shape[1:3] != (height, width):
+ frames = np.stack(
+ [
+ np.asarray(Image.fromarray(frame).resize((width, height), Image.Resampling.LANCZOS))
+ for frame in frames
+ ]
+ )
+ pixels = torch.from_numpy(np.ascontiguousarray(frames)).permute(3, 0, 1, 2).unsqueeze(0)
+ return pixels.to(torch.float32).div_(255.0)
+
+
+def encode_h3_target_video(
+ adapter: Any,
+ pixel_values: torch.Tensor,
+ *,
+ generator: Optional[torch.Generator],
+) -> torch.Tensor:
+ """Apply the framework's sampled-posterior policy for clean H3 targets.
+
+ Diffusers specifies a fixed-seed sample followed by float16 rounding only for
+ H3 *conditions*. Offline clean targets instead use the caller's generator and
+ retain the sampled posterior in float32 before normalization. This distinction
+ is an explicit framework inference from latent training, not an official H3
+ target-encoding recipe.
+
+ Args:
+ adapter: Active MiniMax H3 adapter exposing the video VAE.
+ pixel_values: Normalized-shape pixels ``(1, 3, F, H, W)`` in ``[0, 1]``.
+ generator: Optional posterior sampling generator.
+
+ Returns:
+ Normalized video latents shaped ``(1, 24, F', H', W')``.
+ """
+ if not isinstance(pixel_values, torch.Tensor) or pixel_values.ndim != 5:
+ raise ValueError(
+ "MiniMax H3 target video pixels expected BCFHW tensor, "
+ f"received {type(pixel_values).__name__}/{getattr(pixel_values, 'shape', None)}"
+ )
+ vae = adapter.get_component("vae")
+ device = torch.device(adapter.device)
+ pixel_mean = torch.as_tensor(
+ adapter.pipeline.pixel_mean,
+ device=device,
+ dtype=torch.float32,
+ ).view(1, -1, 1, 1, 1)
+ pixel_std = torch.as_tensor(
+ adapter.pipeline.pixel_std,
+ device=device,
+ dtype=torch.float32,
+ ).view(1, -1, 1, 1, 1)
+ normalized_pixels = (
+ pixel_values.to(device=device, dtype=torch.float32) - pixel_mean
+ ) / pixel_std
+ encoded = vae.encode(normalized_pixels)
+ latents = retrieve_vae_latents(
+ encoded,
+ sample_mode="sample",
+ generator=generator,
+ source="MiniMax H3 target video",
+ ).to(torch.float32)
+ latent_mean, latent_std = _latent_statistics(
+ vae.config,
+ channels=latents.shape[1],
+ rank=5,
+ device=latents.device,
+ source="MiniMax H3 video VAE",
+ )
+ return (latents - latent_mean) / latent_std
+
+
+def prepare_h3_target_audio(
+ payload: Any,
+ *,
+ source_sample_rate: Any,
+ target_sample_rate: int,
+ target_samples: int,
+ target_duration_seconds: float,
+) -> torch.Tensor:
+ """Convert one waveform to exact stereo H3 audio-grid geometry.
+
+ The source-clock truncation precedes resampling, matching Diffusers' H3
+ reference normalization. The final trim/pad then accounts for audio-latent
+ rounding so the waveform lands on the exact model grid.
+
+ Args:
+ payload: Decoded mono or stereo float waveform shaped ``(C, S)``.
+ source_sample_rate: Logical source clock from the media manifest.
+ target_sample_rate: Audio VAE sample rate.
+ target_samples: Exact number of samples required by the audio latent grid.
+ target_duration_seconds: Aligned AV duration used for source-clock truncation.
+
+ Returns:
+ Contiguous float32 stereo waveform shaped ``(2, target_samples)``.
+ """
+ if not isinstance(payload, torch.Tensor):
+ raise TypeError(
+ "MiniMax H3 target audio expected a decoded torch.Tensor, "
+ f"received {type(payload).__name__}"
+ )
+ if payload.ndim != 2 or payload.shape[0] not in (1, 2) or payload.shape[1] < 1:
+ raise ValueError(
+ "MiniMax H3 target audio must be non-empty mono/stereo shaped (C,S), "
+ f"received {tuple(payload.shape)}"
+ )
+ if not payload.is_floating_point():
+ raise TypeError(
+ f"MiniMax H3 target audio expected floating waveform, received {payload.dtype}"
+ )
+ if not torch.isfinite(payload).all():
+ raise ValueError("MiniMax H3 target audio contains non-finite samples")
+ source_sample_rate = _positive_int(source_sample_rate, "target audio sample_rate")
+ target_sample_rate = _positive_int(target_sample_rate, "model audio sample_rate")
+ target_samples = _positive_int(target_samples, "target audio samples")
+ target_duration_seconds = _positive_real(
+ target_duration_seconds,
+ "target audio duration",
+ )
+ source_samples = int(target_duration_seconds * source_sample_rate)
+ if source_samples < 1:
+ raise ValueError(
+ "MiniMax H3 target audio duration resolves to fewer than one source sample: "
+ f"duration={target_duration_seconds}, sample_rate={source_sample_rate}"
+ )
+ source_waveform = payload.detach().to(device="cpu", dtype=torch.float32)[:, :source_samples]
+ waveform = convert_audio(
+ source_waveform,
+ from_rate=source_sample_rate,
+ to_rate=target_sample_rate,
+ to_channels=2,
+ )
+ if waveform.shape[-1] >= target_samples:
+ waveform = waveform[:, :target_samples]
+ else:
+ waveform = torch.nn.functional.pad(waveform, (0, target_samples - waveform.shape[-1]))
+ return waveform.contiguous()
+
+
+def encode_h3_target_audio(adapter: Any, waveform: torch.Tensor) -> torch.Tensor:
+ """Take and normalize the official H3 audio posterior mode.
+
+ Args:
+ adapter: Active MiniMax H3 adapter exposing the audio VAE.
+ waveform: Exact stereo model-rate waveform shaped ``(2, S)``.
+
+ Returns:
+ Normalized channel-major audio latents shaped ``(2, 32, F)``.
+ """
+ audio_vae = adapter.get_component("audio_vae")
+ device = torch.device(adapter.device)
+ posterior = audio_vae.encode(waveform.to(device=device)[:, None])
+ latents = retrieve_vae_latents(
+ posterior,
+ sample_mode="argmax",
+ source="MiniMax H3 target audio",
+ ).to(torch.float32)
+ latent_mean, latent_std = _latent_statistics(
+ audio_vae.config,
+ channels=latents.shape[1],
+ rank=3,
+ device=latents.device,
+ source="MiniMax H3 audio VAE",
+ )
+ return (latents - latent_mean) / latent_std
+
+
+def _validate_h3_t2va_input_layout(
+ condition: Mapping[str, Any],
+ clean_state: LatentState,
+) -> None:
+ """Validate the authoritative flat T2VA layout retained in input conditions."""
+ layout = _normalize_layout(condition)
+ missing = tuple(field for field in _LAYOUT_FIELDS if field not in layout)
+ if missing:
+ raise ValueError(f"MiniMax H3 offline condition layout missing fields={missing}")
+
+ position_ids = layout["position_ids"]
+ token_tags = layout["token_tags"]
+ index_tensors = [layout[field] for field in ("video_indices", "audio_indices", "text_indices")]
+ if (
+ not isinstance(position_ids, torch.Tensor)
+ or position_ids.ndim != 2
+ or position_ids.shape[-1] != 3
+ or position_ids.dtype != torch.float64
+ ):
+ raise ValueError(
+ "MiniMax H3 T2VA position_ids expected float64 shape (N,3), "
+ f"received {type(position_ids).__name__}/{getattr(position_ids, 'shape', None)}/"
+ f"{getattr(position_ids, 'dtype', None)}"
+ )
+ if (
+ not isinstance(token_tags, torch.Tensor)
+ or token_tags.ndim != 1
+ or token_tags.dtype != torch.long
+ ):
+ raise ValueError(
+ "MiniMax H3 T2VA token_tags expected one-dimensional torch.long, "
+ f"received {type(token_tags).__name__}/{getattr(token_tags, 'shape', None)}/"
+ f"{getattr(token_tags, 'dtype', None)}"
+ )
+ for field, values in zip(("video_indices", "audio_indices", "text_indices"), index_tensors):
+ if not isinstance(values, torch.Tensor) or values.ndim != 1 or values.dtype != torch.long:
+ raise ValueError(
+ f"MiniMax H3 T2VA {field} expected one-dimensional torch.long, "
+ f"received {type(values).__name__}/{getattr(values, 'shape', None)}/"
+ f"{getattr(values, 'dtype', None)}"
+ )
+
+ for component in ("video", "audio"):
+ count_field = f"num_condition_{component}_rows"
+ if layout[count_field] != 0:
+ raise ValueError(
+ "MiniMax H3 T2VA offline layout requires no condition rows, "
+ f"received {count_field}={layout[count_field]}"
+ )
+ component_indices = layout[f"{component}_indices"]
+ expected_rows = clean_state.components[component].shape[1]
+ if component_indices.numel() != expected_rows:
+ raise ValueError(
+ f"MiniMax H3 T2VA {component} layout expected {expected_rows} target rows, "
+ f"received {component_indices.numel()} indices"
+ )
+
+ sequence_length = sum(values.numel() for values in index_tensors)
+ if position_ids.shape[0] != sequence_length or token_tags.numel() != sequence_length:
+ raise ValueError(
+ "MiniMax H3 T2VA flat layout sequence lengths disagree: "
+ f"indices={sequence_length}, position_ids={position_ids.shape[0]}, "
+ f"token_tags={token_tags.numel()}"
+ )
+ devices = {
+ position_ids.device,
+ token_tags.device,
+ *(values.device for values in index_tensors),
+ }
+ if len(devices) != 1:
+ raise ValueError(
+ f"MiniMax H3 T2VA flat layout tensors must share one device, got {devices}"
+ )
+ permutation = torch.cat(index_tensors).sort().values
+ expected_permutation = torch.arange(
+ sequence_length,
+ dtype=torch.long,
+ device=permutation.device,
+ )
+ if not torch.equal(permutation, expected_permutation):
+ raise ValueError(
+ "MiniMax H3 T2VA video/audio/text indices must partition the packed sequence"
+ )
+
+
+def validate_h3_encoded_output_geometry(
+ adapter: Any,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+) -> None:
+ """Prove codec geometry and contexts agree with the cached T2VA layout.
+
+ Args:
+ adapter: Active MiniMax H3 adapter.
+ media_batch: Validated exact audiovisual target batch.
+ condition: Cached input condition with geometry and packed-row layout.
+ encoded: Encoded structured target state to validate.
+
+ Returns:
+ None after every geometry, component, and context invariant is proven.
+ """
+ geometry = resolve_h3_output_geometry(adapter, condition)
+ model_shape = _resolve_h3_model_shape(adapter)
+ audio_vae = adapter.get_component("audio_vae")
+ sample_rate = _positive_int(
+ getattr(audio_vae.config, "sampling_rate", None),
+ "audio_vae.config.sampling_rate",
+ )
+ hop_length = _positive_int(
+ getattr(audio_vae, "hop_length", None),
+ "audio_vae.hop_length",
+ )
+ expected_signature = GeometrySignature(
+ media=(
+ MediaGeometrySignature(
+ type=MediaType.VIDEO,
+ height=geometry["height"],
+ width=geometry["width"],
+ frames=geometry["num_frames"],
+ fps=float(adapter.pipeline.fps),
+ ),
+ MediaGeometrySignature(
+ type=MediaType.AUDIO,
+ samples=geometry["num_audio_latents"] * hop_length,
+ sample_rate=sample_rate,
+ ),
+ )
+ )
+ if len(media_batch) != 1 or encoded.geometry_signatures != (expected_signature,):
+ raise ValueError(
+ "MiniMax H3 encoded output geometry disagrees with configured audiovisual geometry: "
+ f"expected {(expected_signature,)!r}, received {encoded.geometry_signatures!r}"
+ )
+ if encoded.decode_context.get("geometry") != geometry:
+ raise ValueError(
+ "MiniMax H3 decode_context geometry disagrees with cached condition geometry: "
+ f"expected {geometry!r}, received {encoded.decode_context.get('geometry')!r}"
+ )
+
+ validate_target_state(encoded.clean_state)
+ if encoded.clean_state.component_names != ("video", "audio"):
+ raise ValueError(
+ "MiniMax H3 clean target components must be ordered ('video', 'audio'), "
+ f"received {encoded.clean_state.component_names}"
+ )
+ patch_time, patch_height, patch_width = model_shape.patch_size
+ expected_shapes = {
+ "video": (
+ 1,
+ geometry["num_latent_frames"]
+ // patch_time
+ * (geometry["latent_height"] // patch_height)
+ * (geometry["latent_width"] // patch_width),
+ model_shape.video_latent_channels * patch_time * patch_height * patch_width,
+ ),
+ "audio": (
+ 1,
+ model_shape.audio_channels * geometry["num_audio_latents"],
+ model_shape.audio_latent_channels,
+ ),
+ }
+ for component, expected_shape in expected_shapes.items():
+ received_shape = tuple(encoded.clean_state.components[component].shape)
+ if received_shape != expected_shape:
+ raise ValueError(
+ f"MiniMax H3 clean {component} rows expected shape {expected_shape}, "
+ f"received {received_shape}"
+ )
+
+ if encoded.forward_context:
+ raise ValueError(
+ "MiniMax H3 T2VA output codec must not duplicate input-owned layout/prefix fields, "
+ f"received keys={tuple(encoded.forward_context)}"
+ )
+ _validate_h3_t2va_input_layout(condition, encoded.clean_state)
+
+
+def _resolve_h3_model_shape(adapter: Any) -> _H3ModelShape:
+ """Validate the packing dimensions of the released H3 checkpoint contract."""
+ pipeline = adapter.pipeline
+ patch_size = getattr(pipeline, "patch_size", None)
+ if not isinstance(patch_size, (tuple, list)) or len(patch_size) != 3:
+ raise TypeError(
+ "pipeline.patch_size expected a length-3 tuple/list, "
+ f"received {type(patch_size).__name__}: {patch_size!r}"
+ )
+ patch_size = tuple(
+ _positive_int(value, f"pipeline.patch_size[{index}]")
+ for index, value in enumerate(patch_size)
+ )
+ video_latent_channels = _positive_int(
+ getattr(pipeline, "vae_latent_channels", None),
+ "pipeline.vae_latent_channels",
+ )
+ audio_channels = _positive_int(
+ getattr(pipeline, "audio_channels", None),
+ "pipeline.audio_channels",
+ )
+ audio_latent_channels = _positive_int(
+ getattr(pipeline, "audio_latent_channels", None),
+ "pipeline.audio_latent_channels",
+ )
+ received = (
+ patch_size,
+ video_latent_channels,
+ audio_channels,
+ audio_latent_channels,
+ )
+ expected = (
+ _RELEASED_PATCH_SIZE,
+ _RELEASED_VIDEO_LATENT_CHANNELS,
+ _RELEASED_AUDIO_CHANNELS,
+ _RELEASED_AUDIO_LATENT_CHANNELS,
+ )
+ if received != expected:
+ raise ValueError(
+ "MiniMax H3 output codec supports the released packing contract "
+ "patch/video_channels/audio_channels/audio_latent_channels="
+ f"{expected!r}, received {received!r}"
+ )
+ return _H3ModelShape(
+ patch_size=patch_size,
+ video_latent_channels=video_latent_channels,
+ audio_channels=audio_channels,
+ audio_latent_channels=audio_latent_channels,
+ )
+
+
+def _latent_statistics(
+ config: Any,
+ *,
+ channels: int,
+ rank: int,
+ device: torch.device,
+ source: str,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+ raw_mean = getattr(config, "latents_mean", None)
+ raw_std = getattr(config, "latents_std", None)
+ if raw_mean is None or raw_std is None:
+ raise ValueError(f"{source} config must define latents_mean and latents_std")
+ mean = torch.as_tensor(raw_mean, device=device, dtype=torch.float32)
+ std = torch.as_tensor(raw_std, device=device, dtype=torch.float32)
+ if mean.shape != (channels,) or std.shape != (channels,):
+ raise ValueError(
+ f"{source} expected per-channel latent statistics shaped ({channels},), "
+ f"received mean={tuple(mean.shape)}, std={tuple(std.shape)}"
+ )
+ if not torch.isfinite(mean).all() or not torch.isfinite(std).all() or torch.any(std <= 0):
+ raise ValueError(f"{source} latent statistics must be finite with strictly positive std")
+ shape = [1, channels, *([1] * (rank - 2))]
+ return mean.view(shape), std.view(shape)
+
+
+def _positive_int(value: Any, source: str) -> int:
+ if not isinstance(value, int) or isinstance(value, bool) or value <= 0:
+ raise ValueError(f"{source} expected positive int, received {value!r}")
+ return value
+
+
+def _positive_real(value: Any, source: str) -> float:
+ if isinstance(value, bool) or not isinstance(value, Real):
+ raise TypeError(
+ f"{source} expected positive finite real, received {type(value).__name__}: {value!r}"
+ )
+ value = float(value)
+ if not math.isfinite(value) or value <= 0:
+ raise ValueError(f"{source} expected positive finite real, received {value!r}")
+ return value
+
+
+__all__ = [
+ "MiniMaxH3AVOutputCodec",
+ "encode_h3_target_audio",
+ "encode_h3_target_video",
+ "prepare_h3_target_audio",
+ "prepare_h3_target_video",
+ "resolve_h3_output_geometry",
+ "validate_h3_encoded_output_geometry",
+]
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index ca2650573..6c0a001c3 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -18,6 +18,11 @@
import torch
+from ...contracts import (
+ BatchCapability,
+ GeometrySource,
+ NegativePromptPolicy,
+)
from ...samples import (
ComponentTimes,
LatentState,
@@ -28,8 +33,11 @@
from ...scheduler import MiniMaxH3SDEScheduler, SchedulerGroup
from ..abc import BaseAdapter
from ..checkpointing import CheckpointUnit
+from ..output_state import DecodedMediaBatch, EncodedOutputState, OutputStateCodec
+from ..pipeline_contracts import audio_video_output_contract
from ..runtime import ModularPipelineRuntime
from ._common import apply_forward_process_noise, draw_forward_process_noise
+from ._output import MiniMaxH3AVOutputCodec, validate_h3_encoded_output_geometry
from .workflow import (
build_h3_component_runtime,
build_h3_replay_forward_kwargs,
@@ -183,7 +191,11 @@ def _forward_state(
compute_log_prob=compute_log_prob,
return_fields=return_fields,
noise_level=noise_level,
- **build_h3_replay_forward_kwargs(forward_kwargs),
+ **build_h3_replay_forward_kwargs(
+ forward_kwargs,
+ state=state,
+ workflow=self.workflow,
+ ),
)
def forward(self, **kwargs: Any) -> MultiModalStepOutput:
@@ -193,9 +205,10 @@ def forward(self, **kwargs: Any) -> MultiModalStepOutput:
class MiniMaxH3T2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"""Load the workflow-pruned MiniMax H3 text-to-video-audio partition."""
- output_state_codec_unavailable_reason = (
- "MiniMax H3 offline targets require lossless audiovisual decoding/alignment, and the "
- "official target-video posterior policy is not defined by the inference encoders"
+ pipeline_io_contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
)
workflow: ClassVar[str] = "t2va"
@@ -208,13 +221,50 @@ class MiniMaxH3T2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"audio_vae",
]
+ def build_output_state_codec(self) -> OutputStateCodec:
+ """Declare the configured audiovisual target codec without loading components.
+
+ Returns:
+ Immutable MiniMax H3 audiovisual output codec declaration.
+ """
+ return MiniMaxH3AVOutputCodec(self)
+
+ def _validate_encoded_output_geometry(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+ ) -> None:
+ """Require encoded rows and rate metadata to match cached T2VA geometry."""
+ validate_h3_encoded_output_geometry(self, media_batch, condition, encoded)
+
+ def _decode_output_state(
+ self,
+ encoded: EncodedOutputState,
+ *,
+ output_type: Literal["pil", "pt", "np"],
+ ) -> Any:
+ """Decode the two-component target state through H3's existing decoder."""
+ geometry = encoded.decode_context.get("geometry")
+ if not isinstance(geometry, Mapping):
+ raise TypeError(
+ "MiniMax H3 decode_context requires a geometry mapping, "
+ f"received {type(geometry).__name__}: {geometry!r}"
+ )
+ return self.decode_latents(
+ encoded.clean_state,
+ geometry=geometry,
+ output_type=output_type,
+ )
+
class MiniMaxH3FL2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"""Load the workflow-pruned MiniMax H3 first/last-frame partition."""
output_state_codec_unavailable_reason = (
- "MiniMax H3 offline targets require lossless audiovisual decoding/alignment, and the "
- "official target-video posterior policy is not defined by the inference encoders"
+ "MiniMax H3 conditioned offline forward requires a shared, reproducible "
+ "conditioned-prefix binder; paired offline-DPO arms must reuse identical "
+ "condition posterior noise"
)
workflow: ClassVar[str] = "fl2va"
@@ -238,8 +288,9 @@ class MiniMaxH3Ref2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"""Load the workflow-pruned MiniMax H3 omni-reference partition."""
output_state_codec_unavailable_reason = (
- "MiniMax H3 offline targets require lossless audiovisual decoding/alignment, and the "
- "official target-video posterior policy is not defined by the inference encoders"
+ "MiniMax H3 conditioned offline forward requires a shared, reproducible "
+ "conditioned-prefix binder; paired offline-DPO arms must reuse identical "
+ "condition posterior noise"
)
workflow: ClassVar[str] = "ref2va"
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index 5edb05f7b..abc53f797 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -25,6 +25,7 @@
MiniMaxH3FL2VASample,
MiniMaxH3Ref2VASample,
MiniMaxH3T2VASample,
+ MultiModalStepOutput,
StructuredTrajectory,
)
from ...scheduler import MiniMaxH3SDEScheduler, SchedulerGroup
@@ -475,7 +476,10 @@ def forward_h3_adapter_state(
f"received prompt B={prompt_embeds.shape[0]}"
)
layout = _normalize_layout({"layout": forward_kwargs["layout"]})
- return forward_h3_state(
+ velocity_only = (
+ not compute_log_prob and next_state is None and tuple(return_fields) == ("velocity",)
+ )
+ result = forward_h3_state(
adapter.get_component(adapter.transformer_component_name),
state,
condition_prefixes,
@@ -488,29 +492,79 @@ def forward_h3_adapter_state(
generator=forward_kwargs.get("generator"),
noise_level=noise_level,
compute_log_prob=compute_log_prob,
+ velocity_only=velocity_only,
attention_kwargs=forward_kwargs.get("attention_kwargs"),
return_kwargs=return_fields,
workflow=adapter.workflow,
)
+ if velocity_only:
+ if not isinstance(result, LatentState):
+ raise TypeError(
+ "MiniMax H3 velocity-only forward expected LatentState, "
+ f"received {type(result).__name__}"
+ )
+ return MultiModalStepOutput(velocity=result)
+ return result
-def build_h3_replay_forward_kwargs(forward_kwargs: Mapping[str, Any]) -> Dict[str, Any]:
+def build_h3_replay_forward_kwargs(
+ forward_kwargs: Mapping[str, Any],
+ *,
+ state: Optional[LatentState] = None,
+ workflow: Optional[str] = None,
+) -> Dict[str, Any]:
"""Select the conditioning arguments H3 ``forward`` accepts from a replay batch.
Replay wrappers receive every collated batch field, while ``forward`` is a strict
- public boundary. Selecting here keeps that boundary strict for rollout callers.
+ public boundary. Offline T2VA cache rows retain their authoritative layout as flat
+ input fields and have no condition latents. This binder nests that layout and builds
+ empty prefixes from the current state, so storage casts cannot leave prefix dtype or
+ device stale. Online replay already supplies both nested fields and keeps that path.
Args:
forward_kwargs: Conditioning arguments resolved from the stored batch.
+ state: Current target state, required to bind missing T2VA prefixes.
+ workflow: H3 workflow identifier, required when prefixes are missing.
Returns:
Conditioning arguments accepted by ``forward``.
"""
- return {
+ selected = {
name: value
for name, value in forward_kwargs.items()
if name in _H3_FORWARD_CONDITIONING_FIELDS
}
+ if "layout" not in selected:
+ layout = _normalize_layout(forward_kwargs)
+ required_layout_fields = (
+ *_LAYOUT_MATRIX_FIELDS,
+ *_LAYOUT_INDEX_FIELDS,
+ *_LAYOUT_COUNT_FIELDS,
+ )
+ missing = tuple(field for field in required_layout_fields if field not in layout)
+ if missing:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} replay flat layout missing fields={missing}"
+ )
+ selected["layout"] = layout
+
+ if "condition_prefixes" not in selected:
+ if workflow != "t2va":
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} replay requires a shared, reproducible "
+ "conditioned-prefix binder"
+ )
+ if not isinstance(state, LatentState):
+ raise TypeError(
+ "MiniMax H3 T2VA replay requires LatentState to bind empty prefixes, "
+ f"received {type(state).__name__}"
+ )
+ validate_target_state(state)
+ selected["condition_prefixes"] = {
+ component: values.new_empty((values.shape[0], 0, values.shape[-1]))
+ for component, values in state.components.items()
+ }
+ return selected
def forward_h3_adapter(adapter: Any, **kwargs: Any) -> Any:
@@ -901,8 +955,10 @@ def _normalize_b1_integer(value: Any, field: str) -> int:
def _decoded_video_sample(video: Any) -> Any:
if isinstance(video, torch.Tensor) and video.shape[0] == 1:
return video[0]
- if isinstance(video, list) and len(video) == 1 and (
- video[0] is None or isinstance(video[0], list)
+ if (
+ isinstance(video, list)
+ and len(video) == 1
+ and (video[0] is None or isinstance(video[0], list))
):
return video[0]
return video
diff --git a/src/flow_factory/models/pipeline_contracts.py b/src/flow_factory/models/pipeline_contracts.py
index f8fe1cce9..3903811d4 100644
--- a/src/flow_factory/models/pipeline_contracts.py
+++ b/src/flow_factory/models/pipeline_contracts.py
@@ -16,7 +16,7 @@
from __future__ import annotations
-from typing import Optional
+from typing import Optional, Tuple
from ..contracts import (
BatchCapability,
@@ -158,11 +158,67 @@ def video_output_contract(
)
+def audio_video_output_contract(
+ *,
+ negative_prompt: NegativePromptPolicy,
+ input_rules: Tuple[InputMediaRule, ...] = (),
+ input_binding: InputMediaBinding = InputMediaBinding.GROUPED_BY_TYPE,
+ input_order: InputMediaOrder = InputMediaOrder.INSENSITIVE,
+ output_fps: RateRequirement = RateRequirement.REQUIRED,
+ output_sample_rate: RateRequirement = RateRequirement.REQUIRED,
+ geometry_source: GeometrySource = GeometrySource.OUTPUT_MEDIA,
+ batch_capability: BatchCapability = BatchCapability.SINGLE_SAMPLE,
+) -> PipelineIOContract:
+ """Build an exact ordered video-and-audio output declaration.
+
+ Supplying explicit input rules keeps this constructor neutral to how a model
+ binds conditions: prompt-only, grouped image conditions, and globally ordered
+ heterogeneous references all use the same output contract. Cross-type total
+ cardinality constraints remain adapter-owned because ``InputMediaSpec``
+ represents per-type bounds only.
+
+ Args:
+ negative_prompt: Whether negative prompts are unsupported, optional, or required.
+ input_rules: Canonically ordered per-type input-media rules.
+ input_binding: How input media are projected into model-facing arguments.
+ input_order: Which input-media ordering carries semantic meaning.
+ output_fps: Whether target-video frame-rate metadata is accepted or required.
+ output_sample_rate: Whether target-audio sample-rate metadata is accepted or required.
+ geometry_source: Boundary that determines aligned output geometry.
+ batch_capability: Whether the adapter accepts uniform, ragged, or single-sample batches.
+
+ Returns:
+ Immutable pipeline I/O contract with exact output order ``(video, audio)``.
+ """
+ video_format = MediaFormat(
+ type=MediaType.VIDEO,
+ fps=output_fps,
+ sample_rate=RateRequirement.NOT_APPLICABLE,
+ )
+ audio_format = MediaFormat(
+ type=MediaType.AUDIO,
+ fps=RateRequirement.NOT_APPLICABLE,
+ sample_rate=output_sample_rate,
+ )
+ return PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=input_rules,
+ binding=input_binding,
+ order=input_order,
+ ),
+ negative_prompt=negative_prompt,
+ output_media=OutputMediaSequence(items=(video_format, audio_format)),
+ geometry_source=geometry_source,
+ batch_capability=batch_capability,
+ )
+
+
__all__ = [
"AUDIO_FORMAT_REQUIRED_RATE",
"IMAGE_FORMAT",
"VIDEO_FORMAT_OPTIONAL_FPS",
"VIDEO_FORMAT_REQUIRED_FPS",
+ "audio_video_output_contract",
"image_output_contract",
"video_output_contract",
]
diff --git a/tests/data_utils/test_offline_dataset.py b/tests/data_utils/test_offline_dataset.py
index 093c9db38..23ce8b16f 100644
--- a/tests/data_utils/test_offline_dataset.py
+++ b/tests/data_utils/test_offline_dataset.py
@@ -38,6 +38,7 @@
PreferenceOutputBatch,
compute_offline_condition_id,
compute_offline_record_id,
+ decode_audio,
decode_video,
load_offline_manifest,
)
@@ -867,16 +868,28 @@ def test_collator_uses_declared_supervision_instead_of_first_item_union(tmp_path
OfflineCollator("demonstration")([demonstration_dataset[0], preference_dataset[0]])
-def test_unsupported_audio_output_fails_explicitly(tmp_path: Path) -> None:
+def test_default_audio_decoder_preserves_samples_and_source_clock_override(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ target_path = tmp_path / "target.wav"
+ target_path.write_bytes(b"identity-only audio payload")
+ calls: List[tuple[str, int | None]] = []
+
+ def fake_load_audio(path: str, sample_rate: int | None = None) -> torch.Tensor:
+ calls.append((str(path), sample_rate))
+ return torch.arange(12, dtype=torch.float64).reshape(2, 6).requires_grad_()
+
+ monkeypatch.setattr(offline_dataset_module, "load_audio", fake_load_audio)
media: Dict[str, Any] = {
"type": "audio",
- "path": "target.audio",
+ "path": "target.wav",
"sample_rate": 16000,
}
record = normalize_v2_record(
{
"schema_version": 2,
- "input": {"prompt": "unsupported", "media": []},
+ "input": {"prompt": "audio target", "media": []},
"supervision": {
"type": "demonstration",
"target": {"media": [media]},
@@ -884,14 +897,27 @@ def test_unsupported_audio_output_fails_explicitly(tmp_path: Path) -> None:
},
dataset_dir=tmp_path,
)
- with pytest.raises(NotImplementedError, match=r"type 'audio'.*no decoder"):
- OfflineDataset(
- [record],
- _condition_cache([record], [{"prompt_embeds": torch.ones(2)}]),
- source_name=SOURCE_NAME,
- source_id=SOURCE_ID,
- supervision_type="demonstration",
- )
+ dataset = OfflineDataset(
+ [record],
+ _condition_cache([record], [{"prompt_embeds": torch.ones(2)}]),
+ source_name=SOURCE_NAME,
+ source_id=SOURCE_ID,
+ supervision_type="demonstration",
+ )
+
+ item = dataset[0]
+
+ assert isinstance(item.output, DemonstrationOutput)
+ decoded = item.output.target_media[0]
+ assert calls == [(str(target_path), None)]
+ assert decoded.payload.shape == (2, 6)
+ assert decoded.payload.dtype is torch.float32
+ assert decoded.payload.device.type == "cpu"
+ assert decoded.payload.is_contiguous()
+ assert decoded.payload.requires_grad is False
+ assert decoded.sample_rate == 16000
+ assert pickle.loads(pickle.dumps(decode_audio)) is decode_audio
+ pickle.dumps(dataset)
def test_default_video_decoder_returns_diffusers_compatible_cpu_frames(tmp_path: Path) -> None:
diff --git a/tests/data_utils/test_offline_train_data.py b/tests/data_utils/test_offline_train_data.py
index 7c9a3d960..42982bdea 100644
--- a/tests/data_utils/test_offline_train_data.py
+++ b/tests/data_utils/test_offline_train_data.py
@@ -142,6 +142,7 @@ def prepare(self, *args: Any) -> None:
class _CountingPreprocessor:
def __init__(self) -> None:
self.calls = 0
+ self.batch_sizes: List[int] = []
self.is_train: bool | None = None
self.guidance_scale: float | None = None
@@ -153,6 +154,7 @@ def preprocess(
guidance_scale: float,
) -> Dict[str, torch.Tensor]:
self.calls += 1
+ self.batch_sizes.append(len(prompt))
self.is_train = is_train
self.guidance_scale = guidance_scale
return {"prompt_embeds": torch.ones(len(prompt), 2)}
@@ -583,6 +585,37 @@ def test_builder_rejects_single_sample_pipeline_batching_before_dataset_io(
assert preprocessor.calls == 0
+def test_single_sample_contract_forces_condition_preprocessing_batch_size_one(
+ tmp_path: Path,
+) -> None:
+ """Single-sample adapters must also preprocess their condition cache at B=1."""
+ dataset_dir = tmp_path / "single-sample-cache"
+ dataset_dir.mkdir()
+ rows = []
+ for index in range(2):
+ target = f"target-{index}.png"
+ _write_image(dataset_dir / target, index)
+ rows.append(_demonstration_row(f"prompt-{index}", target))
+ _write_manifest(dataset_dir, rows)
+ contract = replace(
+ _TEXT_TO_IMAGE_CONTRACT,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
+ preprocessor = _CountingPreprocessor()
+
+ build_offline_train_dataloader(
+ _config(tmp_path, [_source("single-sample-cache", dataset_dir, 0)]),
+ _Accelerator(),
+ preprocessor.preprocess,
+ supervision_type="demonstration",
+ pipeline_io_contract=contract,
+ shuffle=False,
+ )
+
+ assert preprocessor.calls == 2
+ assert preprocessor.batch_sizes == [1, 1]
+
+
def test_builder_rejects_non_unit_weight_and_unresolved_source_id_before_io(
tmp_path: Path,
) -> None:
@@ -610,7 +643,17 @@ def test_builder_rejects_non_unit_weight_and_unresolved_source_id_before_io(
def test_builder_rejects_missing_target_decoder_before_condition_preprocessing(
tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
) -> None:
+ monkeypatch.setattr(
+ offline_train_data,
+ "DEFAULT_MEDIA_DECODERS",
+ {
+ media_type: decoder
+ for media_type, decoder in offline_dataset_module.DEFAULT_MEDIA_DECODERS.items()
+ if media_type != "audio"
+ },
+ )
dataset_dir = tmp_path / "audio"
_write_manifest(
dataset_dir,
diff --git a/tests/models/minimax_h3/test_output_codec.py b/tests/models/minimax_h3/test_output_codec.py
new file mode 100644
index 000000000..fceacbff9
--- /dev/null
+++ b/tests/models/minimax_h3/test_output_codec.py
@@ -0,0 +1,604 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for MiniMax H3 T2VA offline audiovisual target encoding."""
+
+from types import SimpleNamespace
+from typing import Any, Optional
+
+import numpy as np
+import pytest
+import torch
+
+from flow_factory.contracts import (
+ BatchCapability,
+ GeometrySource,
+ MediaType,
+ NegativePromptPolicy,
+ RateRequirement,
+)
+from flow_factory.data_utils.offline_dataset import DecodedMedia
+from flow_factory.models.minimax_h3._output import (
+ MiniMaxH3AVOutputCodec,
+ prepare_h3_target_audio,
+ prepare_h3_target_video,
+ resolve_h3_output_geometry,
+ validate_h3_encoded_output_geometry,
+)
+from flow_factory.models.minimax_h3.adapters import (
+ MiniMaxH3FL2VAAdapter,
+ MiniMaxH3Ref2VAAdapter,
+ MiniMaxH3T2VAAdapter,
+)
+from flow_factory.samples import ComponentTimes, LatentState
+from flow_factory.scheduler import MiniMaxH3SDEScheduler
+from flow_factory.trainers.common.flow_matching import (
+ build_noised_output_state,
+ flow_matching_per_sample_loss,
+ validate_preference_component_times,
+ validate_preference_output_states,
+)
+
+
+class _Posterior:
+ def __init__(self, values: torch.Tensor) -> None:
+ self.values = values
+ self.sample_generator: Optional[torch.Generator] = None
+ self.sample_calls = 0
+ self.mode_calls = 0
+
+ def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ self.sample_generator = generator
+ self.sample_calls += 1
+ return self.values
+
+ def mode(self) -> torch.Tensor:
+ self.mode_calls += 1
+ return self.values
+
+
+class _VideoVAE:
+ def __init__(self, posterior: _Posterior) -> None:
+ self.posterior = posterior
+ self.config = SimpleNamespace(
+ latents_mean=[0.0] * 24,
+ latents_std=[1.0] * 24,
+ )
+ self.encoded_pixels: Optional[torch.Tensor] = None
+
+ def encode(self, pixels: torch.Tensor) -> Any:
+ self.encoded_pixels = pixels
+ return SimpleNamespace(latent_dist=self.posterior)
+
+
+class _AudioVAE:
+ hop_length = 800
+
+ def __init__(self, posterior: _Posterior) -> None:
+ self.posterior = posterior
+ self.config = SimpleNamespace(
+ sampling_rate=32000,
+ latents_mean=[0.0] * 32,
+ latents_std=[1.0] * 32,
+ )
+ self.encoded_waveform: Optional[torch.Tensor] = None
+
+ def encode(self, waveform: torch.Tensor) -> Any:
+ self.encoded_waveform = waveform
+ return SimpleNamespace(latent_dist=self.posterior)
+
+
+class _TinyH3Transformer(torch.nn.Module):
+ """Gradient-bearing stand-in that preserves H3's two-output contract."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.weight = torch.nn.Parameter(torch.tensor(0.25))
+ self.keyframe_noise_aug = 0.999
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ audio_hidden_states: torch.Tensor,
+ **kwargs: Any,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ return hidden_states * self.weight, audio_hidden_states * self.weight
+
+
+class _Adapter:
+ device = torch.device("cpu")
+
+ def __init__(self) -> None:
+ self.video_posterior = _Posterior(torch.full((1, 24, 7, 2, 2), 2.0005))
+ self.audio_posterior = _Posterior(torch.full((2, 32, 37), 3.0))
+ self.vae = _VideoVAE(self.video_posterior)
+ self.audio_vae = _AudioVAE(self.audio_posterior)
+ self.transformer = _TinyH3Transformer()
+ self.pipeline = SimpleNamespace(
+ fps=24,
+ pixel_mean=(0.485, 0.456, 0.406),
+ pixel_std=(0.229, 0.224, 0.225),
+ vae_spatial_compression_ratio=16,
+ vae_frames_per_chunk=17,
+ vae_latents_per_chunk=5,
+ vae_latent_channels=24,
+ patch_size=(1, 2, 2),
+ audio_channels=2,
+ audio_latent_channels=32,
+ audio_sampling_rate=32000,
+ min_duration=0.5,
+ max_duration=15.0,
+ )
+ self.training_args = SimpleNamespace(
+ height=32,
+ width=32,
+ num_frames=22,
+ frame_rate=24.0,
+ )
+
+ def get_component(self, name: str) -> Any:
+ return getattr(self, name)
+
+
+def _base_adapter(*, latent_storage_dtype: str) -> MiniMaxH3T2VAAdapter:
+ components = _Adapter()
+ adapter = object.__new__(MiniMaxH3T2VAAdapter)
+ adapter.accelerator = SimpleNamespace(device=torch.device("cpu"))
+ adapter.training_args = SimpleNamespace(
+ height=32,
+ width=32,
+ num_frames=22,
+ frame_rate=24.0,
+ latent_storage_dtype=latent_storage_dtype,
+ )
+ adapter.pipeline = components.pipeline
+ adapter.component_runtime = SimpleNamespace(
+ get_component=lambda name: components.get_component(name)
+ )
+ adapter.scheduler = MiniMaxH3SDEScheduler(
+ shift=12.0,
+ dynamics_type="Flow-SDE",
+ sde_steps=[0, 1],
+ num_sde_steps=2,
+ )
+ adapter.audio_scheduler = MiniMaxH3SDEScheduler(
+ shift=3.0,
+ dynamics_type="Flow-SDE",
+ sde_steps=[0, 1],
+ num_sde_steps=2,
+ )
+ adapter.scheduler.set_timesteps(2, device="cpu")
+ adapter.audio_scheduler.set_timesteps(2, device="cpu")
+ adapter._output_state_codec = adapter.build_output_state_codec()
+ return adapter
+
+
+def _condition() -> dict[str, Any]:
+ text_rows = 2
+ audio_rows = 74
+ video_rows = 7
+ sequence_length = text_rows + audio_rows + video_rows
+ return {
+ "height": [32],
+ "width": [32],
+ "num_frames": [22],
+ "num_latent_frames": [7],
+ "latent_height": [2],
+ "latent_width": [2],
+ "num_audio_latents": [37],
+ "position_ids": torch.zeros(1, sequence_length, 3, dtype=torch.float64),
+ "token_tags": torch.zeros(1, sequence_length, dtype=torch.long),
+ "text_indices": torch.arange(text_rows).unsqueeze(0),
+ "audio_indices": torch.arange(text_rows, text_rows + audio_rows).unsqueeze(0),
+ "video_indices": torch.arange(text_rows + audio_rows, sequence_length).unsqueeze(0),
+ "num_condition_video_rows": [0],
+ "num_condition_audio_rows": [0],
+ }
+
+
+def _media(*, audio_samples: int = 37 * 800) -> tuple[tuple[DecodedMedia, ...], ...]:
+ return (
+ (
+ DecodedMedia(
+ type="video",
+ path="target.mp4",
+ payload=np.zeros((22, 32, 32, 3), dtype=np.uint8),
+ fps=24.0,
+ ),
+ DecodedMedia(
+ type="audio",
+ path="target.wav",
+ payload=torch.linspace(-1.0, 1.0, audio_samples).unsqueeze(0),
+ sample_rate=32000,
+ ),
+ ),
+ )
+
+
+def test_t2va_declares_exact_configured_single_sample_av_contract() -> None:
+ contract = MiniMaxH3T2VAAdapter.pipeline_io_contract
+
+ assert contract is not None
+ assert contract.input_media.rules == ()
+ assert contract.negative_prompt is NegativePromptPolicy.UNSUPPORTED
+ assert contract.geometry_source is GeometrySource.CONFIGURED
+ assert contract.batch_capability is BatchCapability.SINGLE_SAMPLE
+ assert tuple(item.type for item in contract.output_media.items) == (
+ MediaType.VIDEO,
+ MediaType.AUDIO,
+ )
+ assert contract.output_media.items[0].fps is RateRequirement.REQUIRED
+ assert contract.output_media.items[1].sample_rate is RateRequirement.REQUIRED
+ MiniMaxH3T2VAAdapter.validate_offline_output_capability()
+
+
+@pytest.mark.parametrize("adapter_type", [MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter])
+def test_conditioned_h3_workflows_require_reproducible_prefix_binder(
+ adapter_type: type,
+) -> None:
+ with pytest.raises(NotImplementedError, match="conditioned-prefix binder"):
+ adapter_type.validate_offline_output_capability()
+
+
+def test_codec_uses_framework_video_sample_without_condition_fp16_rounding() -> None:
+ adapter = _Adapter()
+ generator = torch.Generator().manual_seed(17)
+
+ encoded = MiniMaxH3AVOutputCodec(adapter).encode_output_state(
+ _media(audio_samples=100),
+ _condition(),
+ generator,
+ )
+
+ assert adapter.video_posterior.sample_calls == 1
+ assert adapter.video_posterior.sample_generator is generator
+ assert adapter.video_posterior.mode_calls == 0
+ assert adapter.audio_posterior.mode_calls == 1
+ assert adapter.audio_posterior.sample_calls == 0
+ assert adapter.vae.encoded_pixels is not None
+ assert adapter.vae.encoded_pixels.shape == (1, 3, 22, 32, 32)
+ assert adapter.audio_vae.encoded_waveform is not None
+ assert adapter.audio_vae.encoded_waveform.shape == (2, 1, 37 * 800)
+ assert encoded.clean_state.component_names == ("video", "audio")
+ assert encoded.clean_state.components["video"].shape == (1, 7, 96)
+ assert encoded.clean_state.components["audio"].shape == (1, 74, 32)
+ sampled_value = adapter.video_posterior.values.flatten()[0]
+ condition_rounded_value = sampled_value.to(torch.float16).to(torch.float32)
+ assert sampled_value != condition_rounded_value
+ assert encoded.clean_state.components["video"].flatten()[0] == sampled_value
+ assert encoded.forward_context == {}
+ assert encoded.decode_context["geometry"] == {
+ "height": 32,
+ "width": 32,
+ "num_frames": 22,
+ "num_latent_frames": 7,
+ "latent_height": 2,
+ "latent_width": 2,
+ "num_audio_latents": 37,
+ }
+ signature = encoded.geometry_signatures[0]
+ assert signature.media[0].fps == 24.0
+ assert signature.media[1].samples == 37 * 800
+ assert signature.media[1].sample_rate == 32000
+ validate_h3_encoded_output_geometry(adapter, _media(), _condition(), encoded)
+
+
+def test_geometry_hook_independently_enforces_pipeline_duration_bounds() -> None:
+ adapter = _Adapter()
+ adapter.pipeline.min_duration = 5.0
+
+ with pytest.raises(ValueError, match="outside the pipeline contract"):
+ resolve_h3_output_geometry(adapter, _condition())
+
+
+@pytest.mark.parametrize(
+ ("field", "value", "message"),
+ [
+ ("height", 64, "cached output canvas"),
+ ("num_frames", 23, "cached output frame count"),
+ ],
+)
+def test_geometry_hook_rejects_cache_from_another_training_geometry(
+ field: str,
+ value: int,
+ message: str,
+) -> None:
+ adapter = _Adapter()
+ setattr(adapter.training_args, field, value)
+
+ with pytest.raises(ValueError, match=message):
+ resolve_h3_output_geometry(adapter, _condition())
+
+
+def test_geometry_hook_rejects_invalid_input_owned_flat_layout() -> None:
+ adapter = _Adapter()
+ condition = _condition()
+ encoded = MiniMaxH3AVOutputCodec(adapter).encode_output_state(_media(), condition)
+
+ condition["audio_indices"] = condition["audio_indices"][:, :-1]
+ with pytest.raises(ValueError, match="audio layout expected 74 target rows"):
+ validate_h3_encoded_output_geometry(adapter, _media(), condition, encoded)
+
+
+def test_base_lifecycle_casts_clean_state_without_output_context_dtype_drift() -> None:
+ adapter = _base_adapter(latent_storage_dtype="fp16")
+
+ encoded = adapter.encode_output_state(_media(), _condition())
+
+ assert encoded.clean_state.components["video"].dtype is torch.float16
+ assert encoded.clean_state.components["audio"].dtype is torch.float16
+ assert encoded.forward_context == {}
+
+
+def test_adapter_decode_routes_both_components_with_cached_geometry(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ encoded = MiniMaxH3AVOutputCodec(_Adapter()).encode_output_state(_media(), _condition())
+ adapter = object.__new__(MiniMaxH3T2VAAdapter)
+ observed: dict[str, Any] = {}
+
+ def decode(self: Any, state: LatentState, **kwargs: Any) -> str:
+ observed["state"] = state
+ observed.update(kwargs)
+ return "decoded-av"
+
+ monkeypatch.setattr(MiniMaxH3T2VAAdapter, "decode_latents", decode)
+
+ assert adapter.decode_output_state(encoded, output_type="np") == "decoded-av"
+ assert observed["state"] is encoded.clean_state
+ assert observed["geometry"] == encoded.decode_context["geometry"]
+ assert observed["output_type"] == "np"
+
+
+def test_flat_offline_condition_binds_layout_and_state_native_empty_prefixes(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter = object.__new__(MiniMaxH3T2VAAdapter)
+ state = LatentState(
+ {
+ "video": torch.ones(1, 7, 96, dtype=torch.float16),
+ "audio": torch.ones(1, 74, 32, dtype=torch.float16),
+ }
+ )
+ times = ComponentTimes(
+ timestep={"video": torch.tensor([500.0]), "audio": torch.tensor([500.0])},
+ next_timestep={"video": torch.tensor([0.0]), "audio": torch.tensor([0.0])},
+ )
+ observed: dict[str, Any] = {}
+
+ def record_forward(self: Any, **kwargs: Any) -> str:
+ observed.update(kwargs)
+ return "bound"
+
+ monkeypatch.setattr(MiniMaxH3T2VAAdapter, "forward", record_forward)
+ result = adapter._forward_state(
+ batch=SimpleNamespace(),
+ state=state,
+ times=times,
+ next_state=None,
+ compute_log_prob=False,
+ return_fields=("velocity",),
+ noise_level=0.0,
+ forward_kwargs={**_condition(), "prompt_embeds": torch.zeros(1, 2, 4)},
+ )
+
+ assert result == "bound"
+ assert torch.equal(observed["layout"]["video_indices"], _condition()["video_indices"][0])
+ for component in ("video", "audio"):
+ prefix = observed["condition_prefixes"][component]
+ clean = state.components[component]
+ assert prefix.shape == (1, 0, clean.shape[-1])
+ assert prefix.dtype is clean.dtype
+ assert prefix.device == clean.device
+
+
+def test_h3_sft_path_runs_codec_cast_noising_and_velocity_only_forward(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter = _base_adapter(latent_storage_dtype="fp32")
+ condition = {**_condition(), "prompt_embeds": torch.zeros(1, 2, 4)}
+ encoded = adapter.encode_output_state(_media(), condition)
+ times, noised = build_noised_output_state(
+ adapter,
+ encoded.clean_state,
+ torch.tensor([500.0]),
+ batch=condition,
+ generator=torch.Generator().manual_seed(31),
+ )
+
+ def unexpected_scheduler_step(*args: Any, **kwargs: Any) -> None:
+ raise AssertionError("offline velocity-only forward must not step schedulers")
+
+ monkeypatch.setattr(adapter.scheduler, "step", unexpected_scheduler_step)
+ monkeypatch.setattr(adapter.audio_scheduler, "step", unexpected_scheduler_step)
+
+ output = adapter._forward_state(
+ batch=SimpleNamespace(),
+ state=noised.state,
+ times=times,
+ next_state=None,
+ compute_log_prob=False,
+ return_fields=("velocity",),
+ noise_level=0.0,
+ forward_kwargs=condition,
+ )
+ loss = flow_matching_per_sample_loss(adapter, output.velocity, noised).mean()
+ loss.backward()
+
+ transformer = adapter.get_component("transformer")
+ assert output.velocity.component_names == ("video", "audio")
+ assert output.velocity.components["video"].shape == (1, 7, 96)
+ assert output.velocity.components["audio"].shape == (1, 74, 32)
+ assert transformer.weight.grad is not None
+
+
+def test_h3_offline_dpo_arms_reuse_structured_noise_through_real_forward() -> None:
+ adapter = _base_adapter(latent_storage_dtype="fp32")
+ condition = {**_condition(), "prompt_embeds": torch.zeros(1, 2, 4)}
+ chosen = adapter.encode_output_state(_media(), condition)
+ rejected = adapter.encode_output_state(_media(), condition)
+ validate_preference_output_states(chosen, rejected)
+
+ primary_timesteps = torch.tensor([650.0])
+ chosen_times, chosen_noised = build_noised_output_state(
+ adapter,
+ chosen.clean_state,
+ primary_timesteps,
+ batch=condition,
+ generator=torch.Generator().manual_seed(41),
+ )
+ rejected_times, rejected_noised = build_noised_output_state(
+ adapter,
+ rejected.clean_state,
+ primary_timesteps,
+ batch=condition,
+ noise=chosen_noised.noise,
+ )
+ validate_preference_component_times(chosen_times, rejected_times)
+
+ outputs = [
+ adapter._forward_state(
+ batch=SimpleNamespace(),
+ state=noised.state,
+ times=times,
+ next_state=None,
+ compute_log_prob=False,
+ return_fields=("velocity",),
+ noise_level=0.0,
+ forward_kwargs=condition,
+ )
+ for times, noised in (
+ (chosen_times, chosen_noised),
+ (rejected_times, rejected_noised),
+ )
+ ]
+
+ assert rejected_noised.noise is chosen_noised.noise
+ for component in ("video", "audio"):
+ assert torch.equal(
+ rejected_noised.noise.components[component],
+ chosen_noised.noise.components[component],
+ )
+ assert outputs[0].velocity.components[component].shape == (
+ chosen_noised.state.components[component].shape
+ )
+ assert outputs[1].velocity.components[component].shape == (
+ rejected_noised.state.components[component].shape
+ )
+
+
+def test_conditioned_offline_replay_still_requires_prefix_binder() -> None:
+ adapter = object.__new__(MiniMaxH3FL2VAAdapter)
+ state = LatentState(
+ {
+ "video": torch.ones(1, 7, 96),
+ "audio": torch.ones(1, 74, 32),
+ }
+ )
+ times = ComponentTimes(
+ timestep={"video": torch.tensor([500.0]), "audio": torch.tensor([500.0])},
+ next_timestep={"video": torch.tensor([0.0]), "audio": torch.tensor([0.0])},
+ )
+
+ with pytest.raises(ValueError, match="conditioned-prefix binder"):
+ adapter._forward_state(
+ batch=SimpleNamespace(),
+ state=state,
+ times=times,
+ next_state=None,
+ compute_log_prob=False,
+ return_fields=("velocity",),
+ noise_level=0.0,
+ forward_kwargs={**_condition(), "prompt_embeds": torch.zeros(1, 2, 4)},
+ )
+
+
+def test_audio_alignment_trims_or_zero_pads_on_the_exact_latent_grid() -> None:
+ short = prepare_h3_target_audio(
+ torch.tensor([[1.0, 2.0, 3.0]]),
+ source_sample_rate=32000,
+ target_sample_rate=32000,
+ target_samples=5,
+ target_duration_seconds=5 / 32000,
+ )
+ assert short.shape == (2, 5)
+ assert torch.equal(short[0], torch.tensor([1.0, 2.0, 3.0, 0.0, 0.0]))
+ assert torch.equal(short[0], short[1])
+
+ long = prepare_h3_target_audio(
+ torch.arange(14, dtype=torch.float32).reshape(2, 7),
+ source_sample_rate=32000,
+ target_sample_rate=32000,
+ target_samples=5,
+ target_duration_seconds=5 / 32000,
+ )
+ assert torch.equal(long, torch.tensor([[0, 1, 2, 3, 4], [7, 8, 9, 10, 11]]).float())
+
+
+def test_audio_alignment_truncates_on_source_clock_before_resampling(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ observed: dict[str, Any] = {}
+
+ def convert(
+ waveform: torch.Tensor,
+ from_rate: int,
+ to_rate: int,
+ to_channels: int,
+ ) -> torch.Tensor:
+ observed["waveform"] = waveform.clone()
+ observed["rates"] = (from_rate, to_rate, to_channels)
+ return waveform.repeat_interleave(2, dim=-1).expand(2, -1).contiguous()
+
+ monkeypatch.setattr("flow_factory.models.minimax_h3._output.convert_audio", convert)
+ result = prepare_h3_target_audio(
+ torch.arange(8, dtype=torch.float32).unsqueeze(0),
+ source_sample_rate=10,
+ target_sample_rate=20,
+ target_samples=10,
+ target_duration_seconds=0.5,
+ )
+
+ assert torch.equal(observed["waveform"], torch.arange(5, dtype=torch.float32).unsqueeze(0))
+ assert observed["rates"] == (10, 20, 2)
+ assert result.shape == (2, 10)
+
+
+def test_target_video_fails_when_source_cannot_cover_configured_timeline() -> None:
+ with pytest.raises(ValueError, match="too short"):
+ prepare_h3_target_video(
+ np.zeros((4, 8, 8, 3), dtype=np.uint8),
+ source_fps=24.0,
+ target_frames=5,
+ target_fps=24.0,
+ height=8,
+ width=8,
+ )
+
+
+def test_target_video_matches_official_h3_fps_filter_for_30_to_24() -> None:
+ frames = np.zeros((10, 1, 1, 3), dtype=np.uint8)
+ frames[:, 0, 0, :] = np.arange(10, dtype=np.uint8)[:, None]
+
+ pixels = prepare_h3_target_video(
+ frames,
+ source_fps=30.0,
+ target_frames=8,
+ target_fps=24.0,
+ height=1,
+ width=1,
+ )
+
+ values = pixels[0, 0, :, 0, 0].mul(255).round().to(torch.uint8)
+ assert torch.equal(values, torch.tensor([0, 1, 3, 4, 5, 6, 8, 9], dtype=torch.uint8))
diff --git a/tests/models/minimax_h3/test_review_fixes.py b/tests/models/minimax_h3/test_review_fixes.py
index 117b76c5f..91ffd9d76 100644
--- a/tests/models/minimax_h3/test_review_fixes.py
+++ b/tests/models/minimax_h3/test_review_fixes.py
@@ -656,6 +656,8 @@ def _training_batch() -> StackedSampleBatch:
def _patch_h3_forward(monkeypatch) -> None:
def forward(*args, **kwargs):
state = args[1]
+ if kwargs.get("velocity_only"):
+ return state
return MultiModalStepOutput(
next_state=state,
next_state_mean=state,
diff --git a/tests/models/test_offline_output_capability_matrix.py b/tests/models/test_offline_output_capability_matrix.py
index 00afa9fc4..6305f3812 100644
--- a/tests/models/test_offline_output_capability_matrix.py
+++ b/tests/models/test_offline_output_capability_matrix.py
@@ -30,9 +30,8 @@
(Wan2_I2V_Adapter, "first-frame VAE condition"),
(LTX2_T2AV_Adapter, "paired video/audio decoding"),
(LTX2_I2AV_Adapter, "active mask"),
- (MiniMaxH3T2VAAdapter, "target-video posterior policy"),
- (MiniMaxH3FL2VAAdapter, "target-video posterior policy"),
- (MiniMaxH3Ref2VAAdapter, "target-video posterior policy"),
+ (MiniMaxH3FL2VAAdapter, "conditioned-prefix binder"),
+ (MiniMaxH3Ref2VAAdapter, "conditioned-prefix binder"),
],
)
def test_unimplemented_offline_media_semantics_fail_before_model_loading(
@@ -42,3 +41,8 @@ def test_unimplemented_offline_media_semantics_fail_before_model_loading(
"""Expose actionable blockers instead of silently guessing target encoding."""
with pytest.raises(NotImplementedError, match=reason_fragment):
adapter_type.validate_offline_output_capability()
+
+
+def test_minimax_h3_t2va_declares_complete_offline_output_semantics() -> None:
+ """T2VA has no conditioned prefix and can encode paired AV targets on demand."""
+ MiniMaxH3T2VAAdapter.validate_offline_output_capability()
diff --git a/tests/models/test_pipeline_contract_constructors.py b/tests/models/test_pipeline_contract_constructors.py
new file mode 100644
index 000000000..03cb140f2
--- /dev/null
+++ b/tests/models/test_pipeline_contract_constructors.py
@@ -0,0 +1,92 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Test reusable adapter pipeline-contract constructors."""
+
+from flow_factory.contracts import (
+ BatchCapability,
+ GeometrySource,
+ InputMediaBinding,
+ InputMediaOrder,
+ InputMediaRule,
+ MediaType,
+ NegativePromptPolicy,
+ RateRequirement,
+)
+from flow_factory.models.pipeline_contracts import (
+ AUDIO_FORMAT_REQUIRED_RATE,
+ IMAGE_FORMAT,
+ VIDEO_FORMAT_OPTIONAL_FPS,
+ audio_video_output_contract,
+)
+
+
+def test_audio_video_output_contract_defaults_to_exact_required_av_sequence() -> None:
+ contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ )
+
+ assert contract.input_media.rules == ()
+ assert contract.input_media.binding is InputMediaBinding.GROUPED_BY_TYPE
+ assert contract.input_media.order is InputMediaOrder.INSENSITIVE
+ assert tuple(item.type for item in contract.output_media.items) == (
+ MediaType.VIDEO,
+ MediaType.AUDIO,
+ )
+ assert contract.output_media.items[0].fps is RateRequirement.REQUIRED
+ assert contract.output_media.items[1].sample_rate is RateRequirement.REQUIRED
+ assert contract.geometry_source is GeometrySource.OUTPUT_MEDIA
+ assert contract.batch_capability is BatchCapability.SINGLE_SAMPLE
+
+
+def test_audio_video_output_contract_preserves_grouped_image_bounds() -> None:
+ contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ input_rules=(InputMediaRule(format=IMAGE_FORMAT, min_count=1, max_count=2),),
+ input_order=InputMediaOrder.WITHIN_TYPE,
+ geometry_source=GeometrySource.CONFIGURED,
+ )
+
+ assert contract.input_media.rules[0].min_count == 1
+ assert contract.input_media.rules[0].max_count == 2
+ assert contract.input_media.binding is InputMediaBinding.GROUPED_BY_TYPE
+ assert contract.input_media.order is InputMediaOrder.WITHIN_TYPE
+ assert contract.geometry_source is GeometrySource.CONFIGURED
+
+
+def test_audio_video_output_contract_preserves_ordered_reference_rules() -> None:
+ contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ input_rules=(
+ InputMediaRule(format=IMAGE_FORMAT, min_count=0, max_count=12),
+ InputMediaRule(format=VIDEO_FORMAT_OPTIONAL_FPS, min_count=0, max_count=12),
+ InputMediaRule(format=AUDIO_FORMAT_REQUIRED_RATE, min_count=0, max_count=12),
+ ),
+ input_binding=InputMediaBinding.ORDERED_REFERENCES,
+ input_order=InputMediaOrder.GLOBAL,
+ output_fps=RateRequirement.OPTIONAL,
+ output_sample_rate=RateRequirement.OPTIONAL,
+ batch_capability=BatchCapability.RAGGED,
+ )
+
+ assert tuple(rule.format.type for rule in contract.input_media.rules) == (
+ MediaType.IMAGE,
+ MediaType.VIDEO,
+ MediaType.AUDIO,
+ )
+ assert contract.input_media.binding is InputMediaBinding.ORDERED_REFERENCES
+ assert contract.input_media.order is InputMediaOrder.GLOBAL
+ assert contract.output_media.items[0].fps is RateRequirement.OPTIONAL
+ assert contract.output_media.items[1].sample_rate is RateRequirement.OPTIONAL
+ assert contract.batch_capability is BatchCapability.RAGGED
From 8d0ca84d360f815bce8f8eb8c0ba93d38c693514 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Fri, 28 Aug 2026 22:58:58 +0800
Subject: [PATCH 25/76] chore: format offline branch changes
---
src/flow_factory/data_utils/dataset.py | 24 ++++++------------------
src/flow_factory/hparams/model_args.py | 4 +---
src/flow_factory/models/abc.py | 16 ++++------------
tests/docs/test_sensenova_docs.py | 5 +----
4 files changed, 12 insertions(+), 37 deletions(-)
diff --git a/src/flow_factory/data_utils/dataset.py b/src/flow_factory/data_utils/dataset.py
index 97bb2ff96..2a145cfd4 100644
--- a/src/flow_factory/data_utils/dataset.py
+++ b/src/flow_factory/data_utils/dataset.py
@@ -765,11 +765,7 @@ def _preprocess_batch(
# Complex values (nested lists/dicts) in the source JSONL must already be
# stored as JSON strings for Arrow compatibility.
batch_dict[METADATA_COLUMN] = [
- {
- k: v[idx]
- for k, v in batch.items()
- if k not in metadata_excluded_columns
- }
+ {k: v[idx] for k, v in batch.items() if k not in metadata_excluded_columns}
for idx in range(len(batch["prompt"]))
]
@@ -1357,8 +1353,7 @@ def _load_ordered_reference(
loaded["sample_rate"] = effective_sample_rate
else:
raise ValueError(
- "expected ordered reference kind in ('image', 'video', 'audio'), "
- f"got {kind!r}"
+ "expected ordered reference kind in ('image', 'video', 'audio'), " f"got {kind!r}"
)
except (FileNotFoundError, ImportError, OSError, RuntimeError, ValueError) as error:
raise ValueError(
@@ -1407,19 +1402,14 @@ def _decode_ordered_video(
if not container.streams.video:
raise ValueError(f"expected a video stream in {video_path!r}, got none")
video_stream = container.streams.video[0]
- frames = [
- frame.to_ndarray(format="rgb24")
- for frame in container.decode(video_stream)
- ]
+ frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(video_stream)]
reported_frame_rate = video_stream.average_rate or video_stream.guessed_rate
frame_rate = None if reported_frame_rate is None else float(reported_frame_rate)
audio = None
sample_rate = None
if container.streams.audio:
container.seek(0)
- audio, sample_rate = _decode_av_audio_stream(
- container, container.streams.audio[0]
- )
+ audio, sample_rate = _decode_av_audio_stream(container, container.streams.audio[0])
if not frames:
raise ValueError(f"expected video frames in {video_path!r}, decoded none")
return np.stack(frames), frame_rate, audio, sample_rate
@@ -1447,12 +1437,10 @@ def _decode_av_audio_stream(container: Any, stream: Any) -> tuple[torch.Tensor,
chunks = []
for frame in container.decode(stream):
chunks.extend(
- torch.from_numpy(resampled.to_ndarray())
- for resampled in resampler.resample(frame)
+ torch.from_numpy(resampled.to_ndarray()) for resampled in resampler.resample(frame)
)
chunks.extend(
- torch.from_numpy(resampled.to_ndarray())
- for resampled in resampler.resample(None)
+ torch.from_numpy(resampled.to_ndarray()) for resampled in resampler.resample(None)
)
if not chunks:
raise ValueError("expected decoded audio samples, got none")
diff --git a/src/flow_factory/hparams/model_args.py b/src/flow_factory/hparams/model_args.py
index 45e4e5da0..0c7f61921 100644
--- a/src/flow_factory/hparams/model_args.py
+++ b/src/flow_factory/hparams/model_args.py
@@ -95,9 +95,7 @@ def _serialize_dtype_policy(value: DTypePolicy) -> Any:
"""Serialize one normalized dtype policy for YAML output."""
if isinstance(value, dict):
return {
- selector: (
- None if configured_dtype is None else str(configured_dtype).split(".")[-1]
- )
+ selector: (None if configured_dtype is None else str(configured_dtype).split(".")[-1])
for selector, configured_dtype in value.items()
}
if value is not None:
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index 41ee508d1..604e7276f 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -16,11 +16,11 @@
import hashlib
import json
import logging
-import shutil
# src/flow_factory/models/abc.py
import os
import re
+import shutil
from abc import ABC, abstractmethod
from contextlib import ExitStack, contextmanager, nullcontext
from dataclasses import asdict, dataclass, field, fields
@@ -823,20 +823,12 @@ def _load_diffusers_pipeline(
user_policy=user_policy,
manifest_policy=manifest_policy,
component_names=component_names,
- transformer_names=[
- name for name in component_names if "transformer" in name
- ],
- text_encoder_names=[
- name for name in component_names if "text_encoder" in name
- ],
+ transformer_names=[name for name in component_names if "transformer" in name],
+ text_encoder_names=[name for name in component_names if "text_encoder" in name],
preserve_unselected=True,
)
kwargs.update(
- {
- key: value
- for key, value in load_dtype_kwargs.items()
- if key not in kwargs
- }
+ {key: value for key, value in load_dtype_kwargs.items() if key not in kwargs}
)
return pipeline_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
diff --git a/tests/docs/test_sensenova_docs.py b/tests/docs/test_sensenova_docs.py
index 60d6310b4..4e1bbc2b7 100644
--- a/tests/docs/test_sensenova_docs.py
+++ b/tests/docs/test_sensenova_docs.py
@@ -14,7 +14,6 @@
from pathlib import Path
-
ROOT = Path(__file__).resolve().parents[2]
@@ -25,9 +24,7 @@ def _text(path: str) -> str:
def test_main_readme_lists_sensenova_as_t2i_and_multi_reference_i2i() -> None:
readme = _text("README.md")
t2i_start = readme.index('| Text-to-Image | ')
- combined_start = readme.index(
- '
| Text-to-Image & Image(s)-to-Image | '
- )
+ combined_start = readme.index('
| Text-to-Image & Image(s)-to-Image | ')
video_start = readme.index('
| Text-to-Video | ')
t2i_only_rows = readme[t2i_start:combined_start]
From 49bb2f420436d4ac44b36b114d8c6778de497f4e Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Fri, 28 Aug 2026 23:03:55 +0800
Subject: [PATCH 26/76] fix(tests): isolate Bagel optional kernels
---
.agents/knowledge/topics/fix_patterns.md | 8 +++
src/flow_factory/trainers/offline/__init__.py | 14 ++++
tests/models/test_bagel_tdm_contracts.py | 70 +++++++++++++++----
3 files changed, 80 insertions(+), 12 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 9b6f6737b..669bc8d5d 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -198,6 +198,14 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Lesson**: An example configuration is executable documentation. Any recipe migration must update its production parse test, linked data provenance, optional dependency instructions, and validation claims in the same integration change.
- **Related Constraint**: #15
+### Optional-kernel adapter tests must lazy-load behind the dependency seam
+- **Date**: 2026-08-28
+- **Symptom**: Collecting the Bagel TDM contract test on macOS failed before any test ran because `flash-attn>=2.5.8` was unavailable.
+- **Root Cause**: The test imported the Bagel adapter at module scope instead of installing the existing fake optional-kernel modules before the adapter import.
+- **Fix**: `tests/models/test_bagel_tdm_contracts.py` now lazily imports Bagel after stubbing `flash_attn`, OpenCV, and the availability probes; new Python files also receive the required Apache 2.0 headers.
+- **Lesson**: Contract tests for CUDA-only optional adapters must exercise the adapter through its dependency boundary so CPU and macOS collection remains valid; importing such adapters at module scope turns an optional dependency into a repository-wide test dependency.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/trainers/offline/__init__.py b/src/flow_factory/trainers/offline/__init__.py
index 13c2d58e4..c47814e81 100644
--- a/src/flow_factory/trainers/offline/__init__.py
+++ b/src/flow_factory/trainers/offline/__init__.py
@@ -1,3 +1,17 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
"""Finite-dataset training algorithms."""
from .offline_dpo import OfflineDPOTrainer
diff --git a/tests/models/test_bagel_tdm_contracts.py b/tests/models/test_bagel_tdm_contracts.py
index 744084fcf..949359ea5 100644
--- a/tests/models/test_bagel_tdm_contracts.py
+++ b/tests/models/test_bagel_tdm_contracts.py
@@ -1,10 +1,30 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import importlib
+import importlib.machinery
+import sys
+import types
from types import MethodType
from typing import Any
import pytest
import torch
-from flow_factory.models.bagel.bagel import BagelAdapter, BagelSample
+import flow_factory.utils.imports as import_utils
from flow_factory.samples import BaseSample
from flow_factory.trainers.distillation.distillation_runtime import (
validate_media_free_rollout,
@@ -12,8 +32,24 @@
)
-def _adapter(decoder: Any) -> BagelAdapter:
- adapter = object.__new__(BagelAdapter)
+def _load_bagel_types(monkeypatch: pytest.MonkeyPatch) -> tuple[type, type]:
+ """Load Bagel behind the same optional-kernel seam as its adapter tests."""
+ flash_attn = types.ModuleType("flash_attn")
+ flash_attn.__spec__ = importlib.machinery.ModuleSpec("flash_attn", loader=None)
+ flash_attn.flash_attn_varlen_func = lambda *args, **kwargs: None
+ cv2 = types.ModuleType("cv2")
+ cv2.__spec__ = importlib.machinery.ModuleSpec("cv2", loader=None)
+ monkeypatch.setitem(sys.modules, "flash_attn", flash_attn)
+ monkeypatch.setitem(sys.modules, "cv2", cv2)
+ monkeypatch.setattr(import_utils, "is_flash_attn_available", lambda *args: True)
+ monkeypatch.setattr(import_utils, "get_flash_attn_version", lambda: "test")
+
+ module = importlib.import_module("flow_factory.models.bagel.bagel")
+ return module.BagelAdapter, module.BagelSample
+
+
+def _adapter(adapter_type: type, decoder: Any) -> Any:
+ adapter = object.__new__(adapter_type)
adapter.decode_latents = MethodType(decoder, adapter)
return adapter
@@ -33,18 +69,21 @@ def _result(batch_size: int = 2) -> dict[str, Any]:
}
-def test_bagel_assembles_samples_from_one_batched_decode() -> None:
+def test_bagel_assembles_samples_from_one_batched_decode(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter_type, sample_type = _load_bagel_types(monkeypatch)
calls: list[tuple[tuple[int, ...], tuple[int, int] | None]] = []
def decode(
- _self: BagelAdapter,
+ _self: Any,
latents: torch.Tensor,
image_shape: tuple[int, int] | None = None,
) -> list[torch.Tensor]:
calls.append((tuple(latents.shape), image_shape))
return [torch.full((3, 8, 8), index) for index in range(latents.shape[0])]
- adapter = _adapter(decode)
+ adapter = _adapter(adapter_type, decode)
samples = adapter._assemble_samples(
_result(),
prompts=["first", "second"],
@@ -54,21 +93,25 @@ def decode(
)
assert calls == [((2, 4, 8), (64, 64))]
- assert all(isinstance(sample, BagelSample) for sample in samples)
+ assert all(isinstance(sample, sample_type) for sample in samples)
assert torch.equal(samples[0].image, torch.zeros(3, 8, 8))
assert torch.equal(samples[1].image, torch.ones(3, 8, 8))
-def test_bagel_media_free_samples_keep_replay_trajectory() -> None:
+def test_bagel_media_free_samples_keep_replay_trajectory(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter_type, _ = _load_bagel_types(monkeypatch)
+
def decode(
- _self: BagelAdapter,
+ _self: Any,
latents: torch.Tensor,
image_shape: tuple[int, int] | None = None,
) -> list[torch.Tensor]:
del latents, image_shape
raise AssertionError("TDM must not invoke the real Bagel decoder")
- adapter = _adapter(decode)
+ adapter = _adapter(adapter_type, decode)
validate_media_free_rollout(adapter, algorithm_name="TDM")
with without_media_decoding(adapter, algorithm_name="TDM"):
@@ -90,7 +133,10 @@ def decode(
adapter.decode_latents(torch.zeros(2, 4, 8), image_shape=(64, 64))
-def test_bagel_maps_reference_guidance_to_text_cfg() -> None:
- adapter = object.__new__(BagelAdapter)
+def test_bagel_maps_reference_guidance_to_text_cfg(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter_type, _ = _load_bagel_types(monkeypatch)
+ adapter = object.__new__(adapter_type)
assert adapter.reference_guidance_kwargs(4.0) == {"cfg_text_scale": 4.0}
From 1622f5ef2aa350bbe797319cd4c40e95de47ce2d Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sat, 29 Aug 2026 11:12:15 +0800
Subject: [PATCH 27/76] fix(checkpoint): align distillation resume cursor with
rollout GAS
---
.agents/knowledge/topics/fix_patterns.md | 8 +++++
.../distillation/distillation_runtime.py | 31 ++++++++++---------
tests/trainers/test_distillation_metrics.py | 28 ++++++++++++-----
3 files changed, 44 insertions(+), 23 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 669bc8d5d..9aeed8b3b 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -206,6 +206,14 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Lesson**: Contract tests for CUDA-only optional adapters must exercise the adapter through its dependency boundary so CPU and macOS collection remains valid; importing such adapters at module scope turns an optional dependency into a repository-wide test dependency.
- **Related Constraint**: N/A
+### Distillation exact cursors count rollout batches, not backend work items
+- **Date**: 2026-08-29
+- **Symptom**: After timestep-aligned TDM accumulation made each trajectory boundary one backend work item, exact resume skipped `num_inference_steps` times too many prompt batches.
+- **Root Cause**: Cursor reconstruction still multiplied completed rollout iterations by backend `gradient_accumulation_steps`, even though one rollout now contributes multiple boundary losses to that accumulation window.
+- **Fix**: `trainers/distillation/distillation_runtime.py` now derives consumed prompt batches through `resolve_rollout_accumulation_steps()`, and the cursor regression locks `gradient_accumulation_steps=8`, four losses per rollout, and two completed iterations to four consumed batches.
+- **Lesson**: Persisted acquisition progress must be projected through the current acquisition-to-backend work-item ratio. Backend GAS is not a valid dataloader cursor when one acquired batch expands into multiple backward graphs.
+- **Related Constraint**: #18a
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/trainers/distillation/distillation_runtime.py b/src/flow_factory/trainers/distillation/distillation_runtime.py
index daab3ab81..15d8217ac 100644
--- a/src/flow_factory/trainers/distillation/distillation_runtime.py
+++ b/src/flow_factory/trainers/distillation/distillation_runtime.py
@@ -488,7 +488,6 @@ def _completed_rollout_batch_count(trainer: Any) -> int:
progress = getattr(trainer, "progress", None)
completed_iterations = getattr(progress, "rollout_iteration", 0)
training_args = getattr(trainer, "training_args", None)
- accumulation_steps = getattr(training_args, "gradient_accumulation_steps", 1)
if (
not isinstance(completed_iterations, int)
or isinstance(completed_iterations, bool)
@@ -498,16 +497,8 @@ def _completed_rollout_batch_count(trainer: Any) -> int:
"expected rollout_iteration >= 0 as an int, received "
f"{type(completed_iterations).__name__}: {completed_iterations!r}"
)
- if (
- not isinstance(accumulation_steps, int)
- or isinstance(accumulation_steps, bool)
- or accumulation_steps < 1
- ):
- raise ValueError(
- "expected gradient_accumulation_steps >= 1 as an int, received "
- f"{type(accumulation_steps).__name__}: {accumulation_steps!r}"
- )
- return completed_iterations * accumulation_steps
+ rollout_accumulation_steps = resolve_rollout_accumulation_steps(training_args)
+ return completed_iterations * rollout_accumulation_steps
def _collect_rollout_loader_generators(dataloader: Any) -> List[torch.Generator]:
@@ -593,10 +584,11 @@ def _restore_rollout_data_cursor(
"""Rebuild one deterministic loader iterator at a global batch boundary.
Exact checkpoints are published only between acquisition cycles. Each completed
- distillation cycle consumes exactly ``gradient_accumulation_steps`` batches, so
- the persisted rollout-iteration counter is the authoritative cursor. Rebuilding
- from it avoids serializing a Python iterator and works for both finite
- multi-source loaders and the framework's infinite grouped batch samplers.
+ distillation cycle consumes ``gradient_accumulation_steps`` divided by its
+ per-rollout loss count, so the persisted rollout-iteration counter is the
+ authoritative cursor. Rebuilding from it avoids serializing a Python iterator
+ and works for both finite multi-source loaders and the framework's infinite
+ grouped batch samplers.
"""
if (
not isinstance(consumed_batches, int)
@@ -854,6 +846,15 @@ def resolve_rollout_accumulation_steps(training_args: Any) -> int:
"""Recover rollout batches from timestep-aligned backend GAS."""
accumulation_steps = training_args.gradient_accumulation_steps
losses_per_rollout = training_args.get_num_train_timesteps(None)
+ if (
+ not isinstance(accumulation_steps, int)
+ or isinstance(accumulation_steps, bool)
+ or accumulation_steps < 1
+ ):
+ raise ValueError(
+ "expected gradient_accumulation_steps >= 1 as an int, received "
+ f"{type(accumulation_steps).__name__}: {accumulation_steps!r}"
+ )
if (
not isinstance(losses_per_rollout, int)
or isinstance(losses_per_rollout, bool)
diff --git a/tests/trainers/test_distillation_metrics.py b/tests/trainers/test_distillation_metrics.py
index 57bfcf07f..1038ffce2 100644
--- a/tests/trainers/test_distillation_metrics.py
+++ b/tests/trainers/test_distillation_metrics.py
@@ -128,13 +128,19 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
return {"value": torch.tensor(self.values[index])}
-def _rollout_cursor_trainer(progress: TrainingProgress) -> SimpleNamespace:
+def _rollout_cursor_trainer(
+ progress: TrainingProgress,
+ *,
+ accumulation_steps: int = 2,
+ losses_per_rollout: int = 1,
+) -> SimpleNamespace:
"""Build the one-batch rollout surface around an infinite grouped loader."""
return SimpleNamespace(
progress=progress,
training_args=SimpleNamespace(
- gradient_accumulation_steps=2,
+ gradient_accumulation_steps=accumulation_steps,
num_batches_per_epoch=3,
+ get_num_train_timesteps=lambda config: losses_per_rollout,
),
dataloader=_InfiniteGroupedLoader(),
adapter=SimpleNamespace(rollout=lambda: None),
@@ -185,6 +191,7 @@ def sample_batch(batch: Dict[str, torch.Tensor], **kwargs: Any) -> list[tuple]:
training_args=SimpleNamespace(
gradient_accumulation_steps=2,
num_batches_per_epoch=3,
+ get_num_train_timesteps=lambda config: 1,
),
dataloader=dataloader,
adapter=SimpleNamespace(rollout=lambda: None),
@@ -256,6 +263,7 @@ def sample_batch(batch: Dict[str, Any], **kwargs: Any) -> list[tuple[str, int]]:
# independently round their quotas. The finite wrapper length is the
# authoritative result when those two values differ.
num_batches_per_epoch=2,
+ get_num_train_timesteps=lambda config: 1,
),
dataloader=_finite_multi_source_loader(),
adapter=SimpleNamespace(rollout=lambda: None),
@@ -614,9 +622,13 @@ def test_exact_resume_reconstructs_infinite_grouped_rollout_cursor() -> None:
assert resumed.dataloader.batch_sampler.set_epoch_calls == [1]
-def test_exact_resume_uses_gas_to_reconstruct_rollout_batch_count() -> None:
- """Completed rollout iterations expand to the exact number of consumed batches."""
- resumed = _rollout_cursor_trainer(TrainingProgress(rollout_iteration=4))
+def test_exact_resume_uses_rollout_factor_to_reconstruct_batch_count() -> None:
+ """Boundary-aligned GAS must not make exact resume skip extra prompt batches."""
+ resumed = _rollout_cursor_trainer(
+ TrainingProgress(rollout_iteration=2),
+ accumulation_steps=8,
+ losses_per_rollout=4,
+ )
next_batch = generate_one_rollout_batch(
resumed,
@@ -624,9 +636,9 @@ def test_exact_resume_uses_gas_to_reconstruct_rollout_batch_count() -> None:
algorithm_name="TDM",
)
- assert next_batch == [(2, 2)]
- assert resumed._rollout_batches_consumed == 9
- assert resumed.dataloader.batch_sampler.set_epoch_calls == [2]
+ assert next_batch == [(1, 1)]
+ assert resumed._rollout_batches_consumed == 5
+ assert resumed.dataloader.batch_sampler.set_epoch_calls == [1]
@pytest.mark.parametrize("use_explicit_generator", [False, True])
From ea3c1329debd2b91f2bdd8d70fac1370d19ca092 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sat, 29 Aug 2026 11:13:39 +0800
Subject: [PATCH 28/76] chore(checkpointing): add required license headers
---
src/flow_factory/hparams/gradient_checkpointing.py | 14 ++++++++++++++
src/flow_factory/models/checkpointing.py | 14 ++++++++++++++
tests/hparams/test_gradient_checkpointing.py | 14 ++++++++++++++
.../test_selective_gradient_checkpointing.py | 14 ++++++++++++++
4 files changed, 56 insertions(+)
diff --git a/src/flow_factory/hparams/gradient_checkpointing.py b/src/flow_factory/hparams/gradient_checkpointing.py
index e5c476e01..8b21eed19 100644
--- a/src/flow_factory/hparams/gradient_checkpointing.py
+++ b/src/flow_factory/hparams/gradient_checkpointing.py
@@ -1,3 +1,17 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
"""Gradient-checkpointing policy configuration."""
from __future__ import annotations
diff --git a/src/flow_factory/models/checkpointing.py b/src/flow_factory/models/checkpointing.py
index 84a9803ec..6a68475d2 100644
--- a/src/flow_factory/models/checkpointing.py
+++ b/src/flow_factory/models/checkpointing.py
@@ -1,3 +1,17 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
"""Model-level selective gradient-checkpointing utilities."""
from __future__ import annotations
diff --git a/tests/hparams/test_gradient_checkpointing.py b/tests/hparams/test_gradient_checkpointing.py
index 132d95b34..c80f066da 100644
--- a/tests/hparams/test_gradient_checkpointing.py
+++ b/tests/hparams/test_gradient_checkpointing.py
@@ -1,3 +1,17 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
import pytest
from flow_factory.hparams.gradient_checkpointing import (
diff --git a/tests/models/test_selective_gradient_checkpointing.py b/tests/models/test_selective_gradient_checkpointing.py
index 7626da9c4..2af4ba126 100644
--- a/tests/models/test_selective_gradient_checkpointing.py
+++ b/tests/models/test_selective_gradient_checkpointing.py
@@ -1,3 +1,17 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
from types import SimpleNamespace
import pytest
From bcc44677443bf20ba12ac10897d12f5fe8d90246 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sat, 29 Aug 2026 11:15:44 +0800
Subject: [PATCH 29/76] fix(tests): align DMD2 fixture with rollout contract
---
tests/trainers/test_dmd2.py | 1 +
1 file changed, 1 insertion(+)
diff --git a/tests/trainers/test_dmd2.py b/tests/trainers/test_dmd2.py
index 5e698bc9b..f37769760 100644
--- a/tests/trainers/test_dmd2.py
+++ b/tests/trainers/test_dmd2.py
@@ -48,6 +48,7 @@ def _trainer(**training_overrides: Any) -> DMD2Trainer:
"num_inference_steps": 1,
"per_device_batch_size": 1,
"gradient_accumulation_steps": 1,
+ "get_num_train_timesteps": lambda _config: 1,
"ttur_fake_updates": 5,
"num_inner_epochs": 1,
}
From a5d84e63a7907c58aa1e0ff09b52232a83232142 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sat, 29 Aug 2026 11:24:00 +0800
Subject: [PATCH 30/76] chore: normalize imports after parent rebase
---
src/flow_factory/hparams/training_args/__init__.py | 8 ++++----
src/flow_factory/trainers/distillation/tdm.py | 2 +-
src/flow_factory/trainers/distillation/tdm_r1.py | 2 +-
3 files changed, 6 insertions(+), 6 deletions(-)
diff --git a/src/flow_factory/hparams/training_args/__init__.py b/src/flow_factory/hparams/training_args/__init__.py
index cec2f6670..b0955622c 100644
--- a/src/flow_factory/hparams/training_args/__init__.py
+++ b/src/flow_factory/hparams/training_args/__init__.py
@@ -21,6 +21,10 @@
from flow_factory.hparams.training_args import get_training_args_class
"""
+from ..gradient_checkpointing import (
+ GradientCheckpointingPolicy,
+ GradientCheckpointingSpec,
+)
from ._base import EvaluationArguments, TrainingArguments
from ._registry import get_training_args_class, list_registered_training_args
from .awm import AWMTrainingArguments
@@ -36,10 +40,6 @@
from .sft import SFTTrainingArguments
from .tdm import TDMTrainingArguments
from .tdm_r1 import TDMR1TrainingArguments
-from ..gradient_checkpointing import (
- GradientCheckpointingPolicy,
- GradientCheckpointingSpec,
-)
__all__ = [
"EvaluationArguments",
diff --git a/src/flow_factory/trainers/distillation/tdm.py b/src/flow_factory/trainers/distillation/tdm.py
index 5c863d186..0abc6a26b 100644
--- a/src/flow_factory/trainers/distillation/tdm.py
+++ b/src/flow_factory/trainers/distillation/tdm.py
@@ -49,8 +49,8 @@
reference_forward_kwargs,
reject_training_rewards,
replay_forward_kwargs,
- resolve_rollout_accumulation_steps,
require_velocity,
+ resolve_rollout_accumulation_steps,
role_repeat_progress,
run_distillation_training_step,
run_role_phase,
diff --git a/src/flow_factory/trainers/distillation/tdm_r1.py b/src/flow_factory/trainers/distillation/tdm_r1.py
index c00fb8497..554d809cf 100644
--- a/src/flow_factory/trainers/distillation/tdm_r1.py
+++ b/src/flow_factory/trainers/distillation/tdm_r1.py
@@ -34,8 +34,8 @@
generate_one_rollout_batch,
query_score_velocity,
record_distillation_metric,
- resolve_rollout_accumulation_steps,
require_velocity,
+ resolve_rollout_accumulation_steps,
role_repeat_progress,
run_role_phase,
)
From 9857aa71b3ecfe7d2bcc42e358c8d557d9f8f535 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sat, 29 Aug 2026 14:35:44 +0800
Subject: [PATCH 31/76] feat(offline): support Wan LTX2 and MiniMax H3
workflows
---
.agents/knowledge/architecture.md | 25 +-
.agents/knowledge/dependencies.md | 7 +-
.../knowledge/topics/adapter_conventions.md | 27 +-
.agents/knowledge/topics/fix_patterns.md | 103 ++
.agents/knowledge/topics/minimax_h3.md | 51 +-
.../topics/train_inference_consistency.md | 9 +-
README.md | 38 +-
docker/README.md | 7 +-
docker/docker-cuda/Dockerfile | 5 +-
examples/README.md | 6 +-
guidance/datasets.md | 72 +-
guidance/gpu_validation.md | 240 ++++
guidance/new_model.md | 50 +-
guidance/workflow.md | 34 +-
src/flow_factory/contracts/__init__.py | 2 +
src/flow_factory/contracts/pipeline_io.py | 259 +++-
src/flow_factory/data_utils/dataset.py | 53 +-
.../data_utils/offline_condition_cache.py | 188 ++-
.../data_utils/offline_dataset.py | 1 +
.../data_utils/offline_train_data.py | 13 +-
src/flow_factory/data_utils/schema.py | 23 +
src/flow_factory/models/__init__.py | 9 +
src/flow_factory/models/abc.py | 202 +++-
src/flow_factory/models/condition_state.py | 203 ++++
.../models/configured_image_output.py | 2 +-
src/flow_factory/models/ltx2/_output.py | 1066 +++++++++++++++++
src/flow_factory/models/ltx2/ltx2_i2av.py | 124 +-
src/flow_factory/models/ltx2/ltx2_t2av.py | 105 +-
.../models/minimax_h3/__init__.py | 8 +-
.../models/minimax_h3/_condition.py | 99 ++
src/flow_factory/models/minimax_h3/_output.py | 92 +-
.../models/minimax_h3/adapters.py | 151 ++-
src/flow_factory/models/minimax_h3/blocks.py | 252 +++-
.../models/minimax_h3/workflow.py | 81 +-
src/flow_factory/models/pipeline_contracts.py | 20 +-
src/flow_factory/models/wan/_conditioning.py | 527 ++++++++
src/flow_factory/models/wan/_output.py | 208 +++-
src/flow_factory/models/wan/wan2_i2v.py | 421 ++++---
src/flow_factory/models/wan/wan2_t2v.py | 164 +--
src/flow_factory/trainers/abc.py | 1 +
.../trainers/common/flow_matching.py | 10 +-
.../trainers/common/offline_batch.py | 31 +-
.../trainers/common/runtime_identity.py | 6 +
.../trainers/offline/offline_dpo.py | 19 +-
src/flow_factory/trainers/offline/sft.py | 12 +-
tests/contracts/test_pipeline_io_contract.py | 249 +++-
.../test_offline_condition_cache.py | 295 +++++
tests/data_utils/test_offline_train_data.py | 7 +-
tests/data_utils/test_schema.py | 25 +
tests/docs/test_minimax_h3_docs.py | 41 +-
tests/models/ltx2/test_ltx2_adapter_init.py | 29 +-
.../models/ltx2/test_ltx2_component_hooks.py | 29 +
tests/models/ltx2/test_ltx2_output_codec.py | 526 ++++++++
.../models/minimax_h3/test_condition_state.py | 116 ++
.../minimax_h3/test_diffusers_api_contract.py | 154 ++-
tests/models/minimax_h3/test_modular_core.py | 51 +-
tests/models/minimax_h3/test_output_codec.py | 234 +++-
.../minimax_h3/test_workflow_execution.py | 91 +-
.../test_offline_output_capability_matrix.py | 37 +-
.../test_output_state_adapter_lifecycle.py | 126 +-
tests/models/test_wan_output_codec.py | 351 +++++-
.../test_base_adapter_trajectory.py | 43 +
.../trainers/test_offline_batch_primitives.py | 26 +
tests/trainers/test_offline_flow_matching.py | 2 +-
tests/trainers/test_offline_trainers.py | 31 +-
tests/trainers/test_runtime_identity.py | 17 +
66 files changed, 6771 insertions(+), 735 deletions(-)
create mode 100644 guidance/gpu_validation.md
create mode 100644 src/flow_factory/models/condition_state.py
create mode 100644 src/flow_factory/models/ltx2/_output.py
create mode 100644 src/flow_factory/models/minimax_h3/_condition.py
create mode 100644 src/flow_factory/models/wan/_conditioning.py
create mode 100644 tests/models/ltx2/test_ltx2_output_codec.py
create mode 100644 tests/models/minimax_h3/test_condition_state.py
diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md
index 8496d9f7c..84baee949 100644
--- a/.agents/knowledge/architecture.md
+++ b/.agents/knowledge/architecture.md
@@ -61,7 +61,8 @@ exhaustion; offline output media is decoded and encoded on the fly, while only p
conditions enter the preprocessing cache.
Exact runtime identity is built from realized prepared state. It locks optimizer/model/backend
-semantics, ordered training data, and the complete replayed evaluation path (cadence, arguments,
+semantics, the checkpoint-realized pipeline I/O contract, ordered training data, and the complete
+replayed evaluation path (cadence, arguments,
per-dataset overrides, rewards, and ordered prepared loaders). Logging, checkpoint cadence, run
budget, and resume location remain operational. Exact-state save fails before mutation on MPS
because Accelerate does not persist the device RNG needed for exact continuation.
@@ -89,6 +90,9 @@ All four registries map string keys → lazy import paths. Resolution: registry
| `awm` | `AWMTrainer` | Decoupled | `BaseTrainer` |
| `crd` | `CRDTrainer` | Decoupled | `BaseTrainer` |
| `diffusion-opd` | `DiffusionOPDTrainer` | Distillation (on-policy) | `BaseTrainer` |
+| `dmd2` | `DMD2Trainer` | Distillation | `BaseTrainer` |
+| `tdm` | `TDMTrainer` | Distillation | `BaseTrainer` |
+| `tdm-r1` | `TDMR1Trainer` | Distillation + reward | `BaseTrainer` |
**Flat hierarchy**: New trainers inherit from `BaseTrainer` directly. The sanctioned exceptions are `GRPOGuardTrainer → GRPOTrainer` and `DPPOTrainer → GRPOTrainer` (strict GRPO loss variants; see constraint #11).
@@ -107,6 +111,9 @@ All four registries map string keys → lazy import paths. Resolution: registry
| `wan2_i2v` | `Wan2_I2V_Adapter` | Image-to-Video |
| `ltx2_t2av` | `LTX2_T2AV_Adapter` | Text-to-Audio-Video |
| `ltx2_i2av` | `LTX2_I2AV_Adapter` | Image-to-Audio-Video |
+| `minimax-h3-t2va` | `MiniMaxH3T2VAAdapter` | Text-to-Video-Audio |
+| `minimax-h3-fl2va` | `MiniMaxH3FL2VAAdapter` | Sparse First/Last-Frame-to-Video-Audio |
+| `minimax-h3-ref2va` | `MiniMaxH3Ref2VAAdapter` | Ordered-Reference-to-Video-Audio |
| `bagel` | `BagelAdapter` | Text-to-Image & Image(s)-to-Image (T2I & I2I both batched via NaViT packing; subset-round packing handles variable I2I reference-image count, no per-sample fallback — see `topics/adapter_conventions.md`) |
| `sensenova` | `SenseNovaAdapter` | Text-to-Image & Image(s)-to-Image (SenseNova-U1 1.0/1.5; ordered variable-count references remain grouped in `images` and preserve within-type order; independent samples use B=1 prefixes rather than Bagel-style NaViT packing) |
@@ -157,15 +164,25 @@ Timesteps are `[0, 1000]` (scheduler scale); sigmas are `[0, 1]` (flow-matching
Each model adapter wraps a diffusers pipeline into the `BaseAdapter` interface:
- `preprocess_func()` — prompt/input-condition preprocessing and cache projection
- `pipeline_io_contract` — model-neutral input/output modality and geometry declaration
+- `effective_pipeline_io_contract` — checkpoint-realized specialization of the class contract
+- `prepare_condition_state()` — one validated input-owned runtime realization reused across
+ candidate encoding and model forwards
- `encode_output_state()` — validated on-the-fly offline target encoding through an optional codec
- `inference()` — full denoising loop (Stage 3)
- `forward()` — single-step denoising (Stage 6)
**Per-modality encoders** (`encode_prompt`, `encode_image`, `encode_video`, `encode_audio`) are no-op by default on `BaseAdapter` — override only the modalities your model consumes. `preprocess_func` dispatches to all four and skips any that return `None`, so text/image/video-only adapters need no stub overrides for unused modalities.
-Offline codecs declare logical required components without materializing them. Condition/output
-encoders share role-neutral transforms where possible, while callers retain explicit official
-posterior `sample` versus `argmax` semantics.
+Offline condition preparers and codecs declare logical required components without materializing
+them. Condition/output encoders share role-neutral transforms where possible, while callers retain
+explicit official posterior `sample` versus `argmax` semantics. Candidate-specific output context
+cannot overwrite cached or prepared input fields. A separate flow-matching objective reducer lets
+multi-modal SFT/DPO specialize loss aggregation without changing online trajectory reductions.
+
+Input contracts may declare semantic media slots and aggregate cross-type cardinality rules. In
+strict V2 data, an explicit input-only `slot` reserves its argument; unslotted media fills remaining
+slots in declaration order. Outputs reject slots. This keeps algorithm data model-neutral while the
+adapter owns bindings such as first/last frame and ordered heterogeneous references.
**Flat hierarchy**: All adapters inherit directly from `BaseAdapter` — never from another adapter (see constraint #12). Shared logic within a model family uses helper functions, code duplication, or mixins — not adapter subclassing.
diff --git a/.agents/knowledge/dependencies.md b/.agents/knowledge/dependencies.md
index c3e2e41ac..d6d52ef18 100644
--- a/.agents/knowledge/dependencies.md
+++ b/.agents/knowledge/dependencies.md
@@ -42,7 +42,7 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
| `torchvision` | >= 0.19.0 | Vision utilities |
| `torchaudio` | >= 2.4.0 | Audio I/O (audio / audio-video models, CLAP) |
| `transformers` | >= 4.57.1 | Text encoders, tokenizers |
-| `diffusers` | >= 0.37.0 | Diffusion pipelines, schedulers |
+| `diffusers` | >= 0.40.0 | Diffusion pipelines, schedulers, MiniMax H3 and LTX2 APIs |
| `accelerate` | >= 1.11.0 | Distributed training, mixed precision |
| `peft` | >= 0.17.0 | LoRA, parameter-efficient fine-tuning |
| `datasets` | >= 3.3.2 | Dataset loading |
@@ -57,6 +57,8 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
- DeepSpeed is optional — Accelerate alone handles most distributed scenarios.
### diffusers
+- Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule
+ may be used for upstream development, but must not silently override the declared runtime dependency.
- Model adapters depend on specific pipeline classes from diffusers. Major version bumps may rename or remove pipeline classes.
- `load_pipeline()` in each adapter returns a `DiffusionPipeline`-compatible object; breaking changes in diffusers' pipeline API require adapter updates.
@@ -75,7 +77,8 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
### accelerate
- Primary distributed backend. `accelerator.prepare()` wraps a single `ModelBundle` (all target components) plus the optimizer as one root (constraint #9).
-- The dataloader uses custom samplers and is NOT prepared via accelerate.
+- Online generation uses framework samplers. Finite SFT/offline-DPO data uses PyTorch's official
+ `DistributedSampler`; that already-sharded loader is not prepared via Accelerate.
### peft
- Provides LoRA functionality. Applied via `BaseAdapter.apply_lora()` to components listed in `target_components`.
diff --git a/.agents/knowledge/topics/adapter_conventions.md b/.agents/knowledge/topics/adapter_conventions.md
index 8f2ce5dd0..595e281d5 100644
--- a/.agents/knowledge/topics/adapter_conventions.md
+++ b/.agents/knowledge/topics/adapter_conventions.md
@@ -44,18 +44,18 @@ the complete mapping when its forward has different semantics or additional guid
condition `guidance_scale=3.5`; this is a learned model embedding, not classifier-free guidance.
- The currently supported Flux2-Klein forward always passes `guidance=None` into its transformer,
so its `guidance_scale` remains conventional two-pass CFG and inherits the neutral `1.0`.
-- Wan T2V neutralizes both transformer stages with `guidance_scale=guidance_scale_2=1.0`.
+- Wan T2V/I2V neutralize both transformer stages with
+ `guidance_scale=guidance_scale_2=1.0`.
- SenseNova neutralizes text and image guidance together and disables CFG normalization.
- Bagel replaces the base mapping with its actual `cfg_text_scale` / `cfg_img_scale` arguments;
it must not inherit an irrelevant `guidance_scale` key through its permissive `**kwargs`.
-- MiniMax H3 T2VA inherits neutral `guidance_scale=1.0`; its strict forward validates that
- interface value even though the guidance-distilled checkpoint has no CFG branch.
+- MiniMax H3 T2VA/FL2VA/Ref2VA inherit neutral `guidance_scale=1.0`; their strict forward validates
+ that interface value even though the guidance-distilled checkpoint has no CFG branch.
+- LTX2 sets video/audio CFG scales and modality scales to `1.0`, CFG rescale and STG scales to
+ `0.0`, and neutralizes STG block selection with
+ `spatio_temporal_guidance_blocks=None`.
-Wan I2V and LTX2 remain behind explicit offline output-codec blockers. Wan I2V must mirror the two
-neutral Wan transformer scales before it is enabled. LTX2 must set video/audio CFG scales and
-modality scales to `1.0`, CFG rescale and STG scales to `0.0`, and neutralize its STG block
-selection with `spatio_temporal_guidance_blocks=None`. The mapping is adapter-owned model
-conditioning, never a sampling or algorithm knob.
+These mappings are adapter-owned model conditioning, never sampling or algorithm knobs.
### Models with model-specific CFG extensions
@@ -210,6 +210,17 @@ LTX2 packs `[video|audio]` into one `(B, Seq, C)` sequence, so it resolves as PA
14. **Offline forward overrides are model conditioning, not sampling CFG** — SFT and offline DPO expand the adapter's complete immutable `offline_training_forward_overrides` mapping into every policy and reference forward, after batch conditions and configured sampling arguments. Conventional CFG adapters use their CFG-off point; guidance-distilled adapters use the value expected by their learned guidance embedder; multi-branch adapters neutralize every active branch under its real forward argument names. Replace the base mapping rather than adding unrelated keys, and never infer it from cached negative embeddings or expose it as an algorithm knob.
+15. **Runtime condition realization has one input owner** — A conditioned offline adapter declares
+ `build_condition_state_preparer()` only when cached fields are not already the exact forward
+ condition. `prepare_condition_state()` runs once per batch. SFT reuses that realization for
+ target binding; offline DPO reuses the same tensor leaves for chosen/rejected and
+ policy/reference forwards. The input-owned `forward_context` and `output_context` may share
+ prepared tensors. Require collision-free keys only within each consumer's merged view: cached
+ plus forward fields, and input-owned plus candidate-owned output fields. The preparer declaration
+ lists logical components but must not materialize, move, replace, or cast them. Checkpoint
+ variants narrow the class-level I/O superset through `_resolve_pipeline_io_contract()` rather
+ than branching inside the dataset or algorithm.
+
## Fix Records
### Sampling CFG leaked into finite-data velocity matching
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 9aeed8b3b..401c79356 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -214,6 +214,109 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Lesson**: Persisted acquisition progress must be projected through the current acquisition-to-backend work-item ratio. Backend GAS is not a valid dataloader cursor when one acquired batch expands into multiple backward graphs.
- **Related Constraint**: #18a
+### Sparse media arguments require semantic input slots
+- **Date**: 2026-08-29
+- **Symptom**: A last-frame-only MiniMax H3 record could not be represented without pretending its
+ image was the first frame, and heterogeneous Ref2VA cardinality rules could not be expressed by
+ independent per-type counts.
+- **Root Cause**: The public offline schema and pipeline contract treated media position and
+ per-modality cardinality as the complete binding model.
+- **Fix**: V2 input media now accepts an input-only semantic `slot`; contracts declare ordered and
+ required slots plus aggregate count/required-any-type rules; projection resolves explicit slots
+ first and fills remaining slots positionally. Outputs reject slots. Construction rejects
+ multi-slot rules that claim order-insensitivity and aggregate bounds that cannot satisfy their
+ per-type rules.
+- **Lesson**: Use generic argument-binding metadata for sparse conditions and keep model-specific
+ argument names in adapter contracts, not algorithm code or ad-hoc dataset columns.
+- **Related Constraint**: #5
+
+### Offline velocity objectives must bypass unused scheduler transitions
+- **Date**: 2026-08-29
+- **Symptom**: LTX2 near-clean offline targets lost velocity precision after a velocity-to-x0-to-
+ velocity round trip, while exact velocity-only Wan/LTX forwards still invoked scheduler steps that
+ their loss never consumed. LTX2 I2AV also dropped cached negative prompts during preprocessing.
+- **Root Cause**: Generation-oriented forward paths performed transition reconstruction before
+ checking the requested offline component, and the I2AV preprocessing override failed to forward
+ one base prompt argument.
+- **Fix**: LTX2 retains official online reconstruction by default but opts offline forwards into raw
+ model velocity; Wan and LTX return exact velocity requests before scheduler stepping; I2AV now
+ forwards `negative_prompt` explicitly. Component, parity, and initialization regressions cover
+ the split behavior.
+- **Lesson**: Offline objectives may share a model forward with generation but must not inherit
+ numerically lossy or RNG-consuming transition work that is outside their requested output.
+- **Related Constraint**: #7
+
+### Exact resume must lock the checkpoint-realized pipeline contract
+- **Date**: 2026-08-29
+- **Symptom**: Exact resume could accept a checkpoint after an in-place model configuration change
+ switched a Wan adapter between first-only and first/last-frame semantics.
+- **Root Cause**: Runtime identity hashed model arguments and trainer execution but omitted the
+ adapter's resolved `effective_pipeline_io_contract`.
+- **Fix**: The default execution identity now canonicalizes and hashes the realized pipeline I/O
+ contract after adapter initialization; a regression changes only that contract and observes only
+ the execution digest change.
+- **Lesson**: Any checkpoint-dependent specialization that changes legal inputs or forward binding
+ is future-execution state and belongs in exact-resume identity.
+- **Related Constraint**: #18
+
+### Offline condition caches need contract-stable schemas across sources
+- **Date**: 2026-08-29
+- **Symptom**: Changing only semantic slot order could reuse an Arrow cache with the old media
+ projection, while a multi-source batch could fail because an all-empty optional source omitted
+ columns that a populated source emitted.
+- **Root Cause**: The source hash covered record identities but not the effective input projection
+ contract, and projection decided column existence from each source's observed values.
+- **Fix**: Condition source identity now includes the canonical effective input contract. With a
+ contract, negative-prompt, declared media, and semantic-slot columns are projected consistently
+ even when every row in one source is empty. A real two-source `DistributedSampler` loader
+ regression mixes empty and populated optional conditions in one batch.
+- **Lesson**: A concatenated cache schema is defined by the model contract, not by local source
+ sparsity; cache identity must cover every declaration that can reorder or reshape projection.
+- **Related Constraint**: #9
+
+### Distributed Arrow schemas must be inferred before rank sharding
+- **Date**: 2026-08-29
+- **Symptom**: A distributed condition-cache build could write `List(null)` on an all-empty rank
+ and `List(Image)` on a populated rank, then fail when the per-rank Arrow files were consolidated.
+- **Root Cause**: Cross-chunk schema discovery ran after rank sharding, so each process inferred
+ features from only its local value distribution. The standalone cache entry point also ignored a
+ pipeline's single-sample preprocessing capability.
+- **Fix**: Distributed preprocessing now derives one explicit feature schema from the full source
+ before selecting rank-local rows, and both cache entry points force batch size one for ordered or
+ `SINGLE_SAMPLE` contracts. A real two-part Arrow regression separates empty and populated rows
+ across ranks, consolidates the files, and loads the merged dataset.
+- **Lesson**: Distributed writers need a global serialization contract even when their data is
+ disjoint. Batch capability is likewise part of the preprocessing boundary, not only trainer
+ orchestration.
+- **Related Constraint**: #9
+
+### Validate output candidates before stochastic condition preparation
+- **Date**: 2026-08-29
+- **Symptom**: An invalid offline target correctly raised an exception but first consumed condition
+ preparation RNG, so retrying with corrected media no longer reproduced the original encoding.
+- **Root Cause**: `BaseAdapter.encode_output_state()` prepared raw conditions before validating the
+ generator and exact output-media sequence.
+- **Fix**: The lifecycle wrapper now validates generator type and candidate media before invoking
+ any condition preparer or codec. A stochastic-preparer regression proves invalid media leaves the
+ explicit generator unchanged and performs no preparation work.
+- **Lesson**: Pure boundary validation must precede expensive or random transformations. Failed
+ inputs should not mutate the state that determines a later valid retry.
+- **Related Constraint**: #7
+
+### Aggregate media guarantees must be canonical contract state
+- **Date**: 2026-08-29
+- **Symptom**: `INPUT_MEDIA` geometry rejected a valid contract whose aggregate `min_total_count` or
+ `required_any_types` guaranteed a condition, while semantically identical required-type tuples
+ in different orders produced different cache and resume identities.
+- **Root Cause**: Geometry validation recognized only per-type minima, and the set-like aggregate
+ field had no canonical ordering rule.
+- **Fix**: Input-derived geometry now accepts every nonempty guarantee enforced by runtime
+ validation, `required_any_types` must follow canonical media-type order, and required slots must
+ follow their declaration order.
+- **Lesson**: Declarative invariants should be interpreted consistently at construction and runtime,
+ and set-like identity fields require one canonical representation.
+- **Related Constraint**: #5
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/.agents/knowledge/topics/minimax_h3.md b/.agents/knowledge/topics/minimax_h3.md
index fd41f1962..d4f3048da 100644
--- a/.agents/knowledge/topics/minimax_h3.md
+++ b/.agents/knowledge/topics/minimax_h3.md
@@ -45,40 +45,60 @@ and lets that boundary stay strict.
## Input contracts
- T2VA accepts prompt-only workflow input.
-- FL2VA accepts one first image or two images ordered first then last.
-- Ref2VA preserves and hashes ordered image/video/audio manifests.
+- FL2VA accepts semantic `first_frame`, `last_frame`, or both image slots. Unslotted legacy
+ generation input retains first-then-last positional shorthand; strict V2 can express last-only.
+- Ref2VA preserves and hashes 1-12 ordered image/video/audio references and requires at least one
+ image or video.
- Ref2VA declares `supports_ordered_references=True`; all H3 adapters explicitly declare hidden
geometry cache fields and a preprocessing cache version.
- Reference paths are dataset-relative. Positive finite `fps` and `sample_rate` overrides follow
`samples/references.py`.
- PyAV >=18.0.0 decodes video/audio references, including embedded or separate soundtracks.
-## Offline T2VA output contract
+## Offline audiovisual output contract
-`minimax-h3-t2va` supports SFT and offline DPO with one exact ordered output pair:
-video first, then audio. Both `fps` and `sample_rate` are required in V2 supervision.
+All three H3 workflows support SFT and offline DPO with one exact ordered output pair: video first,
+then audio. Both `fps` and `sample_rate` are required in V2 supervision.
Targets are decoded on demand; neither pixels, waveforms, nor VAE latents enter the condition
cache. The pipeline's single-sample capability also forces condition-cache preprocessing to B=1,
independently of the global preprocessing batch-size setting.
-The codec cross-validates cached T2VA layout and geometry against the current training config. It
+The codec cross-validates cached layout and geometry against the current training config. It
resamples video onto the configured fixed 24-fps grid and canvas, truncates audio on its declared
source clock before a single conversion to the audio-VAE rate, and aligns stereo audio to the exact
-latent duration, samples and normalizes the video posterior, takes and normalizes the official
-audio posterior mode, then packs structured rows in `("video", "audio")` order. The codec does
+latent duration, takes and normalizes both posterior modes, then packs structured rows in
+`("video", "audio")` order. The deterministic video mode follows the H3 SFT reference data flow;
+the released Diffusers pipeline does not define an offline target encoder. The codec does
not duplicate input-owned fields in output forward context. Replay nests the flat cached layout
and derives empty T2VA condition prefixes from the current state, preserving storage dtype and
device. Exact velocity-only offline forwards return before either component scheduler steps, so
SFT and offline DPO do not sample unused transitions or perturb scheduler RNG cadence. Every
encoded row count must match the cached layout before transformer execution.
-FL2VA and Ref2VA remain online-only. Their output AV encoding can reuse the T2VA numerical
-codec, but their cached media conditions still need a separately owned, reproducible
-condition-prefix binder. In particular, both offline-DPO arms must consume the same conditioned
-prefix noise; do not generate those prefixes independently inside the chosen/rejected codecs.
+FL2VA and Ref2VA use a separately owned runtime condition-prefix preparer. It realizes the official
+condition noise once per batch, then exposes immutable model-forward and output-binding views. Both
+offline-DPO arms and their policy/reference forwards consume the same prepared prefix object; target
+encoding never draws condition noise independently. Offline flow matching sums the separate video
+and audio means, while online likelihood and distillation retain their existing globally
+element-weighted reducer.
## Fix records
+### H3 offline targets preserve the reference posterior and modality objective
+
+- **Date**: 2026-08-29
+- **Symptom**: H3 clean video targets changed on every encode and audiovisual flow loss weighted
+ modalities by tensor cardinality, heavily downweighting audio.
+- **Root Cause**: The first offline codec inferred stochastic video-posterior sampling where
+ Diffusers has no target recipe, and inherited the online global reducer for a two-term SFT loss.
+- **Fix**: Video and audio target codecs now take deterministic posterior modes, matching the H3
+ SFT reference data flow, and the offline flow hook returns `video_mean + audio_mean` without
+ changing the online reducer. Unequal-cardinality regression tests lock both decisions.
+- **Lesson**: When an inference pipeline omits training semantics, use the nominated training
+ reference for posterior selection and keep objective-specific modality weighting separate from
+ trajectory likelihood aggregation.
+- **Related Constraint**: #7, #8
+
### Offline targets preserve configured geometry and logical source clocks
- **Date**: 2026-08-28
@@ -122,11 +142,12 @@ prefix noise; do not generate those prefixes independently inside the chosen/rej
## Verification boundary
-All workflows have pinned API/schema/no-weight verification. T2VA additionally completed
+All workflows have pinned API/schema/no-weight verification and local offline codec/forward
+coverage. T2VA additionally completed
real-weight LoRA rollout, decode, reward, replay, backward, checkpoint, and resume tests on one
GPU and with FSDP2 on 16 GPUs. The native-resolution path completed initialization, checkpoint,
-decode, and evaluation. FL2VA and Ref2VA remain no-weight validated. Do not claim long-run reward
-improvement, convergence, or numerical parity.
+decode, and evaluation. FL2VA/Ref2VA SFT and offline DPO still require the documented real-weight
+GPU matrix. Do not claim long-run reward improvement, convergence, or numerical parity.
## Upgrade checklist
diff --git a/.agents/knowledge/topics/train_inference_consistency.md b/.agents/knowledge/topics/train_inference_consistency.md
index fac9faa8e..fda66c548 100644
--- a/.agents/knowledge/topics/train_inference_consistency.md
+++ b/.agents/knowledge/topics/train_inference_consistency.md
@@ -39,7 +39,14 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch
7. **Batch/pack composition mismatch (pack-dependent adapters)**: For adapters whose batched `forward()` is *pack-composition-dependent* (e.g. Bagel NaViT sequence packing, where a sample's linear-projection matmuls run over the concatenated `sum_seqlen` of the whole pack), bf16 rounding depends on *which* samples share the pack. If a training micro-batch packs a different sample set than the corresponding rollout pack, the on-policy `forward()` is no longer bit-identical -> `ratio != 1` (~1e-4) even though every stored argument matches. Per-sample (B=1) adapters are immune. The trigger is the optimize-time sample shuffle reordering `samples` before chunking into micro-batches.
8. **Stochastic conditioning encoder recomputed per forward**: When `forward()` *rebuilds* its conditioning from raw inputs each call (rather than replaying a stored embedding) and that encoder is non-deterministic, rollout and training diverge. Bagel I2I is the canonical case: the condition-image VAE (`DiagonalGaussian`, default `sample=True` -> `mean + std*randn`) is encoded **once** at rollout but **re-encoded every training `forward()`** (`_forward_rebuild` -> `_build_gen_context`), so each draws different noise -> different KV context -> on-policy `ratio != 1` (~2e-4). T2I is immune (text-only context, deterministic). Fix: make the condition encode deterministic (`vae.reg.sample = False`, posterior mean) or cache the rollout encoding and replay it. (Only affects `vae.encode` of conditions; generation uses init noise and `vae.decode` is unaffected.) **Bagel applies this fix**: `pipeline.vae.reg.sample = False` in `BagelAdapter.load_pipeline()`.
9. **Preference arms replayed with the wrong batch**: DPO shares forward-process noise across chosen/rejected states, but each arm still owns its conditioning batch. Both policy and reference forwards for the rejected state must receive `rejected_batch`; pairing it with `chosen_batch` evaluates the rejected trajectory under another conditioning context.
-10. **Model-specific velocity direction assumed by a trainer**: Standard flow adapters predict noise-ward velocity (`noise - clean`), while MiniMax H3 predicts data-ward velocity (`clean - noise`). Any `x0` target must use `adapter.project_velocity_to_clean_state()` rather than spelling `xt - sigma * velocity` inside a trainer.
+10. **Condition realization redrawn per candidate**: Geometry-bound VAE conditions and stochastic
+ condition prefixes belong to the input, not to a demonstration/preference candidate. Offline
+ training must call `prepare_condition_state()` once per batch and reuse the resulting tensor
+ leaves for every target arm and policy/reference forward. MiniMax H3 FL2VA/Ref2VA is the
+ stochastic reference: its official condition augmentation must be drawn once before target
+ noise. Wan I2V and LTX2 I2AV use deterministic posterior mode but follow the same ownership
+ boundary so mask/layout binding cannot drift between candidates.
+11. **Model-specific velocity direction assumed by a trainer**: Standard flow adapters predict noise-ward velocity (`noise - clean`), while MiniMax H3 predicts data-ward velocity (`clean - noise`). Any `x0` target must use `adapter.project_velocity_to_clean_state()` rather than spelling `xt - sigma * velocity` inside a trainer.
## Pack-composition-dependent adapters: `shuffle_samples`
diff --git a/README.md b/README.md
index 0b9954858..8c0a23ca1 100644
--- a/README.md
+++ b/README.md
@@ -19,10 +19,9 @@ pip install 'diffusers>=0.40.0'
pip install -e .
```
-* **[2026-04-25]** **LTX-2 Audio-Video** support! Generate synchronized audio-video content with RL fine-tuning. LTX-2 requires the bundled `diffusers` submodule (not yet in the official release):
+* **[2026-04-25]** **LTX-2 Audio-Video** support! Generate synchronized audio-video content with RL fine-tuning through the released Diffusers API:
```bash
-git submodule update --init
-pip install -e ./diffusers
+pip install 'diffusers>=0.40.0'
```
* **[2026-02-01]** Support for multiple **Attention Backends**! Attention-backend selection now lives in the unified `acceleration:` block (the old `model.attn_backend` knob was removed), where it can be combined with `torch.compile` and feature caching — applied in list order:
@@ -80,8 +79,7 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
| Wan2.2-TI2V-5B | 5B | wan2_t2v |
| Wan2.2-T2V-A14B | A14B | wan2_t2v |
- | Image-to-Video | Wan2.1-I2V-14B-480P | 14B | wan2_i2v |
- | Wan2.1-I2V-14B-480P | 14B | wan2_i2v |
+ | Image-to-Video | Wan2.1-I2V-14B-480P | 14B | wan2_i2v |
| Wan2.1-I2V-14B-720P | 14B | wan2_i2v |
| Wan2.2-TI2V-5B | 5B | wan2_i2v |
| Wan2.2-I2V-A14B | A14B | wan2_i2v |
@@ -99,17 +97,18 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
> **Offline output support:** SFT and offline DPO currently support `sd3-5`, `flux1`,
> `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`,
-> `bagel`, `sensenova`, `wan2_t2v`, and `minimax-h3-t2va`. Wan I2V, LTX2, and the
-> conditioned MiniMax H3 FL2VA/Ref2VA workflows fail fast on their currently unresolved
-> output/condition semantics. MiniMax H3 T2VA targets use an exact ordered video/audio pair,
-> encoded on demand into its structured latent state. See the
+> `bagel`, `sensenova`, `wan2_t2v`, `wan2_i2v`, `ltx2_t2av`, `ltx2_i2av`, and all
+> MiniMax H3 workflows. Video/audio targets are encoded on demand and are never written to
+> the preprocessing cache. Conditioned adapters prepare one immutable condition state per batch;
+> offline DPO shares that exact realization across chosen and rejected arms. See the
> [offline model matrix](guidance/datasets.md#offline-model-support).
> **MiniMax H3 status:** the T2VA debug and
> [native-quality FSDP2](examples/grpo/lora/minimax_h3_t2va/quality_720p_fsdp2.yaml)
> paths are real-weight
> validated; a completed long-run reward trend is not claimed. FL2VA and Ref2VA remain
-> schema/API validated. H3 requires B=1, has no CFG, uses neutral guidance `1.0`, and
+> schema/API and local offline-path validated, pending the documented real-weight GPU matrix.
+> H3 requires B=1, has no CFG, uses neutral guidance `1.0`, and
> keeps separate video/audio trajectories.
> Video uses shift 12, audio uses shift 3, and the model predicts data-ward velocity.
> `num_inference_steps=N` means N transitions and N + 1 states.
@@ -160,16 +159,10 @@ pip install -e .[deepspeed]
> **Note**: The Bagel adapter requires `flash-attn` (>= 2.5.8) and `opencv-python`. Install them with `pip install -e .[bagel]` (the `[bagel]` extra is intentionally not part of `[all]` because flash-attn is heavy to build).
-> **Dependency:** MiniMax H3 requires `diffusers>=0.40.0`. PyAV >=18.0.0 decodes
-> ordered video/audio references.
-
-> **Note**: Some models (e.g., LTX-2) require pipeline code not yet released in the official `diffusers` package. For these models, install the bundled diffusers submodule:
-> ```bash
-> git submodule update --init
-> pip install -e ./diffusers
-> ```
+> **Dependency:** MiniMax H3 and LTX2 require the released `diffusers>=0.40.0` API.
+> PyAV >=18.0.0 decodes ordered video/audio references and target media.
-A CUDA training image (Python 3.12, **uv**-based install, PyTorch 2.8 + `cu129`, `deepspeed`, `wandb`, bundled `diffusers`) is defined under [`docker/docker-cuda/`](docker/docker-cuda/Dockerfile). See [`docker/README.md`](docker/README.md) for build and run instructions (including `linux/amd64` on Apple Silicon).
+A CUDA training image (Python 3.12, **uv**-based install, PyTorch 2.8 + `cu129`, `deepspeed`, `wandb`, released `diffusers`) is defined under [`docker/docker-cuda/`](docker/docker-cuda/Dockerfile). See [`docker/README.md`](docker/README.md) for build and run instructions (including `linux/amd64` on Apple Silicon).
## Experiment Trackers
@@ -240,12 +233,17 @@ SFT and offline DPO use strict JSONL with `schema_version: 2`. Public media obje
```jsonl
{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"targets/poster.png"}]}},"metadata":{}}
{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"pairs/chosen.png"}]},"rejected":{"media":[{"type":"image","path":"pairs/rejected.png"}]}},"metadata":{}}
+{"schema_version":2,"input":{"prompt":"Animate toward this ending.","media":[{"type":"image","path":"conditions/end.png","slot":"last_frame"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/story.mp4","fps":24.0},{"type":"audio","path":"targets/story.wav","sample_rate":32000}]}},"metadata":{}}
```
+The optional input-only `slot` field binds sparse conditions to adapter-declared semantic
+arguments. Unslotted media fills remaining slots positionally; supervision outputs reject slots.
+
Prompt and input-condition encodings are cached. Target, chosen, and rejected media are decoded and
encoded on the fly; their VAE latents are never stored in the preprocessing cache. One offline
epoch is one complete dataloader traversal sharded by PyTorch's official `DistributedSampler`. See the
-[dataset guide](guidance/datasets.md#offline-v2-records) for the full schema and cadence rules.
+[dataset guide](guidance/datasets.md#offline-v2-records) for the full schema and cadence rules, and
+the [GPU validation plan](guidance/gpu_validation.md) for the 120-job model/backend/algorithm matrix.
## Text-to-Image & Text-to-Video
diff --git a/docker/README.md b/docker/README.md
index 1196d8d43..ced5c85de 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -31,12 +31,11 @@ docker run --rm -it --gpus all ghcr.io/x-gengroup/flow-factory:0.1.0
### B. Build locally
-Clone with the `diffusers` submodule (required):
+Clone the repository:
```bash
-git clone --recursive https://github.com/X-GenGroup/Flow-Factory.git
+git clone https://github.com/X-GenGroup/Flow-Factory.git
cd Flow-Factory
-# or, if already cloned: git submodule update --init --recursive
```
Build from the **repository root**:
@@ -103,7 +102,7 @@ Do not commit secrets; use environment variables or your orchestrator's secret s
|---|---|---|
| `nvidia-smi` not found in container | NVIDIA Container Toolkit not installed | [Install the toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) and restart Docker |
| `CUDA out of memory` | Batch size too large for GPU VRAM | Reduce batch size, enable DeepSpeed ZeRO-2, or use FSDP |
-| Build fails on `diffusers` install | Submodule not initialized | Run `git submodule update --init --recursive` |
+| Diffusers API import fails | Installed package is older than 0.40.0 | Rebuild the image or install `diffusers>=0.40.0` |
| Every source change triggers full rebuild | Expected with `COPY . /app` | The Dockerfile uses two-phase COPY for layer caching; ensure `pyproject.toml` is unchanged for cache hits |
## Building tips
diff --git a/docker/docker-cuda/Dockerfile b/docker/docker-cuda/Dockerfile
index 0ca46dbc9..91d63bdbf 100644
--- a/docker/docker-cuda/Dockerfile
+++ b/docker/docker-cuda/Dockerfile
@@ -1,4 +1,4 @@
-# Build from repository root (includes diffusers submodule):
+# Build from repository root:
# docker buildx build --platform linux/amd64 -f docker/docker-cuda/Dockerfile -t flow-factory:local --load .
#
# Apple Silicon: always pass --platform linux/amd64 for this training image.
@@ -56,10 +56,9 @@ RUN uv pip install ".[deepspeed,wandb]"
# --- Phase 2: copy full source and re-link editable installs ---
COPY . /app
RUN uv pip install --no-deps -e .
-RUN uv pip install -e ./diffusers
# Fail fast if any layer is broken
-RUN python -c "import deepspeed, torch, wandb; print('torch', torch.__version__)" && ff-train --help >/dev/null
+RUN python -c "import deepspeed, diffusers, torch, wandb; assert tuple(map(int, diffusers.__version__.split('.')[:2])) >= (0, 40); print('torch', torch.__version__, 'diffusers', diffusers.__version__)" && ff-train --help >/dev/null
# Intentionally root: DeepSpeed / NCCL may need elevated IPC access.
CMD ["/bin/bash"]
diff --git a/examples/README.md b/examples/README.md
index 2dea62db8..2765c7800 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -1,6 +1,6 @@
# Examples
-Training configs for all supported algorithm–model combinations.
+Curated training configs for representative supported algorithm–model combinations.
## Directory Structure
@@ -92,7 +92,9 @@ Its 64x96 canvas is intentionally a correctness geometry. The quality-oriented T
default is now the shared-`vid_prompt`, LoRA-rank-64 baseline aligned with the LTX2
T2AV recipe and uses both CLAP and ImageBind rewards. It is configuration/API
validated; no completed long-run reward trend is claimed. FL2VA and Ref2VA are also
-**Schema/API validated only**, rather than claims of training stability or reward improvement.
+schema/API and local offline-path validated, rather than claims of real-weight training stability
+or reward improvement. The complete follow-up campaign is defined in the
+[GPU validation plan](../guidance/gpu_validation.md).
The T2VA `quality_720p_fsdp2.yaml` recipe is the active native-quality path:
768x1344, 124 frames, 24 denoising steps, LoRA rank 64 / alpha 128, and two
diff --git a/guidance/datasets.md b/guidance/datasets.md
index ff0d74726..7a87075e0 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -73,10 +73,17 @@ Every V2 media object uses `type` as its only public discriminator. The accepted
```json
{"type":"image","path":"images/source.png"}
+{"type":"image","path":"images/end.png","slot":"last_frame"}
{"type":"video","path":"videos/clip.mp4","fps":24.0}
{"type":"audio","path":"audios/clip.wav","sample_rate":48000}
```
+`slot` is an optional, input-only semantic binding. It is useful when position alone is ambiguous:
+an explicitly slotted item reserves that adapter-declared slot, while unslotted items fill the
+remaining slots in declaration order. Duplicate, unknown, and wrong-media-type slots fail during
+contract validation. Supervision outputs reject `slot` because their order is declared by the
+pipeline output contract rather than by condition-argument names.
+
Do not write `kind` in a V2 record. Some ordered-reference adapters still consume a validated
legacy `kind` mapping internally; the V2 condition projection creates that private bridge only at
the adapter preprocessing boundary. It is not part of the public V2 schema.
@@ -137,10 +144,18 @@ V2 input
V2 target, chosen, or rejected
-> decode from the source file in Dataset.__getitem__
-> collate decoded CPU media
- -> adapter.encode_output_state under no-grad on every training microbatch
+ -> adapter.prepare_condition_state once under no-grad on every training microbatch
+ -> adapter.encode_output_state with that prepared state under no-grad
-> clean latent state for the objective
```
+The prepared condition keeps immutable cached condition fields and exposes explicit model-forward
+and output-codec views. Those two views may intentionally reference the same prepared input tensor;
+the framework prevents accidental key overwrites rather than requiring artificial tensor copies.
+SFT uses the same object for target binding and the forward.
+Offline DPO uses one object for both chosen/rejected encodes and both policy/reference forwards, so
+stochastic condition realization cannot drift between preference arms.
+
Target, chosen, and rejected payloads, their VAE latents, and supervision metadata are never
written to the Arrow condition cache. There is no target-VAE preprocessing cache. This avoids a
second large media-derived dataset on disk and keeps output geometry and posterior semantics owned
@@ -151,7 +166,10 @@ During offline dataset construction, every unique normalized input and supervisi
streamed once through SHA-256, with digests memoized only for that source build. These digests are
identity metadata: media payloads, decoded pixels, and output latents are neither copied nor cached.
Replacing an input condition file in place therefore changes its condition identity and invalidates
-the Arrow cache automatically. Replacing target, chosen, or rejected media in place changes the
+the Arrow cache automatically. The checkpoint-realized input projection contract is also part of
+the cache key, so changes to slot order, binding, aggregate rules, negative-prompt policy, or batch
+capability cannot reuse an incompatible Arrow schema. Replacing target, chosen, or rejected media
+in place changes the
full record identity used by exact-resume checks; supervision is still decoded afresh and requires
no target-cache invalidation step.
@@ -203,21 +221,40 @@ output semantics.
|---|---|---|
| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. |
| Supported | `wan2_t2v` | Video targets require `fps`; the codec resamples to configured frames/rate and samples the Wan VAE posterior on the fly. |
-| Supported | `minimax-h3-t2va` | Every candidate is an exact ordered `(video, audio)` pair with required `fps` and `sample_rate`. The codec aligns both streams to configured H3 geometry, samples the video posterior, takes the official audio-posterior mode, and packs structured video/audio rows on the fly. Condition preprocessing and training remain B=1. |
-| Blocked | `wan2_i2v` | Output geometry depends on the first-frame VAE latent/mask, while the current condition cache does not preserve the source pixels needed by that binder. |
-| Blocked | `ltx2_t2av`, `ltx2_i2av` | Their adapter codec still needs exact LTX-specific audio/video duration alignment, latent packing, and decode context; I2AV also needs the pinned first-frame active mask. |
-| Blocked | `minimax-h3-fl2va`, `minimax-h3-ref2va` | Their cached input media still need a shared, reproducible offline condition-prefix binder; offline DPO must reuse the same conditioned prefix noise for both preference arms. |
+| Supported | `wan2_i2v` | Input media binds a required `first_frame` image and an optional `last_frame` image. Condition pixels are cached at configured geometry, then encoded with VAE posterior mode once per batch. Expanded-timestep TI2V checkpoints accept the first frame only because official Diffusers ignores a last image in that mode. Video targets require `fps`. Offline execution is B=1. |
+| Supported | `ltx2_t2av`, `ltx2_i2av` | Every candidate is an exact ordered `(video, audio)` pair with required `fps` and `sample_rate`. Both streams are aligned to the official LTX2 clock and encoded/packed on the fly. I2AV requires the `first_frame` image slot, substitutes its posterior-mode first latent into each target, and excludes the pinned tokens with an active mask. |
+| Supported | `minimax-h3-t2va`, `minimax-h3-fl2va`, `minimax-h3-ref2va` | Every candidate is an exact ordered `(video, audio)` pair. FL2VA accepts `first_frame`, `last_frame`, or both slots; Ref2VA accepts 1-12 globally ordered references and requires at least one image or video. Conditioned workflows realize one official prefix per batch, shared by both offline-DPO candidates and policy/reference forwards. H3 remains B=1. |
+
+Wan first/last semantics use generic semantic slots, not model-specific schema keys. The first
+frame is required; the last frame is optional. Unslotted input remains a positional convenience,
+but an explicit slot is recommended for sparse or generated manifests. The target is the complete
+generated video: its first frame, and its final frame when provided, correspond to the conditions.
+
+```jsonl
+{"schema_version":2,"input":{"prompt":"A paper boat crosses the pond.","media":[{"type":"image","path":"conditions/first.png","slot":"first_frame"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/first-only.mp4","fps":24.0}]}},"metadata":{}}
+{"schema_version":2,"input":{"prompt":"Interpolate the changing sky.","media":[{"type":"image","path":"conditions/first.png","slot":"first_frame"},{"type":"image","path":"conditions/last.png","slot":"last_frame"}]},"supervision":{"type":"preference","chosen":{"media":[{"type":"video","path":"pairs/chosen.mp4","fps":24.0}]},"rejected":{"media":[{"type":"video","path":"pairs/rejected.mp4","fps":24.0}]}},"metadata":{}}
+```
-MiniMax H3 T2VA supervision lists video first and audio second. The target video must cover the
-configured 24-fps duration; it is deterministically sampled onto that frame grid and resized to the
-configured canvas. Audio is converted to stereo at the H3 audio-VAE rate, then trimmed or
-right-padded to the exact aligned latent duration. For example:
+LTX2 and MiniMax H3 supervision list video first and audio second. LTX2 aligns audio duration with
+the official `num_frames / frame_rate` clock. H3 target video must cover the configured 24-fps
+duration; it is deterministically sampled onto that frame grid and resized to the configured
+canvas. H3 audio is converted to stereo at the audio-VAE rate, then trimmed or right-padded to the
+exact aligned latent duration. For example:
```jsonl
{"schema_version":2,"input":{"prompt":"Ocean waves beneath an aurora.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/aurora.mp4","fps":24.0},{"type":"audio","path":"targets/aurora.wav","sample_rate":32000}]}},"metadata":{}}
{"schema_version":2,"input":{"prompt":"Ocean waves beneath an aurora.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"video","path":"pairs/chosen.mp4","fps":24.0},{"type":"audio","path":"pairs/chosen.wav","sample_rate":32000}]},"rejected":{"media":[{"type":"video","path":"pairs/rejected.mp4","fps":24.0},{"type":"audio","path":"pairs/rejected.wav","sample_rate":32000}]}},"metadata":{}}
```
+For LTX2 I2AV, bind one image to `first_frame`. H3 FL2VA supports first-only, last-only, and
+first-plus-last records; use explicit slots for the last-only form. H3 Ref2VA puts the complete
+ordered image/video/audio reference sequence in `input.media`; the offline projection bridges those
+public `type` objects to the adapter's private legacy reference representation.
+
+```jsonl
+{"schema_version":2,"input":{"prompt":"Reveal the scene before this ending.","media":[{"type":"image","path":"conditions/end.png","slot":"last_frame"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/story.mp4","fps":24.0},{"type":"audio","path":"targets/story.wav","sample_rate":32000}]}},"metadata":{}}
+```
+
## Common task formats
The following compact formats remain supported for generation acquisition. They are separate from
@@ -367,6 +404,10 @@ FL2VA uses an ordered `"images"` list:
- Any other cardinality is invalid.
- Order must not be sorted, deduplicated, or inferred from filenames.
+This legacy generation form cannot express last-only conditioning. Strict V2 offline records use
+the generic `slot: "last_frame"` binding shown above, so FL2VA can train on a final frame without a
+synthetic first frame.
+
The example stores paths relative to the dataset root, so its YAML sets `image_dir` to the dataset
directory. See the [FL2VA dataset fixture](../dataset/minimax_h3_fl2va/train.jsonl) and
[FL2VA GRPO configuration](../examples/grpo/lora/minimax_h3_fl2va/default.yaml).
@@ -381,8 +422,8 @@ image, video, and audio entries:
```
Array order is semantically significant. It is preserved during validation, encoding, caching, and
-sample identity hashing. At least one image or video reference is required; an audio-only array is
-invalid.
+sample identity hashing. The array accepts 1-12 entries and requires at least one image or video;
+an audio-only array is invalid.
Supported entries:
@@ -394,7 +435,7 @@ Supported entries:
`fps` and `sample_rate` overrides must be finite positive numbers. A video may use its embedded
soundtrack or a separate dataset-relative `audio_path`; a video `sample_rate` override requires
-`audio_path`. Unknown keys and unsupported `kind` values fail before preprocessing.
+`audio_path`. Unknown keys and unsupported legacy `kind` values fail before preprocessing.
This legacy online manifest is distinct from the strict V2 format above. A V2 record always uses
`input.media[*].type`; offline condition projection performs any required legacy `kind` conversion
@@ -419,8 +460,9 @@ TXT/JSONL row
Prompt encoders, condition VAEs, and processors are preprocessing components. Online RL can
offload them after cache creation because optimization consumes cached conditions. Offline SFT and
-offline DPO reload any output-codec components declared by the adapter and encode target,
-chosen, and rejected media on the fly; those output states are never cached.
+offline DPO reload any condition-preparer and output-codec components declared by the adapter.
+They realize one prepared condition per batch, then encode target, chosen, and rejected media on
+the fly; those output states are never cached.
Ref2VA adds an ordered-reference path:
diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md
new file mode 100644
index 000000000..ede5e295f
--- /dev/null
+++ b/guidance/gpu_validation.md
@@ -0,0 +1,240 @@
+# GPU Validation Plan
+
+This document is the handoff contract for real-weight GPU validation. It does
+not claim that an unexecuted combination is supported. A combination becomes
+validated only after its artifacts satisfy the acceptance criteria below.
+
+## Environment gate
+
+Record the following once for every validation campaign:
+
+- Flow-Factory commit SHA and parent stacked-PR SHA.
+- Python, CUDA, PyTorch, Accelerate, DeepSpeed, and Diffusers versions.
+- `diffusers>=0.40.0`; do not mix results from an older official release.
+- GPU model, GPU count, per-GPU memory, driver, and NCCL versions.
+- Exact model revision and dataset-content hash.
+- `PYTHONPATH` and editable-install source must resolve to the tested worktree.
+
+Run a configuration/import preflight before allocating model weights:
+
+```bash
+python -m compileall -q src/flow_factory
+python -c "import diffusers; assert tuple(map(int, diffusers.__version__.split('.')[:2])) >= (0, 40)"
+```
+
+## Main experiment matrix
+
+The main campaign is the Cartesian product of the ten semantic modes, three
+distributed backends, and four algorithms below: **10 x 3 x 4 = 120 jobs**.
+Do not collapse first-frame and first/last-frame Wan rows: they exercise
+different condition layouts and active masks.
+
+### Semantic modes
+
+| ID | Mode | Representative checkpoint | Required input | Required output |
+|---|---|---|---|---|
+| `sd35-t2i` | SD3.5 text-to-image regression anchor | `stabilityai/stable-diffusion-3.5-medium` | prompt | image |
+| `bagel-mri2i` | Bagel ordered multi-reference-images-to-image regression anchor | `ByteDance-Seed/BAGEL-7B-MoT` | prompt plus exactly two ordered images per sample | image |
+| `wan-t2v` | Wan text-to-video | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | prompt | video |
+| `wan-i2v-first` | Wan first-frame-to-video | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | exactly one first-frame image | video |
+| `wan-flf2v` | Wan first/last-frame-to-video | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | ordered first and last images | video |
+| `ltx2-t2av` | LTX2 text-to-audio-video | `Lightricks/LTX-2` | prompt | ordered video and audio |
+| `ltx2-i2av` | LTX2 image-to-audio-video | `Lightricks/LTX-2` | prompt plus one image | ordered video and audio |
+| `h3-t2va` | MiniMax H3 text-to-video-audio | `MiniMaxAI/MiniMax-H3` | prompt | ordered video and audio |
+| `h3-fl2va` | MiniMax H3 sparse first/last-frame-to-video-audio | `MiniMaxAI/MiniMax-H3` | `first_frame`, `last_frame`, or both image slots | ordered video and audio |
+| `h3-ref2va` | MiniMax H3 ordered-reference-to-video-audio | `MiniMaxAI/MiniMax-H3` | 1-12 ordered heterogeneous references, including image or video | ordered video and audio |
+
+The public V2 schema uses the `type` discriminator and an optional input-only
+`slot`. An explicit slot reserves its adapter-declared semantic argument;
+unslotted media is only a positional shorthand that fills remaining slots in
+declaration order. Wan requires `first_frame` and optionally accepts
+`last_frame`; H3 FL accepts either slot or both. Supervision media does not
+repeat independent condition-image objects. Its video is nevertheless the full
+configured output sequence, whose first and/or last endpoint must correspond to
+the supplied endpoint conditions.
+
+### Backends
+
+| ID | `config_file` | Required observation |
+|---|---|---|
+| `ddp` | `config/accelerate_configs/multi_gpu.yaml` | Every rank executes the same branch and step count. |
+| `zero2` | `config/deepspeed/deepspeed_zero2.yaml` | Optimizer state is partitioned and both policy/reference scopes complete. |
+| `fsdp2` | `config/accelerate_configs/fsdp2.yaml` | Adapter wrap plan, DTensor parameters, checkpointing, and component routing remain valid. |
+
+### Algorithms and stopping rules
+
+| ID | `train.trainer_type` | Exact smoke length | Evaluation |
+|---|---|---|---|
+| `grpo` | `grpo` | two training epochs, one optimizer step per epoch | `eval.eval_freq: 0` |
+| `sft` | `sft` | exactly two rank-local dataloader batches | `eval.eval_freq: 0` |
+| `offline-dpo` | `offline-dpo` | exactly two rank-local dataloader batches | `eval.eval_freq: 0` |
+| `tdm` | `tdm` | two training epochs, one generator/fake update cycle per epoch | `eval.eval_freq: 0` |
+
+For SFT and offline DPO, build a finite dataset whose official
+`DistributedSampler` yields exactly two batches per rank, set
+`gradient_accumulation_steps: 1`, and run one complete dataloader epoch. One
+offline epoch means one complete dataloader traversal. Sampler tail padding is
+standard PyTorch behavior and does not change that definition.
+
+For GRPO and TDM, use `max_epochs: 2`,
+`gradient_step_per_epoch: 1`, the smallest valid group size, and no evaluation.
+Reduce resolution, frame count, and inference steps only within each adapter's
+declared geometry constraints. MiniMax H3 must retain at least five seconds of
+24-fps output even in a low-resolution smoke run.
+
+### Concrete smoke profiles
+
+Use these geometry and sampling overlays unless a real checkpoint rejects the
+reduced geometry. Any fallback must stay valid for the adapter contract and be
+recorded in the resolved YAML; do not silently return to a large quality recipe.
+
+| Mode | Starting recipe | Train geometry | Steps | Algorithm-specific notes |
+|---|---|---|---|---|
+| `sd35-t2i` | `examples/grpo/lora/sd3_5/default.yaml` | `resolution: 256` | `num_inference_steps: 2` | Use the checked-in SD3.5 SFT/offline-DPO/TDM recipe as the algorithm overlay. |
+| `bagel-mri2i` | `examples/grpo/lora/bagel/i2i.yaml` | `resolution: 256` | `num_inference_steps: 2` | Every row has exactly two ordered references; keep `shuffle_samples: false`. |
+| `wan-t2v` | `examples/grpo/lora/wan21/t2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Target video has at least five frames and carries its source `fps`. |
+| `wan-i2v-first` | `examples/grpo/lora/wan22/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use TI2V-5B and exactly one condition image. |
+| `wan-flf2v` | `examples/grpo/lora/wan21/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use two ordered condition images; do not use an expanded-timestep checkpoint. |
+| `ltx2-t2av` | `examples/grpo/lora/ltx2/t2av.yaml` | `resolution: [128, 192]`, `num_frames: 9`, `frame_rate: 24.0` | `num_inference_steps: 2` | AV targets cover the exact 9-frame clock; audio carries `sample_rate`. |
+| `ltx2-i2av` | `examples/grpo/lora/ltx2/i2av.yaml` | `resolution: [128, 192]`, `num_frames: 9`, `frame_rate: 24.0` | `num_inference_steps: 2` | One condition image; verify the first latent frame is inactive in the loss. |
+| `h3-t2va` | `examples/grpo/lora/minimax_h3_t2va/debug.yaml` | `resolution: [64, 96]`, `num_frames: 124`, `frame_rate: 24.0` | `num_inference_steps: 2` | Preserve the released five-second minimum and neutral guidance. |
+| `h3-fl2va` | `examples/grpo/lora/minimax_h3_fl2va/default.yaml` plus H3 debug geometry | `resolution: [64, 96]`, `num_frames: 124`, `frame_rate: 24.0` | `num_inference_steps: 2` | The two offline records are explicit `first_frame`-only and `last_frame`-only cases. Cover both slots together in the additional variant gate. |
+| `h3-ref2va` | `examples/grpo/lora/minimax_h3_ref2va/default.yaml` plus H3 debug geometry | `resolution: [64, 96]`, `num_frames: 124`, `frame_rate: 24.0` | `num_inference_steps: 2` | Preserve heterogeneous global reference order; include image, video, and audio references. |
+
+For GRPO set `group_size: 2`, `unique_sample_num_per_epoch: 1`, and
+`gradient_accumulation_steps: auto`; this avoids a degenerate one-candidate
+advantage while retaining one optimizer step per epoch. For TDM set
+`group_size: 1` and `gradient_accumulation_steps: auto`; with the two-step smoke
+profile the resolved accumulation count must be divisible by two. SFT and
+offline DPO use `gradient_accumulation_steps: 1`. Use at least two distributed
+ranks for every backend; increase the rank count only for checkpoint memory
+capacity, without changing the two iteration/batch stopping rule.
+
+The offline fixture must contain exactly `2 * world_size` records and use
+`per_device_batch_size: 1`, so `DistributedSampler(drop_last=False)` yields two
+non-padded batches on every rank. Demonstration and preference fixtures use the
+same input distribution. Preference rows must use distinct chosen/rejected
+media files with identical geometry. Before launch, validate conditioned target
+endpoints against their inputs: LTX2 I2AV and first-conditioned Wan/H3 targets
+must begin with the supplied first frame within the fixture tolerance; a target
+with a last-frame condition must end with that supplied frame. Store the
+comparison metric and tolerance in the fixture manifest.
+
+The campaign generator should materialize one job ID for every Cartesian-product
+cell as `{mode}__{backend}__{algorithm}` and assert that the set has exactly 120
+unique IDs before submission. A skipped or infrastructure-blocked cell remains
+in the result table with its failure classification; it must not silently reduce
+the matrix.
+
+## Checkpoint-variant coverage
+
+The 120-job main matrix uses one checkpoint per semantic mode. Run the following
+additional variant gate before declaring a family generally supported. The
+minimum gate is DDP plus SFT and offline DPO for two batches each; configuration
+construction and static contract validation must also pass under all three
+backends and all four algorithms.
+
+| Family mode | Additional checkpoint variants |
+|---|---|
+| Wan T2V | `Wan2.1-T2V-14B-Diffusers`, `Wan2.2-TI2V-5B-Diffusers`, `Wan2.2-T2V-A14B-Diffusers` |
+| Wan I2V first-only | `Wan2.1-I2V-14B-480P-Diffusers`, `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
+| Wan first/last | `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
+| LTX2 T2AV and I2AV | `Lightricks/LTX-2.3` or its canonical official-Diffusers repository revision |
+| MiniMax H3 | The same checkpoint is covered separately by T2VA, FL2VA, and Ref2VA inputs. Add an FL2VA first-plus-last fixture to complement the first-only/last-only main jobs. |
+
+Wan2.2 TI2V-5B uses expanded timesteps. Official Diffusers ignores a supplied
+last image in that mode, so its effective input contract is first-frame only;
+do not count it as a first/last checkpoint. Other Wan I2V checkpoints must prove
+both first-only and first/last execution.
+
+For the Wan2.2 A14B dual-transformer gate, force or instrument two offline
+timestep samples so that one routes below the transformer boundary and one
+routes at or above it. Record both sampled timesteps and the selected component.
+Two unconstrained random batches are not sufficient evidence that both intended
+trainable transformers updated.
+
+## Per-job overrides
+
+Start from the closest checked-in example and apply a reviewable overlay. Every
+generated job must make these values explicit:
+
+```yaml
+log:
+ logging_backend: none
+ save_freq: 0
+
+train:
+ per_device_batch_size: 1
+ ema_decay: 0
+ enable_gradient_checkpointing: true
+ seed: 42
+ max_epochs: 2 # Use 1 for the finite two-batch offline fixture.
+
+eval:
+ eval_freq: 0
+```
+
+Offline jobs additionally require a homogeneous V2 manifest, unit source
+weights, on-the-fly target media encoding, and `num_train_timesteps: 1`.
+Offline DPO chosen/rejected media must share one prepared condition state, the
+same component times, and the same diffusion noise. Do not add a target-latent
+preprocessing cache.
+
+GRPO jobs should use the smallest reward model that still exercises reward
+routing. TDM is reward-free and needs both generator and fake optimizer
+entries. Disable W&B and checkpoint saving for the smoke campaign.
+
+Bagel workers require the project `bagel` extra in addition to the selected
+distributed backend, including a compatible `flash-attn` build and OpenCV. For
+example, a DeepSpeed worker should install `.[deepspeed,bagel]`. Record the
+installed optional reward extras in the environment manifest as well; a missing
+Bagel or reward dependency is an `import` failure, not a model-support result.
+
+Keep the two-reference Bagel cardinality identical on every rank in the FSDP2
+main-matrix job. Ragged per-rank reference-round counts are a separate stress
+case because a sharded language model must execute the same number of parameter
+all-gathers on every rank.
+
+## Acceptance criteria
+
+A job passes only when all of the following are captured:
+
+1. Configuration parsing, registry lookup, and all imports succeed from the
+ tested worktree.
+2. Model loading preserves the adapter's component-dtype manifest and produces
+ no unexpected FP32/BF16/FP16 coercion.
+3. Every rank reports identical branch routing, batch count, optimizer-step
+ count, and terminal status; no rank hangs at a collective.
+4. Both iterations/batches produce finite loss, finite gradient norm, and a
+ confirmed parameter update for every intended trainable component.
+5. SFT target encoding runs on the fly. Offline DPO reuses one prepared input
+ realization for chosen and rejected outputs and keeps paired noise/time
+ coupling exact.
+6. Conditioned video/audio jobs retain their pinned condition regions through
+ noising and forward loss masks. Unconditioned target regions remain active.
+7. Joint audio-video jobs report both components, aligned durations, valid
+ component-specific scheduler times, and the intended objective reduction.
+8. GRPO replay likelihood and TDM replay checks satisfy the configured
+ tolerances. Condition media and derived prefixes are retained in every
+ conditioned mode. TDM may intentionally skip decoding generated output media
+ when its algorithm contract does not consume it.
+9. Peak allocated/reserved GPU memory and wall time are recorded for capacity
+ planning.
+10. The command, resolved YAML, stdout/stderr log, environment manifest, and a
+ compact metrics JSON are attached to the job result.
+
+## Failure classification
+
+Classify failures before retrying:
+
+- `contract`: schema, media cardinality/order, geometry, or capability mismatch.
+- `import`: dependency, registry, optional component, or worktree-resolution error.
+- `model`: official pipeline/adapter semantic mismatch.
+- `algorithm`: loss, replay, reference-policy, or objective coupling error.
+- `backend`: DDP/DeepSpeed/FSDP2 wrapping, collective, dtype, or checkpoint issue.
+- `capacity`: reproducible OOM after the allowed geometry reductions.
+- `infrastructure`: download, permission, driver, filesystem, or cluster failure.
+
+Do not mark `capacity` or `infrastructure` failures as unsupported model/algorithm
+combinations. Preserve the first failing artifact and link any follow-up fix to
+the exact job ID `{mode}__{backend}__{algorithm}`.
diff --git a/guidance/new_model.md b/guidance/new_model.md
index 025686559..71367702d 100644
--- a/guidance/new_model.md
+++ b/guidance/new_model.md
@@ -218,7 +218,12 @@ preprocess_func(prompt, images, videos, audios, **kwargs):
return results
```
-Text-to-image models override only `encode_prompt` and `encode_image`; image-to-video models add `encode_video`; audio-conditioned models add `encode_audio`. There is no need to add stub `pass` overrides for unused modalities — `BaseAdapter` already provides them.
+Text-to-image, text-to-video, and text-to-audio-video adapters usually override only
+`encode_prompt` because they have no condition media. Image-conditioned tasks add
+`encode_image`; video-conditioned tasks add `encode_video`; audio-conditioned tasks add
+`encode_audio`. These functions encode inputs, not supervised outputs—offline targets belong to
+the output-state codec. There is no need to add stub `pass` overrides for unused modalities;
+`BaseAdapter` already provides them.
#### `encode_prompt`
@@ -554,17 +559,29 @@ Online-only adapters can keep the default `build_output_state_codec() -> None`.
offline DPO, an adapter must additionally declare both sides of its pipeline and provide an
on-the-fly output codec:
-1. Set a class-level `pipeline_io_contract`. It owns input media counts/order/binding, negative
- prompt policy, the exact ordered output media sequence, rate requirements, geometry source, and
- batch capability.
-2. Override `build_output_state_codec()` with a declaration-only codec. Its
+1. Set a class-level `pipeline_io_contract`. It owns per-type and aggregate input counts,
+ order/binding, optional semantic input slots, negative prompt policy, the exact ordered output
+ media sequence, rate requirements, geometry source, and batch capability. Explicit V2 slots
+ reserve declared arguments; unslotted inputs fill the remaining slots in declaration order, and
+ output media must never carry slots. If checkpoints behind one adapter expose narrower behavior, override
+ `_resolve_pipeline_io_contract()` and return an immutable instance-specific specialization;
+ offline data validation consumes `effective_pipeline_io_contract`.
+2. If cached input fields are not already the exact forward condition, override
+ `build_condition_state_preparer()` with a declaration-only preparer. Its
+ `required_components` lists runtime encoders and `prepare_condition_state()` returns one
+ `PreparedConditionState` per batch. Put input-owned model fields in `forward_context` and
+ input-owned target-binding fields in `output_context`. The two runtime consumer views may
+ intentionally share an input-owned tensor (for example a mask or layout), but each merged
+ consumer view must remain collision-free with cached fields and later candidate-owned output
+ fields.
+3. Override `build_output_state_codec()` with a declaration-only codec. Its
`required_components` names logical runtime components such as `("vae",)`; construction must
not load, materialize, move, replace, or cast them.
-3. Return an `EncodedOutputState` containing a detached `LatentState`, output-derived forward and
+4. Return an `EncodedOutputState` containing a detached `LatentState`, output-derived forward and
decode contexts, and one exact geometry signature per sample.
-4. Override `_validate_encoded_output_geometry()` so configured, condition-derived, and
+5. Override `_validate_encoded_output_geometry()` so configured, condition-derived, and
output-derived dimensions cannot drift silently.
-5. Declare a complete immutable `offline_training_forward_overrides` mapping whenever the base
+6. Declare a complete immutable `offline_training_forward_overrides` mapping whenever the base
`{"guidance_scale": 1.0}` contract does not describe the adapter. Offline trainers apply this
mapping after sampling configuration and cached batch conditions, so it owns loss-time model
conditioning. Conventional CFG branches must all be set to their neutral point (for example,
@@ -573,10 +590,11 @@ on-the-fly output codec:
the complete mapping so permissive `**kwargs` forwards do not receive unrelated base keys.
The dataset remains responsible only for strict V2 parsing and CPU media decoding. The adapter
-owns numerical output semantics. The SFT/offline-DPO trainer calls `encode_output_state()` on every
-microbatch under `torch.no_grad`; target, chosen, and rejected latents are not preprocessing-cache
-columns. The declared output components are loaded through `ModelLoadCoordinator`, never from
-inside the codec.
+owns numerical condition/output semantics. The SFT/offline-DPO trainer first calls
+`prepare_condition_state()` once, then calls `encode_output_state()` under `torch.no_grad`; offline
+DPO passes the same prepared object to both preference candidates. Target, chosen, and rejected
+latents are not preprocessing-cache columns. Declared condition and output components are loaded
+through `ModelLoadCoordinator`, never from inside a preparer or codec.
Condition encoding and target encoding should share role-neutral numerical transforms instead of
duplicating VAE math. Extract helpers for pixel preprocessing, posterior extraction, latent
@@ -605,6 +623,11 @@ multi-component state order, active masks, rate alignment, or forward context th
encoding does not own. If those semantics are not lossless, set a concrete
`output_state_codec_unavailable_reason` so offline selection fails before downloading weights.
+Multi-modal objectives may need a reduction different from trajectory likelihoods. Override the
+protected `_reduce_flow_matching_objective_values()` hook only for that objective. Do not change
+`reduce_latent_values()` merely to implement SFT: online policy gradients, replay log-probability,
+and distillation continue to rely on their established trajectory-wide reduction.
+
SenseNova is an example of an important boundary: its existing condition schema uses grouped
`images` with within-type order. Do not advertise heterogeneous ordered references merely because
several images are accepted. The public V2 discriminator remains `type`; conversion to a legacy
@@ -901,7 +924,10 @@ Before submitting a new model adapter, verify:
- [ ] **`inference()`** — Accepts both raw and pre-encoded inputs; returns `List[Sample]`
- [ ] **`forward()`** — Single denoising step; ends with `self.scheduler.step()`; returns `SDESchedulerOutput`
- [ ] **Pipeline I/O contract** — Declares exact input/output media, rate, geometry, and batch semantics before enabling offline training
+- [ ] **Effective checkpoint contract (when needed)** — `_resolve_pipeline_io_contract()` narrows a class-level superset without changing public dataset or algorithm code
+- [ ] **Condition-state preparer (when needed)** — Declaration-only logical component requirements; one input realization reused by every candidate/forward in the batch
- [ ] **Output-state codec (when supported)** — Declaration-only logical component requirements; on-the-fly detached target encoding; exact geometry validation
+- [ ] **Objective reduction (when specialized)** — Override only the offline flow-matching hook; online trajectory reduction remains unchanged
- [ ] **Offline forward overrides (when supported)** — Complete immutable adapter mapping; sampling controls never define offline loss semantics; every CFG branch is neutralized or every distilled guidance condition is explicitly pinned
- [ ] **Role-neutral encoder math** — Condition/output paths reuse transforms but explicitly preserve official posterior `sample` versus `argmax` policy and generator routing
- [ ] **Explicit offline blocker (when unsupported)** — `output_state_codec_unavailable_reason` names the missing lossless semantic boundary
diff --git a/guidance/workflow.md b/guidance/workflow.md
index 7af9c133b..491eb3d14 100644
--- a/guidance/workflow.md
+++ b/guidance/workflow.md
@@ -139,8 +139,8 @@ def preprocess_func(self, prompt, images, ...):
- **No HF default-cache copy**: Because each `map()` call sets `cache_file_name`, HuggingFace does **not** also write a duplicate `cache-*.arrow` under `~/.cache/huggingface/datasets/...`.
- **Intelligent caching**: A hash fingerprint of `(dataset, split, max_dataset_size, preprocess_func source, preprocess_kwargs, extra_hash_strs)` (the last includes `model_type` and `model_name_or_path`) determines the cache path. Subsequent runs that match the fingerprint take the fast path without any `Dataset.map` invocation.
- **Component offloading**: Text and condition encoders can be offloaded after cache creation. An
- offline adapter's declared output codec components are reloaded through the component lifecycle
- for on-the-fly target encoding.
+ offline adapter's declared runtime condition-preparer and output-codec components are reloaded
+ through the component lifecycle for on-the-fly condition realization and target encoding.
- **No target cache**: SFT targets and offline-DPO chosen/rejected candidates are decoded per
dataset access and encoded per training microbatch. Target payloads, latent states, and metadata
never enter the Arrow condition cache.
@@ -171,10 +171,15 @@ sampler.set_epoch(data_epoch)
for batch in dataloader:
condition = cached prompt/input tensors
output = freshly decoded target or chosen/rejected media
- trainer.optimize_batch(batch)
+ trainer.optimize_batch(batch) # prepares condition once inside the microbatch
data_epoch += 1 # only after clean exhaustion
```
+Inside `optimize_batch`, the trainer calls
+`adapter.prepare_condition_state(condition)` exactly once, binds every target
+candidate through that prepared object, and reuses its model-forward view for all
+policy/reference passes. The driver must not prepare it a second time.
+
One complete dataloader traversal is one offline epoch. Source weights must be `1`,
`data.sampler_type` remains `auto`, and `gradient_accumulation_steps` is an explicit integer. The
rank-local batch count must be divisible by gradient accumulation; the framework does not add
@@ -202,15 +207,26 @@ evaluation adapters or rewards may consume global RNG before the next training a
cannot currently save exact state because Accelerate does not serialize MPS RNG; use
`log.save_model_only: true` on Apple Silicon.
+The adapter prepares an input-owned condition state exactly once per offline batch. An identity
+preparer returns cached fields unchanged. A conditioned model may instead realize geometry-bound
+VAE latents, masks, or stochastic prefixes and split them into model-forward and output-codec
+contexts. SFT reuses that state for its target; offline DPO reuses the same object for chosen and
+rejected encoding and for both policy/reference arms. Candidate-specific output context is bound
+only after this input state exists, so neither candidate can accidentally own or redraw an input
+condition.
+
The target codec is adapter-owned but role-neutral at its numerical core. Condition and output
encoding reuse the same pixel preprocessing, VAE transform, normalization, and packing helpers.
Their semantic policies remain explicit: official condition paths commonly use posterior
-`argmax`, while stochastic target training uses posterior `sample` with an optional generator.
-Sharing a transform must never silently erase that role boundary.
-
-For offline DPO, chosen and rejected arms share the primary timestep, component-time mapping, and
-diffusion noise. Both policy arms run before one frozen-reference scope covers both reference
-forwards. SFT has no reference branch.
+`argmax`, while stochastic target training may use posterior `sample` with an optional generator.
+Sharing a transform must never silently erase that role boundary. Target media remains on demand;
+the prepared-condition boundary does not introduce a target-pixel or target-latent cache.
+
+For offline DPO, chosen and rejected arms share one prepared input state, the primary timestep,
+component-time mapping, and diffusion noise. Both policy arms run before one frozen-reference
+scope covers both reference forwards. SFT has no reference branch. Multi-component adapters may
+specialize only the offline flow-matching objective reduction (for example, a sum of per-modality
+means) without changing the trajectory-wide reducer used by online policy likelihoods.
## Stage 2: K-Repeat Sampling
diff --git a/src/flow_factory/contracts/__init__.py b/src/flow_factory/contracts/__init__.py
index 1da8d56d5..a77326a04 100644
--- a/src/flow_factory/contracts/__init__.py
+++ b/src/flow_factory/contracts/__init__.py
@@ -47,6 +47,7 @@
OutputMediaSequence,
PipelineIOContract,
RateRequirement,
+ resolve_pipeline_input_media_slots,
validate_pipeline_model_input,
validate_pipeline_output_candidate,
)
@@ -78,6 +79,7 @@
"OutputMediaSequence",
"PipelineIOContract",
"RateRequirement",
+ "resolve_pipeline_input_media_slots",
"ROLLOUT_STORAGE_KEYS",
"TRAINER_METADATA_KEYS",
"validate_pipeline_model_input",
diff --git a/src/flow_factory/contracts/pipeline_io.py b/src/flow_factory/contracts/pipeline_io.py
index bbc15316e..7e0006148 100644
--- a/src/flow_factory/contracts/pipeline_io.py
+++ b/src/flow_factory/contracts/pipeline_io.py
@@ -198,16 +198,20 @@ def __post_init__(self) -> None:
@dataclass(frozen=True, slots=True)
class InputMediaRule:
- """Declare the accepted count for one input media format."""
+ """Declare count and optional semantic slots for one input media format."""
format: MediaFormat
min_count: int
max_count: int | None
+ slots: tuple[str, ...] = ()
+ required_slots: tuple[str, ...] = ()
def __post_init__(self) -> None:
"""Validate strict cardinality types and bounds."""
_require_instance(self.format, MediaFormat, "format")
_require_non_negative_int(self.min_count, "min_count")
+ _require_string_tuple(self.slots, "slots")
+ _require_string_tuple(self.required_slots, "required_slots")
if self.max_count is not None:
_require_non_negative_int(self.max_count, "max_count")
if self.max_count == 0:
@@ -217,6 +221,37 @@ def __post_init__(self) -> None:
f"expected max_count >= min_count, received "
f"min_count={self.min_count} and max_count={self.max_count}"
)
+ if len(set(self.slots)) != len(self.slots):
+ raise ValueError("input media slots must be unique within one media rule")
+ if len(set(self.required_slots)) != len(self.required_slots):
+ raise ValueError("required input media slots must be unique")
+ unknown_required_slots = tuple(
+ slot for slot in self.required_slots if slot not in self.slots
+ )
+ if unknown_required_slots:
+ raise ValueError(
+ "required input media slots must be declared in slots; "
+ f"unknown={unknown_required_slots!r}"
+ )
+ canonical_required_slots = tuple(slot for slot in self.slots if slot in self.required_slots)
+ if self.required_slots != canonical_required_slots:
+ raise ValueError(
+ "required input media slots must use declared slot order "
+ f"{canonical_required_slots!r}, received {self.required_slots!r}"
+ )
+ if self.slots:
+ if self.max_count != len(self.slots):
+ raise ValueError(
+ "a slotted input media rule requires max_count to equal the number "
+ f"of slots, received max_count={self.max_count!r}, slots={self.slots!r}"
+ )
+ if len(self.required_slots) > self.min_count:
+ raise ValueError(
+ "required slot count cannot exceed min_count, received "
+ f"required_slots={self.required_slots!r}, min_count={self.min_count}"
+ )
+ elif self.required_slots:
+ raise ValueError("required input media slots cannot be declared without slots")
@dataclass(frozen=True, slots=True)
@@ -226,12 +261,44 @@ class InputMediaSpec:
rules: tuple[InputMediaRule, ...]
binding: InputMediaBinding
order: InputMediaOrder
+ min_total_count: int | None = None
+ max_total_count: int | None = None
+ required_any_types: tuple[MediaType, ...] = ()
def __post_init__(self) -> None:
"""Validate a canonical and coherent input-media declaration."""
_require_tuple(self.rules, InputMediaRule, "rules")
_require_enum(self.binding, InputMediaBinding, "binding")
_require_enum(self.order, InputMediaOrder, "order")
+ if self.min_total_count is not None:
+ _require_non_negative_int(self.min_total_count, "min_total_count")
+ if self.max_total_count is not None:
+ _require_non_negative_int(self.max_total_count, "max_total_count")
+ if self.max_total_count == 0:
+ raise ValueError(
+ "max_total_count=0 is not canonical; omit all input media rules instead"
+ )
+ if (
+ self.min_total_count is not None
+ and self.max_total_count is not None
+ and self.max_total_count < self.min_total_count
+ ):
+ raise ValueError(
+ "expected max_total_count >= min_total_count, received "
+ f"min_total_count={self.min_total_count} and "
+ f"max_total_count={self.max_total_count}"
+ )
+ _require_enum_tuple(self.required_any_types, MediaType, "required_any_types")
+ if len(set(self.required_any_types)) != len(self.required_any_types):
+ raise ValueError("required_any_types must contain each media type at most once")
+ canonical_required_any_types = tuple(
+ media_type for media_type in MediaType if media_type in self.required_any_types
+ )
+ if self.required_any_types != canonical_required_any_types:
+ raise ValueError(
+ "required_any_types must use canonical type order "
+ f"{canonical_required_any_types}, received {self.required_any_types}"
+ )
media_types = tuple(rule.format.type for rule in self.rules)
if len(set(media_types)) != len(media_types):
@@ -244,7 +311,59 @@ def __post_init__(self) -> None:
"input media rules must use canonical type order "
f"{canonical_media_types}, received {media_types}"
)
+ declared_slots = tuple(slot for rule in self.rules for slot in rule.slots)
+ if len(set(declared_slots)) != len(declared_slots):
+ raise ValueError("input media slot names must be unique across media rules")
+ if declared_slots and self.binding is not InputMediaBinding.GROUPED_BY_TYPE:
+ raise ValueError("semantic input media slots require grouped_by_type binding")
+ if (
+ any(len(rule.slots) > 1 for rule in self.rules)
+ and self.order is not InputMediaOrder.WITHIN_TYPE
+ ):
+ raise ValueError(
+ "multi-slot input media rules require within_type ordering because "
+ "unslotted media uses positional fallback"
+ )
+ unknown_required_types = tuple(
+ media_type for media_type in self.required_any_types if media_type not in media_types
+ )
+ if unknown_required_types:
+ raise ValueError(
+ "required_any_types must be declared by input media rules; "
+ f"unknown={unknown_required_types!r}"
+ )
+ minimum_from_rules = sum(rule.min_count for rule in self.rules)
+ if (
+ self.rules
+ and self.max_total_count is not None
+ and self.max_total_count < minimum_from_rules
+ ):
+ raise ValueError(
+ "max_total_count cannot be smaller than the sum of per-type minimums, "
+ f"received max_total_count={self.max_total_count}, "
+ f"per_type_minimum={minimum_from_rules}"
+ )
+ if (
+ self.rules
+ and self.min_total_count is not None
+ and all(rule.max_count is not None for rule in self.rules)
+ ):
+ maximum_from_rules = sum(
+ rule.max_count for rule in self.rules if rule.max_count is not None
+ )
+ if self.min_total_count > maximum_from_rules:
+ raise ValueError(
+ "min_total_count cannot exceed the sum of finite per-type maximums, "
+ f"received min_total_count={self.min_total_count}, "
+ f"per_type_maximum={maximum_from_rules}"
+ )
if not self.rules:
+ if (
+ self.min_total_count is not None
+ or self.max_total_count is not None
+ or self.required_any_types
+ ):
+ raise ValueError("media-free inputs cannot declare aggregate media constraints")
if self.binding is not InputMediaBinding.GROUPED_BY_TYPE:
raise ValueError("media-free inputs must use grouped_by_type binding")
if self.order is not InputMediaOrder.INSENSITIVE:
@@ -287,11 +406,15 @@ def __post_init__(self) -> None:
_require_instance(self.output_media, OutputMediaSequence, "output_media")
_require_enum(self.geometry_source, GeometrySource, "geometry_source")
_require_enum(self.batch_capability, BatchCapability, "batch_capability")
- if self.geometry_source is GeometrySource.INPUT_MEDIA and not any(
- rule.min_count > 0 for rule in self.input_media.rules
- ):
+ guarantees_input_media = (
+ any(rule.min_count > 0 for rule in self.input_media.rules)
+ or bool(self.input_media.min_total_count)
+ or bool(self.input_media.required_any_types)
+ )
+ if self.geometry_source is GeometrySource.INPUT_MEDIA and not guarantees_input_media:
raise ValueError(
- "input_media geometry requires at least one input media rule with min_count > 0"
+ "input_media geometry requires input constraints that guarantee at least "
+ "one media item"
)
@@ -385,6 +508,98 @@ def validate_pipeline_model_input(
f"received {count}"
)
+ total_count = len(media)
+ input_media = contract.input_media
+ if input_media.min_total_count is not None and total_count < input_media.min_total_count:
+ raise ValueError(
+ f"pipeline requires at least {input_media.min_total_count} input media item(s) "
+ f"in total, received {total_count}"
+ )
+ if input_media.max_total_count is not None and total_count > input_media.max_total_count:
+ raise ValueError(
+ f"pipeline accepts at most {input_media.max_total_count} input media item(s) "
+ f"in total, received {total_count}"
+ )
+ if input_media.required_any_types and not any(
+ counts[media_type.value] > 0 for media_type in input_media.required_any_types
+ ):
+ accepted = tuple(media_type.value for media_type in input_media.required_any_types)
+ raise ValueError(
+ "pipeline requires at least one input media item whose type is in " f"{accepted!r}"
+ )
+ _resolve_pipeline_input_media_slots_unchecked(media, contract)
+
+
+def resolve_pipeline_input_media_slots(
+ model_input: ModelInputLike,
+ contract: PipelineIOContract,
+) -> tuple[str | None, ...]:
+ """Resolve explicit and positional input media onto adapter-declared slots.
+
+ Explicit slot values reserve their named positions first. Unslotted media
+ then fill the remaining positions in declaration order, preserving the V2
+ positional shorthand while making sparse bindings such as a last-frame-only
+ request unambiguous.
+ """
+ validate_pipeline_model_input(model_input, contract)
+ return _resolve_pipeline_input_media_slots_unchecked(model_input.media, contract)
+
+
+def _resolve_pipeline_input_media_slots_unchecked(
+ media: tuple[InputMediaLike, ...],
+ contract: PipelineIOContract,
+) -> tuple[str | None, ...]:
+ assignments: list[str | None] = [None] * len(media)
+ rules_by_type = {rule.format.type.value: rule for rule in contract.input_media.rules}
+ for media_type, rule in rules_by_type.items():
+ indices = [index for index, item in enumerate(media) if item.type == media_type]
+ if not rule.slots:
+ for index in indices:
+ slot = getattr(media[index], "slot", None)
+ if slot is not None:
+ raise ValueError(
+ f"pipeline input media[{index}] declares slot={slot!r}, but "
+ f"media type {media_type!r} has no semantic slots"
+ )
+ continue
+
+ claimed: dict[str, int] = {}
+ unassigned_indices = []
+ for index in indices:
+ slot = getattr(media[index], "slot", None)
+ if slot is None:
+ unassigned_indices.append(index)
+ continue
+ if type(slot) is not str:
+ raise TypeError(
+ f"expected model_input.media[{index}].slot to be str or None, "
+ f"received {type(slot).__name__}: {slot!r}"
+ )
+ if slot not in rule.slots:
+ raise ValueError(
+ f"pipeline input media[{index}] slot {slot!r} is not accepted for "
+ f"media type {media_type!r}; accepted slots={rule.slots!r}"
+ )
+ if slot in claimed:
+ raise ValueError(
+ f"pipeline input media slot {slot!r} is assigned more than once at "
+ f"indices {claimed[slot]} and {index}"
+ )
+ assignments[index] = slot
+ claimed[slot] = index
+
+ remaining_slots = [slot for slot in rule.slots if slot not in claimed]
+ for index, slot in zip(unassigned_indices, remaining_slots):
+ assignments[index] = slot
+ claimed[slot] = index
+ missing_required_slots = tuple(slot for slot in rule.required_slots if slot not in claimed)
+ if missing_required_slots:
+ raise ValueError(
+ f"pipeline requires input media slots {missing_required_slots!r} for "
+ f"media type {media_type!r}"
+ )
+ return tuple(assignments)
+
def validate_pipeline_output_candidate(
media: tuple[OutputMediaLike, ...],
@@ -425,6 +640,11 @@ def validate_pipeline_output_candidate(
f"expected output candidate media[{index}] to implement OutputMediaLike, "
f"received {type(item).__name__}: {item!r}"
)
+ slot = getattr(item, "slot", None)
+ if slot is not None:
+ raise ValueError(
+ f"output candidate media[{index}] cannot declare input-only slot={slot!r}"
+ )
media_type = item.type
if type(media_type) is not str:
raise TypeError(
@@ -526,6 +746,34 @@ def _require_tuple(value: object, item_type: type[object], field_name: str) -> N
_require_instance(item, item_type, f"{field_name}[{index}]")
+def _require_string_tuple(value: object, field_name: str) -> None:
+ if type(value) is not tuple:
+ raise TypeError(
+ f"expected {field_name} to be tuple, received {type(value).__name__}: {value!r}"
+ )
+ for index, item in enumerate(value):
+ if type(item) is not str:
+ raise TypeError(
+ f"expected {field_name}[{index}] to be str, received "
+ f"{type(item).__name__}: {item!r}"
+ )
+ if not item.strip():
+ raise ValueError(f"expected {field_name}[{index}] to be a non-empty string")
+
+
+def _require_enum_tuple(
+ value: object,
+ enum_type: type[Enum],
+ field_name: str,
+) -> None:
+ if type(value) is not tuple:
+ raise TypeError(
+ f"expected {field_name} to be tuple, received {type(value).__name__}: {value!r}"
+ )
+ for index, item in enumerate(value):
+ _require_enum(item, enum_type, f"{field_name}[{index}]")
+
+
def _require_non_negative_int(value: object, field_name: str) -> None:
if type(value) is not int:
raise TypeError(
@@ -552,6 +800,7 @@ def _require_non_negative_int(value: object, field_name: str) -> None:
"OutputMediaLike",
"PipelineIOContract",
"RateRequirement",
+ "resolve_pipeline_input_media_slots",
"validate_pipeline_model_input",
"validate_pipeline_output_candidate",
]
diff --git a/src/flow_factory/data_utils/dataset.py b/src/flow_factory/data_utils/dataset.py
index 2a145cfd4..f7af6eccd 100644
--- a/src/flow_factory/data_utils/dataset.py
+++ b/src/flow_factory/data_utils/dataset.py
@@ -350,6 +350,21 @@ def _preprocess_dataset(
source_hash_override=source_hash_override,
)
+ # Every distributed rank must write the same Arrow schema. Infer it from
+ # the unsharded source before selecting the rank-local rows; otherwise an
+ # all-empty optional-media shard can become ``List(null)`` while another
+ # rank writes ``List(Image)`` and the merged cache cannot be loaded.
+ preprocess_features = self._infer_cross_chunk_features(
+ raw_dataset=raw_dataset,
+ preprocessing_batch_size=preprocessing_batch_size,
+ image_dir=self.image_dir,
+ video_dir=self.video_dir,
+ audio_dir=self.audio_dir,
+ force_reprocess=force_reprocess,
+ target_arrow_path=target_arrow_path,
+ require_explicit_features=bool(self.num_shards and self.num_shards > 1),
+ )
+
if self.num_shards and self.num_shards > 1:
if self.shard_index is None:
raise ValueError(
@@ -375,16 +390,6 @@ def _preprocess_dataset(
if target_arrow_path is not None:
os.makedirs(os.path.dirname(os.path.abspath(target_arrow_path)), exist_ok=True)
- preprocess_features = self._infer_cross_chunk_features(
- raw_dataset=raw_dataset,
- preprocessing_batch_size=preprocessing_batch_size,
- image_dir=self.image_dir,
- video_dir=self.video_dir,
- audio_dir=self.audio_dir,
- force_reprocess=force_reprocess,
- target_arrow_path=target_arrow_path,
- )
-
processed_dataset = raw_dataset.map(
self._preprocess_batch,
batched=True,
@@ -417,6 +422,7 @@ def _infer_cross_chunk_features(
audio_dir: Optional[str],
force_reprocess: bool,
target_arrow_path: Optional[str],
+ require_explicit_features: bool = False,
) -> Optional[HFFeatures]:
"""Infer one explicit schema when later map chunks introduce typed values.
@@ -429,13 +435,16 @@ def _infer_cross_chunk_features(
adapter-owned mutable state is intentionally outside that guarantee.
Args:
- raw_dataset: Input dataset after any distributed sharding.
+ raw_dataset: Full input dataset before distributed sharding.
preprocessing_batch_size: Map chunk size.
image_dir: Image root forwarded to preprocessing.
video_dir: Video root forwarded to preprocessing.
audio_dir: Audio root forwarded to preprocessing.
force_reprocess: Whether an existing explicit Arrow target is rebuilt.
target_arrow_path: Optional explicit Arrow cache target.
+ require_explicit_features: Whether to infer a schema even when the
+ first source chunk is already representative. Distributed ranks
+ use this to share one global schema across disjoint shards.
Returns:
Explicit output features for a cross-chunk nullable transition, or
@@ -453,7 +462,9 @@ def _infer_cross_chunk_features(
preprocessing_batch_size=preprocessing_batch_size,
)
if probe_batches is None:
- return None
+ if not require_explicit_features or not len(raw_dataset):
+ return None
+ probe_batches = [list(range(min(preprocessing_batch_size, len(raw_dataset))))]
explicit_generators = list(_iter_torch_generators(self._preprocess_kwargs))
generator_states = [generator.get_state() for generator in explicit_generators]
@@ -581,7 +592,16 @@ def _preprocess_batch(
f"received B={len(batch['prompt'])} for split={self.split!r}"
)
# The columns that are used in preprocess and maintained in the final results.
- PREPROCESS_COLUMNS = ("prompt", "negative_prompt", "images", "videos", "audios")
+ PREPROCESS_COLUMNS = (
+ "prompt",
+ "negative_prompt",
+ "images",
+ "videos",
+ "audios",
+ "image_slots",
+ "video_slots",
+ "audio_slots",
+ )
metadata_excluded_columns = set(PREPROCESS_COLUMNS)
metadata_excluded_columns.update(self._passthrough_columns)
if self._uses_ordered_references:
@@ -703,6 +723,12 @@ def _preprocess_batch(
reference_args["references"] = loaded_reference_batch
batch["reference_manifest"] = canonical_manifests
+ slot_args = {
+ column: batch[column]
+ for column in ("image_slots", "video_slots", "audio_slots")
+ if column in batch
+ }
+
# 5. Call preprocess function with filtered kwargs
input_args = {
**prompt_args,
@@ -710,6 +736,7 @@ def _preprocess_batch(
**video_args,
**audio_args,
**reference_args,
+ **slot_args,
**self._preprocess_kwargs,
}
filtered_args = filter_kwargs(self._preprocess_func, **input_args)
diff --git a/src/flow_factory/data_utils/offline_condition_cache.py b/src/flow_factory/data_utils/offline_condition_cache.py
index 4d0ea3dcf..94b85d4e6 100644
--- a/src/flow_factory/data_utils/offline_condition_cache.py
+++ b/src/flow_factory/data_utils/offline_condition_cache.py
@@ -30,6 +30,13 @@
from datasets import Dataset as HFDataset
+from ..contracts import (
+ BatchCapability,
+ InputMediaBinding,
+ NegativePromptPolicy,
+ PipelineIOContract,
+ resolve_pipeline_input_media_slots,
+)
from ..samples.references import canonicalize_reference_manifest
from .dataset import (
METADATA_COLUMN,
@@ -43,7 +50,7 @@
)
from .schema import MediaAsset, NormalizedDatasetRecord
-_CONDITION_SOURCE_FORMAT = "flow-factory-offline-condition-v1"
+_CONDITION_SOURCE_FORMAT = "flow-factory-offline-condition-v3"
def project_offline_condition_dataset(
@@ -51,6 +58,7 @@ def project_offline_condition_dataset(
*,
source_name: str,
ordered_references: bool,
+ pipeline_io_contract: PipelineIOContract | None = None,
_media_digest_cache: MutableMapping[str, str] | None = None,
) -> HFDataset:
"""Build an input-only raw dataset for ``GeneralDataset`` preprocessing.
@@ -65,6 +73,23 @@ def project_offline_condition_dataset(
raise TypeError(
"ordered_references must be a bool, " f"got {type(ordered_references).__name__}"
)
+ if pipeline_io_contract is not None and not isinstance(
+ pipeline_io_contract, PipelineIOContract
+ ):
+ raise TypeError(
+ "pipeline_io_contract must be a PipelineIOContract or None, "
+ f"got {type(pipeline_io_contract).__name__}"
+ )
+ if pipeline_io_contract is not None:
+ expected_ordered_references = (
+ pipeline_io_contract.input_media.binding is InputMediaBinding.ORDERED_REFERENCES
+ )
+ if ordered_references != expected_ordered_references:
+ raise ValueError(
+ "offline condition projection binding disagrees with pipeline contract: "
+ f"ordered_references={ordered_references}, "
+ f"contract={pipeline_io_contract.input_media.binding.value!r}"
+ )
stable_records = tuple(records)
if not stable_records:
raise ValueError("offline condition projection requires at least one record")
@@ -91,9 +116,28 @@ def project_offline_condition_dataset(
"prompt": [record.model_input.prompt for record in stable_records],
OFFLINE_CONDITION_ID_COLUMN: condition_ids,
}
+ resolved_slots = []
+ for record in stable_records:
+ if pipeline_io_contract is None:
+ slots = tuple(None for _ in record.model_input.media)
+ if any(asset.slot is not None for asset in record.model_input.media):
+ raise ValueError(
+ "offline condition projection requires pipeline_io_contract when "
+ "input media declares semantic slots"
+ )
+ else:
+ slots = resolve_pipeline_input_media_slots(
+ record.model_input,
+ pipeline_io_contract,
+ )
+ resolved_slots.append(slots)
negative_prompts = [record.model_input.negative_prompt for record in stable_records]
- if any(value is not None for value in negative_prompts):
+ projects_negative_prompt = (
+ pipeline_io_contract is not None
+ and pipeline_io_contract.negative_prompt is not NegativePromptPolicy.UNSUPPORTED
+ )
+ if projects_negative_prompt or any(value is not None for value in negative_prompts):
# Adapter tokenizers consume a homogeneous text batch. In a mixed V2
# batch, an omitted optional negative prompt is semantically the empty
# prompt, not a tokenizer-level ``None`` value.
@@ -110,37 +154,60 @@ def project_offline_condition_dataset(
for index, record in enumerate(stable_records)
]
else:
+ grouped_rows = [
+ _group_media_by_type_and_slot(
+ record.model_input.media,
+ slots,
+ pipeline_io_contract,
+ )
+ for record, slots in zip(stable_records, resolved_slots)
+ ]
grouped_columns = {
- "images": [
- [asset.path for asset in record.model_input.media if asset.type == "image"]
- for record in stable_records
- ],
+ "images": [[asset.path for asset, _ in row["image"]] for row in grouped_rows],
"videos": [
- [
- _to_grouped_rate_spec(asset, rate_name="fps")
- for asset in record.model_input.media
- if asset.type == "video"
- ]
- for record in stable_records
+ [_to_grouped_rate_spec(asset, rate_name="fps") for asset, _ in row["video"]]
+ for row in grouped_rows
],
"audios": [
- [
- _to_grouped_rate_spec(asset, rate_name="sample_rate")
- for asset in record.model_input.media
- if asset.type == "audio"
- ]
- for record in stable_records
+ [_to_grouped_rate_spec(asset, rate_name="sample_rate") for asset, _ in row["audio"]]
+ for row in grouped_rows
],
+ "image_slots": [[slot for _, slot in row["image"]] for row in grouped_rows],
+ "video_slots": [[slot for _, slot in row["video"]] for row in grouped_rows],
+ "audio_slots": [[slot for _, slot in row["audio"]] for row in grouped_rows],
}
- columns.update(
- {column_name: values for column_name, values in grouped_columns.items() if any(values)}
+ declared_rules = (
+ {}
+ if pipeline_io_contract is None
+ else {rule.format.type.value: rule for rule in pipeline_io_contract.input_media.rules}
)
+ column_media_types = {
+ "images": "image",
+ "videos": "video",
+ "audios": "audio",
+ "image_slots": "image",
+ "video_slots": "video",
+ "audio_slots": "audio",
+ }
+ for column_name, values in grouped_columns.items():
+ media_type = column_media_types[column_name]
+ rule = declared_rules.get(media_type)
+ if column_name.endswith("_slots"):
+ include = rule is not None and bool(rule.slots)
+ else:
+ include = rule is not None if pipeline_io_contract is not None else any(values)
+ if include:
+ columns[column_name] = values
return HFDataset.from_dict(columns)
-def compute_offline_condition_source_hash(condition_ids: Sequence[str]) -> str:
- """Hash ordered input identities for an input-only cache fingerprint."""
+def compute_offline_condition_source_hash(
+ condition_ids: Sequence[str],
+ *,
+ pipeline_io_contract: PipelineIOContract | None = None,
+) -> str:
+ """Hash ordered inputs and their effective projection contract."""
stable_ids = tuple(condition_ids)
if not stable_ids:
raise ValueError("offline condition source hash requires at least one condition id")
@@ -150,10 +217,19 @@ def compute_offline_condition_source_hash(condition_ids: Sequence[str]) -> str:
"offline condition ids must be non-empty strings, "
f"got {condition_id!r} at index {index}"
)
+ if pipeline_io_contract is not None and not isinstance(
+ pipeline_io_contract,
+ PipelineIOContract,
+ ):
+ raise TypeError(
+ "pipeline_io_contract must be a PipelineIOContract or None, "
+ f"got {type(pipeline_io_contract).__name__}"
+ )
payload = json.dumps(
{
"format": _CONDITION_SOURCE_FORMAT,
"condition_ids": stable_ids,
+ "input_projection_contract": _input_projection_contract_identity(pipeline_io_contract),
},
ensure_ascii=False,
allow_nan=False,
@@ -169,6 +245,7 @@ def build_offline_condition_cache(
dataset_dir: str | os.PathLike[str],
preprocess_func: PreprocessCallable,
preprocess_kwargs: Mapping[str, Any] | None = None,
+ pipeline_io_contract: PipelineIOContract | None = None,
cache_dir: str | os.PathLike[str] = "~/.cache/flow_factory/datasets",
force_reprocess: bool = False,
preprocessing_batch_size: int | None = None,
@@ -193,16 +270,24 @@ def build_offline_condition_cache(
records,
source_name=source_name,
ordered_references=ordered_references,
+ pipeline_io_contract=pipeline_io_contract,
)
condition_ids = tuple(raw_dataset[OFFLINE_CONDITION_ID_COLUMN])
- source_hash = compute_offline_condition_source_hash(condition_ids)
+ source_hash = compute_offline_condition_source_hash(
+ condition_ids,
+ pipeline_io_contract=pipeline_io_contract,
+ )
normalized_dataset_dir = os.path.expanduser(os.fspath(dataset_dir))
normalized_cache_dir = os.path.expanduser(os.fspath(cache_dir))
normalized_preprocess_kwargs = dict(preprocess_kwargs or {})
normalized_extra_hash_strs = list(extra_hash_strs or ())
+ requires_single_sample_batches = ordered_references or (
+ pipeline_io_contract is not None
+ and pipeline_io_contract.batch_capability is BatchCapability.SINGLE_SAMPLE
+ )
if preprocessing_batch_size is None:
- preprocessing_batch_size = 1 if ordered_references else 16
+ preprocessing_batch_size = 1 if requires_single_sample_batches else 16
if (
not isinstance(preprocessing_batch_size, int)
or isinstance(preprocessing_batch_size, bool)
@@ -212,6 +297,8 @@ def build_offline_condition_cache(
"preprocessing_batch_size must be a positive integer, "
f"got {preprocessing_batch_size!r}"
)
+ if requires_single_sample_batches:
+ preprocessing_batch_size = 1
normalized_target_arrow_path: str
if target_arrow_path is None:
@@ -261,6 +348,36 @@ def _to_legacy_reference(asset: MediaAsset) -> Dict[str, Any]:
return reference
+def _input_projection_contract_identity(
+ contract: PipelineIOContract | None,
+) -> Dict[str, Any] | None:
+ """Return the canonical contract fields that can alter input preprocessing."""
+ if contract is None:
+ return None
+ input_media = contract.input_media
+ return {
+ "binding": input_media.binding.value,
+ "order": input_media.order.value,
+ "min_total_count": input_media.min_total_count,
+ "max_total_count": input_media.max_total_count,
+ "required_any_types": [value.value for value in input_media.required_any_types],
+ "rules": [
+ {
+ "type": rule.format.type.value,
+ "fps": rule.format.fps.value,
+ "sample_rate": rule.format.sample_rate.value,
+ "min_count": rule.min_count,
+ "max_count": rule.max_count,
+ "slots": list(rule.slots),
+ "required_slots": list(rule.required_slots),
+ }
+ for rule in input_media.rules
+ ],
+ "negative_prompt": contract.negative_prompt.value,
+ "batch_capability": contract.batch_capability.value,
+ }
+
+
def _to_grouped_rate_spec(asset: MediaAsset, *, rate_name: str) -> Dict[str, Any]:
spec: Dict[str, Any] = {"path": asset.path}
rate = getattr(asset, rate_name)
@@ -269,6 +386,29 @@ def _to_grouped_rate_spec(asset: MediaAsset, *, rate_name: str) -> Dict[str, Any
return spec
+def _group_media_by_type_and_slot(
+ media: Sequence[MediaAsset],
+ slots: Sequence[str | None],
+ contract: PipelineIOContract | None,
+) -> Dict[str, List[tuple[MediaAsset, str | None]]]:
+ """Group one record and canonicalize slotted media into declaration order."""
+ grouped: Dict[str, List[tuple[MediaAsset, str | None]]] = {
+ "image": [],
+ "video": [],
+ "audio": [],
+ }
+ for asset, slot in zip(media, slots):
+ grouped[asset.type].append((asset, slot))
+ if contract is None:
+ return grouped
+ for rule in contract.input_media.rules:
+ if not rule.slots:
+ continue
+ slot_order = {slot: index for index, slot in enumerate(rule.slots)}
+ grouped[rule.format.type.value].sort(key=lambda item: slot_order[item[1]])
+ return grouped
+
+
__all__ = [
"build_offline_condition_cache",
"compute_offline_condition_source_hash",
diff --git a/src/flow_factory/data_utils/offline_dataset.py b/src/flow_factory/data_utils/offline_dataset.py
index 9043a15bc..767257ffd 100644
--- a/src/flow_factory/data_utils/offline_dataset.py
+++ b/src/flow_factory/data_utils/offline_dataset.py
@@ -700,6 +700,7 @@ def _media_identity(
"path": media.path,
"fps": media.fps,
"sample_rate": media.sample_rate,
+ "slot": media.slot,
"content_sha256": _media_content_sha256(
media.path,
media_digest_cache=media_digest_cache,
diff --git a/src/flow_factory/data_utils/offline_train_data.py b/src/flow_factory/data_utils/offline_train_data.py
index 4fc16575a..f2b66dfd5 100644
--- a/src/flow_factory/data_utils/offline_train_data.py
+++ b/src/flow_factory/data_utils/offline_train_data.py
@@ -223,7 +223,7 @@ def build_offline_train_dataloader(
preprocess_func=preprocess_func,
preprocess_kwargs=normalized_preprocess_kwargs,
preprocessing_batch_size=data_args.preprocessing_batch_size,
- batch_capability=pipeline_io_contract.batch_capability,
+ pipeline_io_contract=pipeline_io_contract,
force_reprocess=data_args.force_reprocess,
extra_hash_strs=[*normalized_extra_hash_strs, f"offline_train_source:{source.name}"],
preprocess_parallelism=data_args.preprocess_parallelism,
@@ -315,7 +315,7 @@ def _build_distributed_condition_cache(
preprocess_func: PreprocessCallable,
preprocess_kwargs: Mapping[str, Any],
preprocessing_batch_size: int,
- batch_capability: BatchCapability,
+ pipeline_io_contract: PipelineIOContract,
force_reprocess: bool,
extra_hash_strs: Sequence[str],
preprocess_parallelism: Literal["global", "local"],
@@ -326,17 +326,22 @@ def _build_distributed_condition_cache(
ordered_references = _supports_ordered_references(preprocess_func)
effective_batch_size = (
1
- if ordered_references or batch_capability is BatchCapability.SINGLE_SAMPLE
+ if ordered_references
+ or pipeline_io_contract.batch_capability is BatchCapability.SINGLE_SAMPLE
else preprocessing_batch_size
)
raw_dataset = project_offline_condition_dataset(
records,
source_name=source_name,
ordered_references=ordered_references,
+ pipeline_io_contract=pipeline_io_contract,
_media_digest_cache=_media_digest_cache,
)
condition_ids = tuple(raw_dataset[OFFLINE_CONDITION_ID_COLUMN])
- source_hash = compute_offline_condition_source_hash(condition_ids)
+ source_hash = compute_offline_condition_source_hash(
+ condition_ids,
+ pipeline_io_contract=pipeline_io_contract,
+ )
dataset_builder = _create_or_load_dataset(
split=split,
accelerator=accelerator,
diff --git a/src/flow_factory/data_utils/schema.py b/src/flow_factory/data_utils/schema.py
index de80b12f4..941dd32b6 100644
--- a/src/flow_factory/data_utils/schema.py
+++ b/src/flow_factory/data_utils/schema.py
@@ -32,6 +32,7 @@
Field,
JsonValue,
field_validator,
+ model_validator,
)
MediaType = Literal["image", "video", "audio"]
@@ -52,6 +53,7 @@ class _MediaRefBase(_StrictFrozenModel):
"""Shared path contract for the exact-key media variants."""
path: str
+ slot: str | None = None
@field_validator("path")
@classmethod
@@ -60,6 +62,13 @@ def _validate_path(cls, value: str) -> str:
raise ValueError("media path must be a non-empty string")
return value
+ @field_validator("slot")
+ @classmethod
+ def _validate_slot(cls, value: str | None) -> str | None:
+ if value is not None and not value.strip():
+ raise ValueError("media slot must be a non-empty string when provided")
+ return value
+
class ImageRef(_MediaRefBase):
"""Image asset with no rate fields."""
@@ -100,6 +109,18 @@ class OutputCandidateSpec(_StrictFrozenModel):
media: List[MediaRef] = Field(min_length=1)
+ @model_validator(mode="after")
+ def _reject_input_only_slots(self) -> "OutputCandidateSpec":
+ slotted_indices = tuple(
+ index for index, media in enumerate(self.media) if media.slot is not None
+ )
+ if slotted_indices:
+ raise ValueError(
+ "media slot is input-only and cannot appear in an output candidate; "
+ f"indices={slotted_indices!r}"
+ )
+ return self
+
class DemonstrationSpec(_StrictFrozenModel):
"""A single supervised target without naming a training algorithm."""
@@ -139,6 +160,7 @@ class MediaAsset:
path: str
fps: float | None = None
sample_rate: int | None = None
+ slot: str | None = None
@dataclass(frozen=True, slots=True)
@@ -261,6 +283,7 @@ def _normalize_media(
path=path,
fps=media.fps if isinstance(media, VideoRef) else None,
sample_rate=media.sample_rate if isinstance(media, AudioRef) else None,
+ slot=media.slot,
)
diff --git a/src/flow_factory/models/__init__.py b/src/flow_factory/models/__init__.py
index 9dffe36b0..39e172aa6 100644
--- a/src/flow_factory/models/__init__.py
+++ b/src/flow_factory/models/__init__.py
@@ -21,6 +21,11 @@
"""
from .abc import BaseAdapter
+from .condition_state import (
+ ConditionStatePreparer,
+ PreparedConditionState,
+ validate_condition_preparer_required_components,
+)
from .latent_geometry import LatentAxes, LatentLayout, infer_latent_axes
from .loader import load_model
from .model_bundle import ModelBundle, RoutedComponentProxy
@@ -42,6 +47,10 @@
__all__ = [
# Core classes
"BaseAdapter",
+ # Runtime input-condition preparation
+ "ConditionStatePreparer",
+ "PreparedConditionState",
+ "validate_condition_preparer_required_components",
# Offline target encoding
"DecodedMediaBatch",
"EncodedOutputState",
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index 604e7276f..d110e78a1 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -110,6 +110,11 @@
select_gradient_checkpointing_units,
selective_gradient_checkpointing_function,
)
+from .condition_state import (
+ ConditionStatePreparer,
+ PreparedConditionState,
+ validate_condition_preparer_required_components,
+)
from .latent_geometry import LatentAxes, infer_latent_axes
from .model_bundle import RoutedComponentProxy
from .output_state import (
@@ -267,10 +272,12 @@ class BaseAdapter(ABC):
# name. Overriding one would silently bypass that contract, so subclasses are
# rejected at class creation instead of at training time.
_BOUNDARY_OWNING_METHODS: ClassVar[Tuple[str, ...]] = (
+ "prepare_condition_state",
"encode_output_state",
"decode_output_state",
"forward_state",
"reduce_component_latent_values",
+ "reduce_flow_matching_objective_values",
"reduce_latent_values",
)
@@ -283,6 +290,8 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
"Provide build_output_state_codec() and "
"_validate_encoded_output_geometry() instead."
)
+ elif name == "prepare_condition_state":
+ override_hint = "Provide build_condition_state_preparer() instead."
elif name == "decode_output_state":
override_hint = "Override the protected hook _decode_output_state instead."
else:
@@ -329,6 +338,17 @@ def __init__(self, config: Arguments, accelerator: Accelerator):
"expected SchedulerGroup.primary to be the canonical pipeline scheduler, "
f"received primary component {self.scheduler_group.primary_name!r}"
)
+ self._effective_pipeline_io_contract = self._resolve_pipeline_io_contract()
+ if self._effective_pipeline_io_contract is not None and not isinstance(
+ self._effective_pipeline_io_contract,
+ PipelineIOContract,
+ ):
+ raise TypeError(
+ f"adapter {type(self).__name__} expected _resolve_pipeline_io_contract() "
+ "to return PipelineIOContract or None, received "
+ f"{type(self._effective_pipeline_io_contract).__name__}: "
+ f"{self._effective_pipeline_io_contract!r}"
+ )
# Compatibility alias: the runtime override mapping is the sole authoritative cache.
self._components: Dict[str, torch.nn.Module] = cast(
@@ -338,6 +358,12 @@ def __init__(self, config: Arguments, accelerator: Accelerator):
self.model_args.target_components
)
+ # Build per-request input-condition realization after the component runtime
+ # exists, but before any codec may consume its declaration. Like the output
+ # codec, the preparer declares lifecycle metadata only.
+ self._condition_state_preparer = self._build_condition_state_preparer_declaration()
+ self._condition_state_encoding_modules = self._validate_condition_state_preparer_lifecycle()
+
# Build target-media encoding only after load-dtype policy, component runtime,
# scheduler group, and target-name canonicalization are established. The codec
# declaration is immutable lifecycle metadata; it must not materialize, load,
@@ -455,6 +481,116 @@ def cast_latent_state(
return state
return LatentState(components, active_masks=state.active_masks)
+ # =========================== Condition-State Preparation =======================
+ @property
+ def effective_pipeline_io_contract(self) -> Optional[PipelineIOContract]:
+ """Return the checkpoint-realized pipeline input/output contract."""
+ return getattr(
+ self,
+ "_effective_pipeline_io_contract",
+ type(self).pipeline_io_contract,
+ )
+
+ def _resolve_pipeline_io_contract(self) -> Optional[PipelineIOContract]:
+ """Resolve checkpoint-specific capabilities from the class declaration.
+
+ Most adapters use one immutable class-level contract. An adapter wrapping
+ checkpoint variants with narrower input semantics may override this hook
+ and return a validated specialization after its pipeline config is known.
+
+ Returns:
+ Effective contract for this adapter instance, or ``None``.
+ """
+ return type(self).pipeline_io_contract
+
+ def _build_condition_state_preparer_declaration(
+ self,
+ ) -> Optional[ConditionStatePreparer]:
+ """Build preparer metadata without changing component runtime state."""
+ materialized_before = tuple(self.component_runtime.materialized_component_names)
+ overrides_before = tuple(self.component_runtime.override_components)
+ preparer = self.build_condition_state_preparer()
+ materialized_after = tuple(self.component_runtime.materialized_component_names)
+ overrides_after = tuple(self.component_runtime.override_components)
+ if materialized_after != materialized_before or overrides_after != overrides_before:
+ raise RuntimeError(
+ f"adapter {type(self).__name__}.build_condition_state_preparer() must be "
+ "declaration-only and cannot materialize or replace components: "
+ f"materialized_before={materialized_before}, "
+ f"materialized_after={materialized_after}, "
+ f"overrides_before={overrides_before}, overrides_after={overrides_after}"
+ )
+ return preparer
+
+ @property
+ def condition_state_preparer(self) -> Optional[ConditionStatePreparer]:
+ """Return the immutable preparer selected during adapter construction."""
+ return getattr(self, "_condition_state_preparer", None)
+
+ @property
+ def condition_state_encoding_modules(self) -> Tuple[str, ...]:
+ """Return validated component names required for condition realization."""
+ return self._condition_state_encoding_modules
+
+ def build_condition_state_preparer(self) -> Optional[ConditionStatePreparer]:
+ """Build the adapter-owned runtime condition preparer, if required.
+
+ The default identity path needs no declaration. A model whose input
+ condition depends on runtime geometry, stochastic augmentation, or an
+ on-device encoder returns a declaration-only preparer here.
+
+ Returns:
+ Adapter-owned preparer, or ``None`` for identity preparation.
+ """
+ return None
+
+ def _validate_condition_state_preparer_lifecycle(self) -> Tuple[str, ...]:
+ """Validate condition-preparer lifecycle metadata."""
+ preparer = self.condition_state_preparer
+ if preparer is None:
+ return ()
+ return validate_condition_preparer_required_components(
+ preparer,
+ tuple(self.component_runtime.declared_component_names),
+ )
+
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ """Realize one input-owned condition state for a request or batch.
+
+ Args:
+ condition: Cached input-only model condition.
+ generator: Optional generator for adapter-owned stochastic realization.
+
+ Returns:
+ Validated prepared condition reused by every candidate and forward
+ derived from this request.
+ """
+ if not isinstance(condition, Mapping):
+ raise TypeError(
+ "expected condition-state input to be Mapping[str, Any], "
+ f"received {type(condition).__name__}: {condition!r}"
+ )
+ if generator is not None and not isinstance(generator, torch.Generator):
+ raise TypeError(
+ "expected condition-state generator to be torch.Generator or None, "
+ f"received {type(generator).__name__}: {generator!r}"
+ )
+ preparer = self.condition_state_preparer
+ if preparer is None:
+ return PreparedConditionState.identity(condition)
+ with torch.no_grad():
+ prepared = preparer.prepare_condition_state(condition, generator)
+ if not isinstance(prepared, PreparedConditionState):
+ raise TypeError(
+ "condition-state preparer must return PreparedConditionState, "
+ f"received {type(prepared).__name__}"
+ )
+ return prepared
+
# ============================ Output-State Encoding ============================
def _build_output_state_codec_declaration(self) -> Optional[OutputStateCodec]:
"""Build codec metadata without changing component materialization or overrides."""
@@ -562,7 +698,7 @@ def validate_offline_output_capability(cls) -> None:
def _validate_output_state_codec_lifecycle(self) -> Tuple[str, ...]:
"""Validate the adapter's pipeline contract and codec declaration."""
unavailable_reason = type(self)._validated_output_state_codec_unavailable_reason()
- contract = self.pipeline_io_contract
+ contract = self.effective_pipeline_io_contract
if contract is not None and not isinstance(contract, PipelineIOContract):
raise TypeError(
f"adapter {type(self).__name__} expected pipeline_io_contract to be "
@@ -590,14 +726,15 @@ def _validate_output_state_codec_lifecycle(self) -> Tuple[str, ...]:
def encode_output_state(
self,
media_batch: DecodedMediaBatch,
- condition: Mapping[str, Any],
+ condition: Union[Mapping[str, Any], PreparedConditionState],
generator: Optional[torch.Generator] = None,
) -> EncodedOutputState:
"""Encode decoded targets through the adapter-owned validated boundary.
Args:
media_batch: Exact output-media sequence for every batch sample.
- condition: Model-input condition for the same batch.
+ condition: Cached or already-prepared model-input condition for the
+ same batch.
generator: Optional deterministic generator used by stochastic encoders.
Returns:
@@ -614,7 +751,7 @@ def encode_output_state(
f"offline output-state encoding is unavailable for adapter "
f"{type(self).__name__}: {unavailable_reason}"
)
- contract = self.pipeline_io_contract
+ contract = self.effective_pipeline_io_contract
if contract is None:
raise RuntimeError(
f"adapter {type(self).__name__} cannot encode output state because it does not "
@@ -626,11 +763,6 @@ def encode_output_state(
f"adapter {type(self).__name__} declares pipeline_io_contract but does not "
"provide an output-state codec through build_output_state_codec()"
)
- if not isinstance(condition, Mapping):
- raise TypeError(
- "expected output-state condition to be Mapping[str, Any], "
- f"received {type(condition).__name__}: {condition!r}"
- )
if generator is not None and not isinstance(generator, torch.Generator):
raise TypeError(
"expected output-state generator to be torch.Generator or None, "
@@ -638,10 +770,21 @@ def encode_output_state(
)
validated_media = validate_output_candidate_batch(media_batch, contract)
+ if isinstance(condition, PreparedConditionState):
+ prepared_condition = condition
+ elif isinstance(condition, Mapping):
+ prepared_condition = self.prepare_condition_state(condition, generator)
+ else:
+ raise TypeError(
+ "expected output-state condition to be Mapping[str, Any] or "
+ "PreparedConditionState, "
+ f"received {type(condition).__name__}: {condition!r}"
+ )
+ codec_condition = prepared_condition.output_codec_condition()
with torch.no_grad():
encoded = codec.encode_output_state(
validated_media,
- condition,
+ codec_condition,
generator,
)
@@ -672,7 +815,7 @@ def encode_output_state(
device=self.device,
)
- self._validate_encoded_output_geometry(validated_media, condition, encoded)
+ self._validate_encoded_output_geometry(validated_media, codec_condition, encoded)
return encoded
def decode_output_state(
@@ -754,7 +897,8 @@ def _validate_encoded_output_geometry(
raise NotImplementedError(
f"adapter {type(self).__name__} provides an output-state codec but must override "
"_validate_encoded_output_geometry() to validate geometry signatures against "
- f"geometry_source={self.pipeline_io_contract.geometry_source.value!r}"
+ "geometry_source="
+ f"{self.effective_pipeline_io_contract.geometry_source.value!r}"
)
# ============================== Loading Components ==============================
@@ -4371,6 +4515,40 @@ def _reduce_latent_values(
state=state,
)
+ def reduce_flow_matching_objective_values(
+ self,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState] = None,
+ ) -> torch.Tensor:
+ """Reduce offline flow-matching errors to one scalar per sample.
+
+ This objective-specific boundary is deliberately separate from the
+ trajectory-wide element-weighted reducer used by online policy and
+ distillation algorithms. Most adapters inherit the existing global
+ reduction unchanged; multi-modal training recipes may override the
+ protected hook without changing rollout likelihood semantics.
+
+ Args:
+ values: Per-element squared errors in component order.
+ state: Noised state supplying active masks.
+
+ Returns:
+ One flow-matching objective value per batch sample.
+ """
+ batch_size = bridge.validate_reduction_inputs(self, values, state)
+ reduced = self._reduce_flow_matching_objective_values(values, state=state)
+ return bridge.validate_reduced_latent_values(self, reduced, batch_size)
+
+ def _reduce_flow_matching_objective_values(
+ self,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState] = None,
+ ) -> torch.Tensor:
+ """Use the existing globally element-weighted reduction by default."""
+ return self.reduce_latent_values(values, state=state)
+
# ======================================= Sampling & Training =======================================
@abstractmethod
def forward(
diff --git a/src/flow_factory/models/condition_state.py b/src/flow_factory/models/condition_state.py
new file mode 100644
index 000000000..6d386a896
--- /dev/null
+++ b/src/flow_factory/models/condition_state.py
@@ -0,0 +1,203 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Adapter-owned runtime realization of cached model-input conditions."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Any, Optional, Protocol, Tuple, runtime_checkable
+
+import torch
+
+from ..contracts import NON_MODEL_CONDITION_KEYS
+
+
+@dataclass(frozen=True, slots=True)
+class PreparedConditionState:
+ """Bundle cached conditions with one realized model-input state.
+
+ Args:
+ condition: Cached input-only fields retained for model forward.
+ forward_context: Runtime input-owned fields added to model forward.
+ output_context: Runtime input-owned fields consumed only while encoding
+ and binding an output target.
+
+ Note:
+ The ownership shell and outer mappings are copied and frozen. Tensor
+ leaves are retained without cloning so online rollout and offline DPO can
+ reuse one exact stochastic condition realization.
+ """
+
+ condition: Mapping[str, Any]
+ forward_context: Mapping[str, Any]
+ output_context: Mapping[str, Any]
+
+ def __post_init__(self) -> None:
+ """Freeze mappings and reject ambiguous field ownership."""
+ condition = _freeze_string_mapping(self.condition, "PreparedConditionState.condition")
+ forward_context = _freeze_string_mapping(
+ self.forward_context,
+ "PreparedConditionState.forward_context",
+ )
+ output_context = _freeze_string_mapping(
+ self.output_context,
+ "PreparedConditionState.output_context",
+ )
+ _reject_non_model_keys(condition, "PreparedConditionState.condition")
+ _reject_non_model_keys(forward_context, "PreparedConditionState.forward_context")
+
+ forward_collisions = tuple(sorted(set(condition).intersection(forward_context)))
+ if forward_collisions:
+ raise ValueError(
+ "prepared condition forward context collides with cached condition keys "
+ f"{forward_collisions}; every model-forward field must have one owner"
+ )
+ output_collisions = tuple(sorted(set(condition).intersection(output_context)))
+ if output_collisions:
+ raise ValueError(
+ "prepared condition output context collides with cached condition keys "
+ f"{output_collisions}; every output-binding field must have one owner"
+ )
+
+ object.__setattr__(self, "condition", condition)
+ object.__setattr__(self, "forward_context", forward_context)
+ object.__setattr__(self, "output_context", output_context)
+
+ @classmethod
+ def identity(cls, condition: Mapping[str, Any]) -> "PreparedConditionState":
+ """Create a realization that preserves the cached condition unchanged.
+
+ Args:
+ condition: Cached input-only model fields.
+
+ Returns:
+ Identity prepared condition with no runtime-owned contexts.
+ """
+ return cls(condition=condition, forward_context={}, output_context={})
+
+ def model_forward_condition(self) -> Mapping[str, Any]:
+ """Return the collision-free model-forward condition mapping."""
+ return MappingProxyType({**self.condition, **self.forward_context})
+
+ def output_codec_condition(self) -> Mapping[str, Any]:
+ """Return cached and output-binding fields for target encoding."""
+ return MappingProxyType({**self.condition, **self.output_context})
+
+
+@runtime_checkable
+class ConditionStatePreparer(Protocol):
+ """Define adapter-owned per-request condition realization."""
+
+ @property
+ def required_components(self) -> Tuple[str, ...]:
+ """Return adapter component names required while realizing conditions."""
+ ...
+
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ """Prepare one exact model-input state from cached conditions."""
+ ...
+
+
+def validate_condition_preparer_required_components(
+ preparer: object,
+ available_components: Sequence[str],
+) -> Tuple[str, ...]:
+ """Validate preparer component requirements against an adapter runtime.
+
+ Args:
+ preparer: Structural condition-state preparer instance.
+ available_components: Canonical component names exposed by the runtime.
+
+ Returns:
+ The preparer's validated required component tuple.
+ """
+ prepare = getattr(preparer, "prepare_condition_state", None)
+ if not callable(prepare):
+ raise TypeError(
+ "expected condition-state preparer with callable prepare_condition_state, "
+ f"received {type(preparer).__name__}"
+ )
+ required_components = getattr(preparer, "required_components", None)
+ _validate_component_names(
+ required_components,
+ "condition preparer.required_components",
+ allow_empty=True,
+ )
+ if isinstance(available_components, (str, bytes)) or not isinstance(
+ available_components, Sequence
+ ):
+ raise TypeError(
+ "expected available_components to be a sequence of strings, "
+ f"received {type(available_components).__name__}: {available_components!r}"
+ )
+ available = tuple(available_components)
+ _validate_component_names(available, "available_components", allow_empty=True)
+ unknown = tuple(name for name in required_components if name not in available)
+ if unknown:
+ raise ValueError(
+ "condition-state preparer requires unknown adapter components "
+ f"{unknown}; available components={available}"
+ )
+ return tuple(required_components)
+
+
+def _freeze_string_mapping(value: object, identifier: str) -> Mapping[str, Any]:
+ if not isinstance(value, Mapping):
+ raise TypeError(
+ f"expected Mapping[str, Any] for {identifier}, received "
+ f"{type(value).__name__}: {value!r}"
+ )
+ invalid_keys = tuple(key for key in value if not isinstance(key, str) or not key)
+ if invalid_keys:
+ raise TypeError(f"expected non-empty string keys for {identifier}, got {invalid_keys!r}")
+ return MappingProxyType(dict(value))
+
+
+def _reject_non_model_keys(value: Mapping[str, Any], identifier: str) -> None:
+ rejected = tuple(sorted(set(value).intersection(NON_MODEL_CONDITION_KEYS)))
+ if rejected:
+ raise ValueError(
+ f"{identifier} contains fields that cannot enter model forward: {rejected}"
+ )
+
+
+def _validate_component_names(
+ value: object,
+ identifier: str,
+ *,
+ allow_empty: bool,
+) -> None:
+ if type(value) is not tuple:
+ raise TypeError(f"expected tuple[str, ...] for {identifier}, received {value!r}")
+ if not value and not allow_empty:
+ raise ValueError(f"expected {identifier} to contain at least one component")
+ invalid = tuple(name for name in value if not isinstance(name, str) or not name)
+ if invalid:
+ raise TypeError(f"expected non-empty strings for {identifier}, received {invalid!r}")
+ if len(set(value)) != len(value):
+ raise ValueError(f"expected unique component names for {identifier}, received {value!r}")
+
+
+__all__ = [
+ "ConditionStatePreparer",
+ "PreparedConditionState",
+ "validate_condition_preparer_required_components",
+]
diff --git a/src/flow_factory/models/configured_image_output.py b/src/flow_factory/models/configured_image_output.py
index 39b75e568..a2a5729cb 100644
--- a/src/flow_factory/models/configured_image_output.py
+++ b/src/flow_factory/models/configured_image_output.py
@@ -189,7 +189,7 @@ def build_output_state_codec(self) -> OutputStateCodec:
TypeError: If no pipeline contract is declared.
ValueError: If output media or geometry ownership is incompatible.
"""
- contract = self.pipeline_io_contract
+ contract = self.effective_pipeline_io_contract
if contract is None:
raise TypeError(
f"adapter {type(self).__name__} must declare pipeline_io_contract before "
diff --git a/src/flow_factory/models/ltx2/_output.py b/src/flow_factory/models/ltx2/_output.py
new file mode 100644
index 000000000..c38c9e0a1
--- /dev/null
+++ b/src/flow_factory/models/ltx2/_output.py
@@ -0,0 +1,1066 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""On-the-fly audiovisual target encoding shared by the LTX2 adapters."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping
+from dataclasses import dataclass
+from numbers import Real
+from types import MappingProxyType
+from typing import Any, ClassVar, Literal, Optional, Tuple
+
+import numpy as np
+import torch
+import torchaudio
+
+from ...contracts import MediaType
+from ...samples import LatentState
+from ...utils.audio import convert_audio
+from ..condition_state import PreparedConditionState
+from ..configured_image_output import retrieve_vae_latents
+from ..output_state import (
+ DecodedMediaBatch,
+ EncodedOutputState,
+ GeometrySignature,
+ MediaGeometrySignature,
+)
+
+LTX2_OFFLINE_FORWARD_OVERRIDES = MappingProxyType(
+ {
+ "guidance_scale": 1.0,
+ "audio_guidance_scale": 1.0,
+ "guidance_rescale": 0.0,
+ "audio_guidance_rescale": 0.0,
+ "stg_scale": 0.0,
+ "audio_stg_scale": 0.0,
+ "spatio_temporal_guidance_blocks": None,
+ "modality_scale": 1.0,
+ "audio_modality_scale": 1.0,
+ "preserve_raw_model_velocity": True,
+ }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class LTX2VideoGeometry:
+ """Canonical configured video and packed-latent geometry."""
+
+ height: int
+ width: int
+ num_frames: int
+ frame_rate: float
+ latent_frames: int
+ latent_height: int
+ latent_width: int
+ latent_channels: int
+ patch_size: int
+ patch_size_t: int
+ sequence_length: int
+ feature_dim: int
+
+
+@dataclass(frozen=True, slots=True)
+class LTX2AudioGeometry:
+ """Canonical configured waveform, mel, and packed-latent geometry."""
+
+ sample_rate: int
+ hop_length: int
+ waveform_channels: int
+ target_samples: int
+ mel_bins: int
+ latent_mel_bins: int
+ latent_channels: int
+ latent_frames: int
+ temporal_compression_ratio: int
+ mel_compression_ratio: int
+ sequence_length: int
+ feature_dim: int
+
+
+@dataclass(frozen=True, slots=True)
+class LTX2OutputGeometry:
+ """One aligned LTX2 video/audio output geometry."""
+
+ video: LTX2VideoGeometry
+ audio: LTX2AudioGeometry
+
+
+@dataclass(frozen=True, slots=True)
+class LTX2FirstFrameConditionPreparer:
+ """Encode an I2AV condition image once for all offline target candidates."""
+
+ adapter: Any
+ required_components: ClassVar[Tuple[str, ...]] = ("vae",)
+
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ """Realize the deterministic VideoVAE posterior mode for the first frame."""
+ del generator
+ if "condition_images" not in condition:
+ raise ValueError(
+ "LTX2 I2AV offline condition requires cached `condition_images` pixels"
+ )
+ video_geometry = resolve_ltx2_video_geometry(self.adapter)
+ pixels = condition["condition_images"]
+ if not isinstance(pixels, torch.Tensor):
+ raise TypeError(
+ "LTX2 I2AV condition_images must be a torch.Tensor, "
+ f"received {type(pixels).__name__}"
+ )
+ expected_shape = (
+ pixels.shape[0] if pixels.ndim == 4 else None,
+ 3,
+ video_geometry.height,
+ video_geometry.width,
+ )
+ if pixels.ndim != 4 or tuple(pixels.shape) != expected_shape:
+ raise ValueError(
+ "LTX2 I2AV condition_images must use configured BCHW geometry: "
+ f"expected {expected_shape}, received {tuple(pixels.shape)}"
+ )
+ if not pixels.is_floating_point():
+ raise TypeError(
+ "LTX2 I2AV condition_images must be floating pixels, " f"received {pixels.dtype}"
+ )
+ _require_finite_tensor(pixels, "LTX2 I2AV condition_images")
+
+ vae = self.adapter.get_component("vae")
+ vae_dtype = _floating_module_dtype(vae, "LTX2 VideoVAE")
+ encoded = vae.encode(pixels.to(device=self.adapter.device, dtype=vae_dtype).unsqueeze(2))
+ condition_latents = retrieve_vae_latents(
+ encoded,
+ sample_mode="argmax",
+ source="LTX2 I2AV condition image",
+ ).to(device=self.adapter.device, dtype=torch.float32)
+ expected_latent_shape = (
+ pixels.shape[0],
+ video_geometry.latent_channels,
+ 1,
+ video_geometry.latent_height,
+ video_geometry.latent_width,
+ )
+ if tuple(condition_latents.shape) != expected_latent_shape:
+ raise ValueError(
+ "LTX2 I2AV condition VideoVAE latent geometry mismatch: "
+ f"expected {expected_latent_shape}, received {tuple(condition_latents.shape)}"
+ )
+ _require_finite_tensor(condition_latents, "LTX2 I2AV condition latents")
+
+ # Pixels are an encoder input, not a transformer input. Dropping them from
+ # the realized condition also avoids retaining a full-resolution tensor
+ # through both chosen and rejected offline forwards.
+ model_condition = {
+ key: value for key, value in condition.items() if key != "condition_images"
+ }
+ return PreparedConditionState(
+ condition=model_condition,
+ forward_context={},
+ output_context={"condition_video_latents": condition_latents.detach()},
+ )
+
+
+@dataclass(frozen=True, slots=True)
+class LTX2AVOutputCodec:
+ """Encode exact ``(video, audio)`` targets into LTX2's joint latent state."""
+
+ adapter: Any
+ conditioned: bool = False
+ required_components: ClassVar[Tuple[str, ...]] = ("vae", "audio_vae")
+
+ def encode_output_state(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> EncodedOutputState:
+ """Encode VideoVAE and AudioVAE posterior modes on the configured AV grid."""
+ del generator
+ geometry = resolve_ltx2_output_geometry(self.adapter, conditioned=self.conditioned)
+ videos = []
+ waveforms = []
+ for sample_index, candidate in enumerate(media_batch):
+ if len(candidate) != 2:
+ raise ValueError(
+ "LTX2 output codec expected exact (video, audio) media, "
+ f"received {len(candidate)} items for sample {sample_index}"
+ )
+ video_media, audio_media = candidate
+ videos.append(
+ prepare_ltx2_target_video(
+ video_media.payload,
+ source_fps=video_media.fps,
+ geometry=geometry.video,
+ )
+ )
+ waveforms.append(
+ prepare_ltx2_target_audio(
+ audio_media.payload,
+ source_sample_rate=audio_media.sample_rate,
+ geometry=geometry.audio,
+ duration_seconds=geometry.video.num_frames / geometry.video.frame_rate,
+ )
+ )
+
+ video_latents = encode_ltx2_target_video(self.adapter, videos, geometry.video)
+ if self.conditioned:
+ condition_latents = condition.get("condition_video_latents")
+ expected_condition_shape = (
+ len(media_batch),
+ geometry.video.latent_channels,
+ 1,
+ geometry.video.latent_height,
+ geometry.video.latent_width,
+ )
+ if not isinstance(condition_latents, torch.Tensor):
+ raise TypeError(
+ "LTX2 I2AV output binding requires condition_video_latents from "
+ "prepare_condition_state()"
+ )
+ if tuple(condition_latents.shape) != expected_condition_shape:
+ raise ValueError(
+ "LTX2 I2AV condition latent batch/geometry mismatch: "
+ f"expected {expected_condition_shape}, received "
+ f"{tuple(condition_latents.shape)}"
+ )
+ condition_latents = condition_latents.to(
+ device=self.adapter.device,
+ dtype=torch.float32,
+ )
+ _require_finite_tensor(condition_latents, "LTX2 I2AV condition latents")
+ video_latents = torch.cat(
+ [condition_latents, video_latents[:, :, 1:]],
+ dim=2,
+ )
+
+ packed_video = normalize_and_pack_ltx2_video(
+ self.adapter,
+ video_latents,
+ geometry.video,
+ )
+ packed_audio = encode_ltx2_target_audio(
+ self.adapter,
+ torch.stack(waveforms),
+ geometry.audio,
+ )
+
+ forward_context = _ltx2_forward_context(geometry)
+ active_masks = None
+ if self.conditioned:
+ unpacked_mask = packed_video.new_zeros(
+ (
+ len(media_batch),
+ 1,
+ geometry.video.latent_frames,
+ geometry.video.latent_height,
+ geometry.video.latent_width,
+ )
+ )
+ unpacked_mask[:, :, 0] = 1.0
+ conditioning_mask = self.adapter.pipeline._pack_latents(
+ unpacked_mask,
+ geometry.video.patch_size,
+ geometry.video.patch_size_t,
+ )
+ if conditioning_mask.shape[-1] != 1:
+ raise ValueError(
+ "LTX2 I2AV conditioning mask requires one scalar per packed token, "
+ f"received packed shape {tuple(conditioning_mask.shape)}"
+ )
+ conditioning_mask = conditioning_mask.squeeze(-1)
+ forward_context["conditioning_mask"] = conditioning_mask
+ active_masks = {
+ "video": (~conditioning_mask.bool()).unsqueeze(-1),
+ "audio": torch.ones(
+ (len(media_batch), geometry.audio.sequence_length, 1),
+ device=packed_audio.device,
+ dtype=torch.bool,
+ ),
+ }
+
+ signature = GeometrySignature(
+ media=(
+ MediaGeometrySignature(
+ type=MediaType.VIDEO,
+ height=geometry.video.height,
+ width=geometry.video.width,
+ frames=geometry.video.num_frames,
+ fps=geometry.video.frame_rate,
+ ),
+ MediaGeometrySignature(
+ type=MediaType.AUDIO,
+ samples=geometry.audio.target_samples,
+ sample_rate=geometry.audio.sample_rate,
+ ),
+ )
+ )
+ return EncodedOutputState(
+ clean_state=LatentState(
+ {"video": packed_video.detach(), "audio": packed_audio.detach()},
+ active_masks=active_masks,
+ ),
+ forward_context=forward_context,
+ decode_context={
+ "height": geometry.video.height,
+ "width": geometry.video.width,
+ "num_frames": geometry.video.num_frames,
+ "frame_rate": geometry.video.frame_rate,
+ },
+ geometry_signatures=tuple(signature for _ in media_batch),
+ )
+
+
+def resolve_ltx2_video_geometry(adapter: Any) -> LTX2VideoGeometry:
+ """Resolve configured video geometry from runtime VAE and transformer metadata."""
+ height = _positive_int(getattr(adapter.training_args, "height", None), "training_args.height")
+ width = _positive_int(getattr(adapter.training_args, "width", None), "training_args.width")
+ num_frames = _positive_int(
+ getattr(adapter.training_args, "num_frames", None),
+ "training_args.num_frames",
+ )
+ frame_rate = _positive_real(
+ getattr(adapter.training_args, "frame_rate", None),
+ "training_args.frame_rate",
+ )
+ pipeline = adapter.pipeline
+ spatial_ratio = _positive_int(
+ getattr(pipeline, "vae_spatial_compression_ratio", None),
+ "pipeline.vae_spatial_compression_ratio",
+ )
+ temporal_ratio = _positive_int(
+ getattr(pipeline, "vae_temporal_compression_ratio", None),
+ "pipeline.vae_temporal_compression_ratio",
+ )
+ if height % spatial_ratio or width % spatial_ratio:
+ raise ValueError(
+ "LTX2 configured height/width must be divisible by the VideoVAE spatial "
+ f"compression ratio {spatial_ratio}, received {(height, width)}"
+ )
+ if (num_frames - 1) % temporal_ratio:
+ raise ValueError(
+ "LTX2 configured num_frames must satisfy "
+ f"(num_frames - 1) % {temporal_ratio} == 0, received {num_frames}"
+ )
+ latent_frames = (num_frames - 1) // temporal_ratio + 1
+ latent_height = height // spatial_ratio
+ latent_width = width // spatial_ratio
+ patch_size = _positive_int(
+ getattr(pipeline, "transformer_spatial_patch_size", None),
+ "pipeline.transformer_spatial_patch_size",
+ )
+ patch_size_t = _positive_int(
+ getattr(pipeline, "transformer_temporal_patch_size", None),
+ "pipeline.transformer_temporal_patch_size",
+ )
+ if latent_frames % patch_size_t or latent_height % patch_size or latent_width % patch_size:
+ raise ValueError(
+ "LTX2 latent video grid must be divisible by transformer patches "
+ f"(time={patch_size_t}, spatial={patch_size}), received "
+ f"{(latent_frames, latent_height, latent_width)}"
+ )
+ vae = adapter.get_component("vae")
+ latent_channels = _positive_int(
+ getattr(getattr(vae, "config", None), "latent_channels", None),
+ "vae.config.latent_channels",
+ )
+ feature_dim = latent_channels * patch_size_t * patch_size * patch_size
+ sequence_length = (
+ latent_frames // patch_size_t * (latent_height // patch_size) * (latent_width // patch_size)
+ )
+ return LTX2VideoGeometry(
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ frame_rate=frame_rate,
+ latent_frames=latent_frames,
+ latent_height=latent_height,
+ latent_width=latent_width,
+ latent_channels=latent_channels,
+ patch_size=patch_size,
+ patch_size_t=patch_size_t,
+ sequence_length=sequence_length,
+ feature_dim=feature_dim,
+ )
+
+
+def resolve_ltx2_output_geometry(
+ adapter: Any,
+ *,
+ conditioned: bool,
+) -> LTX2OutputGeometry:
+ """Resolve and cross-check the complete config-driven LTX2 AV geometry."""
+ video = resolve_ltx2_video_geometry(adapter)
+ if conditioned and (video.patch_size != 1 or video.patch_size_t != 1):
+ raise ValueError(
+ "LTX2 I2AV's official scalar conditioning mask requires video patch_size=1 "
+ f"and patch_size_t=1, received {(video.patch_size, video.patch_size_t)}"
+ )
+
+ pipeline = adapter.pipeline
+ audio_vae = adapter.get_component("audio_vae")
+ audio_config = getattr(audio_vae, "config", None)
+ sample_rate = _positive_int(
+ getattr(audio_config, "sample_rate", None),
+ "audio_vae.config.sample_rate",
+ )
+ hop_length = _positive_int(
+ getattr(audio_config, "mel_hop_length", None),
+ "audio_vae.config.mel_hop_length",
+ )
+ mel_bins = _positive_int(
+ getattr(audio_config, "mel_bins", None),
+ "audio_vae.config.mel_bins",
+ )
+ waveform_channels = _positive_int(
+ getattr(audio_config, "in_channels", None),
+ "audio_vae.config.in_channels",
+ )
+ if waveform_channels not in (1, 2):
+ raise ValueError(
+ "LTX2 audio frontend supports mono or stereo AudioVAE inputs, "
+ f"received in_channels={waveform_channels}"
+ )
+ latent_channels = _positive_int(
+ getattr(audio_config, "latent_channels", None),
+ "audio_vae.config.latent_channels",
+ )
+ temporal_ratio = _positive_int(
+ getattr(pipeline, "audio_vae_temporal_compression_ratio", None),
+ "pipeline.audio_vae_temporal_compression_ratio",
+ )
+ mel_ratio = _positive_int(
+ getattr(pipeline, "audio_vae_mel_compression_ratio", None),
+ "pipeline.audio_vae_mel_compression_ratio",
+ )
+ if mel_bins % mel_ratio:
+ raise ValueError(
+ "LTX2 AudioVAE mel bins must be divisible by mel compression ratio, "
+ f"received mel_bins={mel_bins}, ratio={mel_ratio}"
+ )
+ pipeline_sample_rate = _positive_int(
+ getattr(pipeline, "audio_sampling_rate", None),
+ "pipeline.audio_sampling_rate",
+ )
+ pipeline_hop_length = _positive_int(
+ getattr(pipeline, "audio_hop_length", None),
+ "pipeline.audio_hop_length",
+ )
+ if (sample_rate, hop_length) != (pipeline_sample_rate, pipeline_hop_length):
+ raise ValueError(
+ "LTX2 pipeline and AudioVAE audio clocks disagree: "
+ f"pipeline={(pipeline_sample_rate, pipeline_hop_length)}, "
+ f"audio_vae={(sample_rate, hop_length)}"
+ )
+
+ duration_seconds = video.num_frames / video.frame_rate
+ target_samples = max(round(duration_seconds * sample_rate), 1)
+ latent_frames = round(duration_seconds * sample_rate / hop_length / temporal_ratio)
+ if latent_frames < 1:
+ raise ValueError(
+ "LTX2 configured AV duration resolves to no audio latent frames: "
+ f"duration={duration_seconds}, sample_rate={sample_rate}, "
+ f"hop_length={hop_length}, temporal_ratio={temporal_ratio}"
+ )
+ latent_mel_bins = mel_bins // mel_ratio
+ feature_dim = latent_channels * latent_mel_bins
+ transformer_config = _component_config(adapter, "transformer")
+ transformer_video_dim = _positive_int(
+ getattr(transformer_config, "in_channels", None),
+ "transformer.config.in_channels",
+ )
+ transformer_audio_dim = _positive_int(
+ getattr(transformer_config, "audio_in_channels", None),
+ "transformer.config.audio_in_channels",
+ )
+ if video.feature_dim != transformer_video_dim:
+ raise ValueError(
+ "LTX2 packed video feature width disagrees with transformer config: "
+ f"expected {transformer_video_dim}, resolved {video.feature_dim}"
+ )
+ if feature_dim != transformer_audio_dim:
+ raise ValueError(
+ "LTX2 packed audio feature width disagrees with transformer config: "
+ f"expected {transformer_audio_dim}, resolved {feature_dim}"
+ )
+ audio_patch_size = _positive_int(
+ getattr(transformer_config, "audio_patch_size", 1),
+ "transformer.config.audio_patch_size",
+ )
+ audio_patch_size_t = _positive_int(
+ getattr(transformer_config, "audio_patch_size_t", 1),
+ "transformer.config.audio_patch_size_t",
+ )
+ if (audio_patch_size, audio_patch_size_t) != (1, 1):
+ raise ValueError(
+ "LTX2 Diffusers 0.40 packs audio as one full-mel token per latent time; "
+ "non-unit audio patching is not supported by the online pipeline, received "
+ f"{(audio_patch_size, audio_patch_size_t)}"
+ )
+
+ audio = LTX2AudioGeometry(
+ sample_rate=sample_rate,
+ hop_length=hop_length,
+ waveform_channels=waveform_channels,
+ target_samples=target_samples,
+ mel_bins=mel_bins,
+ latent_mel_bins=latent_mel_bins,
+ latent_channels=latent_channels,
+ latent_frames=latent_frames,
+ temporal_compression_ratio=temporal_ratio,
+ mel_compression_ratio=mel_ratio,
+ sequence_length=latent_frames,
+ feature_dim=feature_dim,
+ )
+ return LTX2OutputGeometry(video=video, audio=audio)
+
+
+def prepare_ltx2_target_video(
+ payload: Any,
+ *,
+ source_fps: Any,
+ geometry: LTX2VideoGeometry,
+) -> np.ndarray:
+ """Select deterministic nearest-time RGB frames on the configured cadence."""
+ if not isinstance(payload, np.ndarray):
+ raise TypeError(
+ "LTX2 target video expected a decoded NumPy array, "
+ f"received {type(payload).__name__}"
+ )
+ if payload.dtype != np.uint8 or payload.ndim != 4 or payload.shape[-1] != 3:
+ raise ValueError(
+ "LTX2 target video must be uint8 RGB shaped (F,H,W,3), "
+ f"received dtype={payload.dtype}, shape={tuple(payload.shape)}"
+ )
+ if payload.shape[0] < 1:
+ raise ValueError("LTX2 target video must contain at least one frame")
+ source_fps = _positive_real(source_fps, "target video fps")
+ indices = np.rint(
+ np.arange(geometry.num_frames, dtype=np.float64) * source_fps / geometry.frame_rate
+ ).astype(np.int64)
+ if indices[-1] >= payload.shape[0]:
+ required_duration = (geometry.num_frames - 1) / geometry.frame_rate
+ available_duration = (payload.shape[0] - 1) / source_fps
+ raise ValueError(
+ "LTX2 target video is too short for configured temporal geometry: "
+ f"requires {required_duration:.6f}s, has {available_duration:.6f}s"
+ )
+ return np.ascontiguousarray(payload[indices])
+
+
+def prepare_ltx2_target_audio(
+ payload: Any,
+ *,
+ source_sample_rate: Any,
+ geometry: LTX2AudioGeometry,
+ duration_seconds: float,
+) -> torch.Tensor:
+ """Convert one waveform to the exact official LTX2 model-rate audio clock."""
+ if not isinstance(payload, torch.Tensor):
+ raise TypeError(
+ "LTX2 target audio expected a decoded torch.Tensor, "
+ f"received {type(payload).__name__}"
+ )
+ if payload.ndim != 2 or payload.shape[0] not in (1, 2) or payload.shape[1] < 1:
+ raise ValueError(
+ "LTX2 target audio must be non-empty mono/stereo shaped (C,S), "
+ f"received {tuple(payload.shape)}"
+ )
+ if not payload.is_floating_point():
+ raise TypeError(f"LTX2 target audio expected floating waveform, received {payload.dtype}")
+ _require_finite_tensor(payload, "LTX2 target audio")
+ source_sample_rate = _positive_int(source_sample_rate, "target audio sample_rate")
+ duration_seconds = _positive_real(duration_seconds, "target AV duration")
+ source_samples = int(duration_seconds * source_sample_rate)
+ if source_samples < 1:
+ raise ValueError("LTX2 target AV duration resolves to fewer than one source audio sample")
+ source_waveform = payload.detach().to(device="cpu", dtype=torch.float32)[:, :source_samples]
+ waveform = convert_audio(
+ source_waveform,
+ from_rate=source_sample_rate,
+ to_rate=geometry.sample_rate,
+ to_channels=geometry.waveform_channels,
+ )
+ if waveform.shape[-1] >= geometry.target_samples:
+ waveform = waveform[:, : geometry.target_samples]
+ else:
+ waveform = torch.nn.functional.pad(
+ waveform,
+ (0, geometry.target_samples - waveform.shape[-1]),
+ )
+ return waveform.contiguous()
+
+
+def ltx2_log_mel_spectrogram(
+ waveforms: torch.Tensor,
+ *,
+ sample_rate: int,
+ hop_length: int,
+ mel_bins: int,
+) -> torch.Tensor:
+ """Apply Lightricks' official magnitude Slaney log-mel frontend.
+
+ The frontend is intentionally independent from the vocoder's inverse-STFT
+ helpers: those use a different FFT and hop for bandwidth extension and are
+ not the AudioVAE training representation.
+ """
+ if not isinstance(waveforms, torch.Tensor) or waveforms.ndim != 3:
+ raise ValueError(
+ "LTX2 log-mel frontend expected waveform tensor shaped (B,C,S), "
+ f"received {type(waveforms).__name__}/{getattr(waveforms, 'shape', None)}"
+ )
+ if not waveforms.is_floating_point():
+ raise TypeError(
+ f"LTX2 log-mel frontend expected floating waveform, received {waveforms.dtype}"
+ )
+ _require_finite_tensor(waveforms, "LTX2 log-mel waveform")
+ sample_rate = _positive_int(sample_rate, "LTX2 log-mel sample_rate")
+ hop_length = _positive_int(hop_length, "LTX2 log-mel hop_length")
+ mel_bins = _positive_int(mel_bins, "LTX2 log-mel mel_bins")
+ if waveforms.shape[-1] <= 512:
+ raise ValueError(
+ "LTX2 official centered 1024-point log-mel frontend requires more than "
+ f"512 waveform samples for reflect padding, received {waveforms.shape[-1]}"
+ )
+ mel_spectrogram = getattr(getattr(torchaudio, "transforms", None), "MelSpectrogram", None)
+ if not callable(mel_spectrogram):
+ raise RuntimeError(
+ "LTX2 target audio encoding requires torchaudio.transforms.MelSpectrogram"
+ )
+ frontend = mel_spectrogram(
+ sample_rate=sample_rate,
+ n_fft=1024,
+ win_length=1024,
+ hop_length=hop_length,
+ f_min=0.0,
+ f_max=sample_rate / 2,
+ n_mels=mel_bins,
+ window_fn=torch.hann_window,
+ center=True,
+ pad_mode="reflect",
+ power=1.0,
+ norm="slaney",
+ mel_scale="slaney",
+ ).to(device=waveforms.device, dtype=torch.float32)
+ magnitude_mel = frontend(waveforms.to(torch.float32))
+ log_mel = magnitude_mel.clamp_min(1e-5).log().permute(0, 1, 3, 2).contiguous()
+ _require_finite_tensor(log_mel, "LTX2 log-mel spectrogram")
+ return log_mel
+
+
+def encode_ltx2_target_video(
+ adapter: Any,
+ videos: list[np.ndarray],
+ geometry: LTX2VideoGeometry,
+) -> torch.Tensor:
+ """Preprocess videos and take the deterministic VideoVAE posterior mode."""
+ pixels = adapter.pipeline.video_processor.preprocess_video(
+ videos,
+ height=geometry.height,
+ width=geometry.width,
+ )
+ if not isinstance(pixels, torch.Tensor):
+ raise TypeError(
+ "LTX2 video_processor.preprocess_video must return torch.Tensor, "
+ f"received {type(pixels).__name__}"
+ )
+ expected_shape = (len(videos), 3, geometry.num_frames, geometry.height, geometry.width)
+ if tuple(pixels.shape) != expected_shape:
+ raise ValueError(
+ "LTX2 target video preprocessing changed configured geometry: "
+ f"expected {expected_shape}, received {tuple(pixels.shape)}"
+ )
+ if not pixels.is_floating_point():
+ raise TypeError(
+ f"LTX2 target video preprocessing must return floating pixels, got {pixels.dtype}"
+ )
+ _require_finite_tensor(pixels, "LTX2 target video pixels")
+ vae = adapter.get_component("vae")
+ encoded = vae.encode(
+ pixels.to(
+ device=adapter.device,
+ dtype=_floating_module_dtype(vae, "LTX2 VideoVAE"),
+ )
+ )
+ latents = retrieve_vae_latents(
+ encoded,
+ sample_mode="argmax",
+ source="LTX2 target video",
+ ).to(device=adapter.device, dtype=torch.float32)
+ expected_latent_shape = (
+ len(videos),
+ geometry.latent_channels,
+ geometry.latent_frames,
+ geometry.latent_height,
+ geometry.latent_width,
+ )
+ if tuple(latents.shape) != expected_latent_shape:
+ raise ValueError(
+ "LTX2 target VideoVAE latent geometry mismatch: "
+ f"expected {expected_latent_shape}, received {tuple(latents.shape)}"
+ )
+ _require_finite_tensor(latents, "LTX2 target video latents")
+ return latents
+
+
+def normalize_and_pack_ltx2_video(
+ adapter: Any,
+ latents: torch.Tensor,
+ geometry: LTX2VideoGeometry,
+) -> torch.Tensor:
+ """Apply Diffusers' LTX2 video normalization and patch packing."""
+ vae = adapter.get_component("vae")
+ scaling_factor = _positive_real(
+ getattr(getattr(vae, "config", None), "scaling_factor", None),
+ "vae.config.scaling_factor",
+ )
+ normalized = adapter.pipeline._normalize_latents(
+ latents.to(torch.float32),
+ vae.latents_mean,
+ vae.latents_std,
+ scaling_factor,
+ )
+ packed = adapter.pipeline._pack_latents(
+ normalized,
+ geometry.patch_size,
+ geometry.patch_size_t,
+ ).to(device=adapter.device, dtype=torch.float32)
+ expected_shape = (
+ latents.shape[0],
+ geometry.sequence_length,
+ geometry.feature_dim,
+ )
+ if tuple(packed.shape) != expected_shape:
+ raise ValueError(
+ "LTX2 packed target video geometry mismatch: "
+ f"expected {expected_shape}, received {tuple(packed.shape)}"
+ )
+ _require_finite_tensor(packed, "LTX2 packed target video")
+ return packed
+
+
+def encode_ltx2_target_audio(
+ adapter: Any,
+ waveforms: torch.Tensor,
+ geometry: LTX2AudioGeometry,
+) -> torch.Tensor:
+ """Apply the official frontend, AudioVAE mode, conformance, packing, and normalization."""
+ expected_waveform_shape = (
+ waveforms.shape[0] if waveforms.ndim == 3 else None,
+ geometry.waveform_channels,
+ geometry.target_samples,
+ )
+ if waveforms.ndim != 3 or tuple(waveforms.shape) != expected_waveform_shape:
+ raise ValueError(
+ "LTX2 target waveform batch geometry mismatch: "
+ f"expected {expected_waveform_shape}, received {tuple(waveforms.shape)}"
+ )
+ log_mel = ltx2_log_mel_spectrogram(
+ waveforms.to(device="cpu", dtype=torch.float32),
+ sample_rate=geometry.sample_rate,
+ hop_length=geometry.hop_length,
+ mel_bins=geometry.mel_bins,
+ )
+ audio_vae = adapter.get_component("audio_vae")
+ encoded = audio_vae.encode(
+ log_mel.to(
+ device=adapter.device,
+ dtype=_floating_module_dtype(audio_vae, "LTX2 AudioVAE"),
+ )
+ )
+ latents = retrieve_vae_latents(
+ encoded,
+ sample_mode="argmax",
+ source="LTX2 target audio",
+ ).to(device=adapter.device, dtype=torch.float32)
+ if (
+ latents.ndim != 4
+ or latents.shape[0] != waveforms.shape[0]
+ or latents.shape[1] != geometry.latent_channels
+ or latents.shape[3] != geometry.latent_mel_bins
+ ):
+ raise ValueError(
+ "LTX2 target AudioVAE latent geometry mismatch: expected "
+ f"(B={waveforms.shape[0]}, C={geometry.latent_channels}, T, "
+ f"M={geometry.latent_mel_bins}), received {tuple(latents.shape)}"
+ )
+ if latents.shape[2] >= geometry.latent_frames:
+ latents = latents[:, :, : geometry.latent_frames]
+ else:
+ latents = torch.nn.functional.pad(
+ latents,
+ (0, 0, 0, geometry.latent_frames - latents.shape[2]),
+ )
+ packed = adapter.pipeline._pack_audio_latents(latents)
+ normalized = adapter.pipeline._normalize_audio_latents(
+ packed,
+ audio_vae.latents_mean,
+ audio_vae.latents_std,
+ ).to(device=adapter.device, dtype=torch.float32)
+ expected_shape = (
+ waveforms.shape[0],
+ geometry.sequence_length,
+ geometry.feature_dim,
+ )
+ if tuple(normalized.shape) != expected_shape:
+ raise ValueError(
+ "LTX2 packed target audio geometry mismatch: "
+ f"expected {expected_shape}, received {tuple(normalized.shape)}"
+ )
+ _require_finite_tensor(normalized, "LTX2 packed target audio")
+ return normalized
+
+
+def validate_ltx2_encoded_output_geometry(
+ adapter: Any,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+ *,
+ conditioned: bool,
+) -> None:
+ """Validate self-reported codec geometry against runtime configs and I2AV binding."""
+ geometry = resolve_ltx2_output_geometry(adapter, conditioned=conditioned)
+ batch_size = len(media_batch)
+ expected_video_shape = (
+ batch_size,
+ geometry.video.sequence_length,
+ geometry.video.feature_dim,
+ )
+ expected_audio_shape = (
+ batch_size,
+ geometry.audio.sequence_length,
+ geometry.audio.feature_dim,
+ )
+ video = encoded.clean_state.components["video"]
+ audio = encoded.clean_state.components["audio"]
+ if tuple(video.shape) != expected_video_shape or tuple(audio.shape) != expected_audio_shape:
+ raise ValueError(
+ "LTX2 encoded component geometry mismatch: expected "
+ f"video={expected_video_shape}, audio={expected_audio_shape}; received "
+ f"video={tuple(video.shape)}, audio={tuple(audio.shape)}"
+ )
+
+ expected_forward = _ltx2_forward_context(geometry)
+ for key, value in expected_forward.items():
+ if encoded.forward_context.get(key) != value:
+ raise ValueError(
+ f"LTX2 encoded forward_context[{key!r}] mismatch: "
+ f"expected {value!r}, received {encoded.forward_context.get(key)!r}"
+ )
+ expected_decode = {
+ "height": geometry.video.height,
+ "width": geometry.video.width,
+ "num_frames": geometry.video.num_frames,
+ "frame_rate": geometry.video.frame_rate,
+ }
+ for key, value in expected_decode.items():
+ if encoded.decode_context.get(key) != value:
+ raise ValueError(
+ f"LTX2 encoded decode_context[{key!r}] mismatch: "
+ f"expected {value!r}, received {encoded.decode_context.get(key)!r}"
+ )
+
+ expected_signature = GeometrySignature(
+ media=(
+ MediaGeometrySignature(
+ type=MediaType.VIDEO,
+ height=geometry.video.height,
+ width=geometry.video.width,
+ frames=geometry.video.num_frames,
+ fps=geometry.video.frame_rate,
+ ),
+ MediaGeometrySignature(
+ type=MediaType.AUDIO,
+ samples=geometry.audio.target_samples,
+ sample_rate=geometry.audio.sample_rate,
+ ),
+ )
+ )
+ if encoded.geometry_signatures != tuple(expected_signature for _ in media_batch):
+ raise ValueError("LTX2 encoded geometry signatures do not match the configured AV geometry")
+
+ if not conditioned:
+ if encoded.clean_state.active_masks is not None:
+ raise ValueError("LTX2 T2AV clean targets must not carry active masks")
+ return
+
+ condition_latents = condition.get("condition_video_latents")
+ expected_condition_shape = (
+ batch_size,
+ geometry.video.latent_channels,
+ 1,
+ geometry.video.latent_height,
+ geometry.video.latent_width,
+ )
+ if (
+ not isinstance(condition_latents, torch.Tensor)
+ or tuple(condition_latents.shape) != expected_condition_shape
+ ):
+ raise ValueError(
+ "LTX2 I2AV encoded output is not bound to the prepared first-frame condition: "
+ f"expected {expected_condition_shape}, received "
+ f"{getattr(condition_latents, 'shape', None)}"
+ )
+ conditioning_mask = encoded.forward_context.get("conditioning_mask")
+ if not isinstance(conditioning_mask, torch.Tensor) or tuple(conditioning_mask.shape) != (
+ batch_size,
+ geometry.video.sequence_length,
+ ):
+ raise ValueError(
+ "LTX2 I2AV conditioning_mask geometry mismatch: expected "
+ f"{(batch_size, geometry.video.sequence_length)}, received "
+ f"{getattr(conditioning_mask, 'shape', None)}"
+ )
+ first_frame_tokens = (
+ geometry.video.latent_height
+ // geometry.video.patch_size
+ * (geometry.video.latent_width // geometry.video.patch_size)
+ )
+ expected_conditioning_mask = torch.zeros_like(conditioning_mask, dtype=torch.bool)
+ expected_conditioning_mask[:, :first_frame_tokens] = True
+ if not torch.equal(conditioning_mask.bool(), expected_conditioning_mask):
+ raise ValueError(
+ "LTX2 I2AV conditioning_mask must pin exactly the first latent video frame"
+ )
+ active_masks = encoded.clean_state.active_masks
+ if active_masks is None:
+ raise ValueError("LTX2 I2AV clean targets require video/audio active masks")
+ if not torch.equal(
+ active_masks["video"].reshape_as(conditioning_mask),
+ ~expected_conditioning_mask,
+ ):
+ raise ValueError("LTX2 I2AV video active mask must be the inverse conditioning mask")
+ if not bool(active_masks["audio"].all()):
+ raise ValueError("LTX2 I2AV audio target must remain fully active")
+
+
+def reduce_ltx2_flow_matching_objective_values(
+ adapter: Any,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState],
+) -> torch.Tensor:
+ """Sum official per-modality means only for the offline flow objective.
+
+ LTX2 trains the video and audio denoisers as two equally weighted terms.
+ Keeping this policy behind the flow-matching-specific adapter hook preserves
+ the element-weighted joint likelihood used by online RL and distillation.
+ """
+ component_means = adapter.reduce_component_latent_values(values, state=state)
+ return component_means["video"] + component_means["audio"]
+
+
+def decode_ltx2_output_state(
+ adapter: Any,
+ encoded: EncodedOutputState,
+ *,
+ output_type: Literal["pil", "pt", "np"],
+) -> Any:
+ """Route the shared video/audio state through the existing LTX2 decoder."""
+ if encoded.clean_state.component_names != ("video", "audio"):
+ raise ValueError(
+ "LTX2 output decoding requires component order ('video', 'audio'), "
+ f"received {encoded.clean_state.component_names}"
+ )
+ context = encoded.decode_context
+ return adapter.decode_latents(
+ encoded.clean_state.components["video"],
+ encoded.clean_state.components["audio"],
+ height=context["height"],
+ width=context["width"],
+ num_frames=context["num_frames"],
+ frame_rate=context["frame_rate"],
+ output_type=output_type,
+ )
+
+
+def _ltx2_forward_context(geometry: LTX2OutputGeometry) -> dict[str, Any]:
+ return {
+ "height": geometry.video.height,
+ "width": geometry.video.width,
+ "num_frames": geometry.video.num_frames,
+ "frame_rate": geometry.video.frame_rate,
+ "video_seq_len": geometry.video.sequence_length,
+ "audio_num_frames": geometry.audio.latent_frames,
+ }
+
+
+def _component_config(adapter: Any, name: str) -> Any:
+ getter = getattr(adapter, "get_component_config", None)
+ if callable(getter):
+ return getter(name)
+ component = adapter.get_component(name)
+ config = getattr(component, "config", None)
+ if config is None:
+ raise TypeError(f"LTX2 component {name!r} does not expose config")
+ return config
+
+
+def _floating_module_dtype(module: Any, identifier: str) -> torch.dtype:
+ dtype = getattr(module, "dtype", None)
+ if isinstance(dtype, torch.dtype) and dtype.is_floating_point:
+ return dtype
+ parameters = getattr(module, "parameters", None)
+ if callable(parameters):
+ first = next(iter(parameters()), None)
+ if isinstance(first, torch.Tensor) and first.dtype.is_floating_point:
+ return first.dtype
+ raise TypeError(f"{identifier} must expose a floating dtype, received {dtype!r}")
+
+
+def _positive_int(value: Any, identifier: str) -> int:
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise TypeError(
+ f"expected positive int for {identifier}, received "
+ f"{type(value).__name__}: {value!r}"
+ )
+ if value <= 0:
+ raise ValueError(f"expected positive int for {identifier}, received {value}")
+ return value
+
+
+def _positive_real(value: Any, identifier: str) -> float:
+ if isinstance(value, bool) or not isinstance(value, Real):
+ raise TypeError(
+ f"expected positive real for {identifier}, received "
+ f"{type(value).__name__}: {value!r}"
+ )
+ value = float(value)
+ if not math.isfinite(value) or value <= 0:
+ raise ValueError(f"expected positive finite real for {identifier}, received {value!r}")
+ return value
+
+
+def _require_finite_tensor(value: torch.Tensor, identifier: str) -> None:
+ if not bool(torch.isfinite(value).all()):
+ raise ValueError(f"{identifier} contains non-finite values")
+
+
+__all__ = [
+ "LTX2AVOutputCodec",
+ "LTX2FirstFrameConditionPreparer",
+ "LTX2OutputGeometry",
+ "LTX2_OFFLINE_FORWARD_OVERRIDES",
+ "decode_ltx2_output_state",
+ "ltx2_log_mel_spectrogram",
+ "reduce_ltx2_flow_matching_objective_values",
+ "resolve_ltx2_output_geometry",
+ "resolve_ltx2_video_geometry",
+ "validate_ltx2_encoded_output_geometry",
+]
diff --git a/src/flow_factory/models/ltx2/ltx2_i2av.py b/src/flow_factory/models/ltx2/ltx2_i2av.py
index 9fa9261df..fdcba6d65 100644
--- a/src/flow_factory/models/ltx2/ltx2_i2av.py
+++ b/src/flow_factory/models/ltx2/ltx2_i2av.py
@@ -26,6 +26,12 @@
)
from PIL import Image
+from ...contracts import (
+ BatchCapability,
+ GeometrySource,
+ InputMediaRule,
+ NegativePromptPolicy,
+)
from ...hparams import *
from ...samples import (
ComponentTimes,
@@ -58,6 +64,8 @@
create_trajectory_collector,
)
from ..abc import BaseAdapter
+from ..output_state import DecodedMediaBatch, EncodedOutputState
+from ..pipeline_contracts import IMAGE_FORMAT, audio_video_output_contract
from ._common import (
LTX2_COMPONENT_ORDER,
attach_ltx2_state_masks,
@@ -73,6 +81,14 @@
split_ltx2_callback_results,
validate_i2av_forward_state_inputs,
)
+from ._output import (
+ LTX2_OFFLINE_FORWARD_OVERRIDES,
+ LTX2AVOutputCodec,
+ LTX2FirstFrameConditionPreparer,
+ decode_ltx2_output_state,
+ reduce_ltx2_flow_matching_objective_values,
+ validate_ltx2_encoded_output_geometry,
+)
logger = setup_logger(__name__)
@@ -177,11 +193,22 @@ class LTX2_I2AV_Adapter(BaseAdapter):
Shared logic with LTX2_T2AV_Adapter is handled via code duplication.
"""
- output_state_codec_unavailable_reason = (
- "LTX2 I2AV offline targets require paired video/audio decoding, exact duration "
- "alignment, and an active mask for the pinned first-frame condition; those lossless "
- "audiovisual output semantics are not yet implemented"
+ component_load_dtype_defaults = {"audio_vae": torch.float32}
+ pipeline_io_contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ input_rules=(
+ InputMediaRule(
+ format=IMAGE_FORMAT,
+ min_count=1,
+ max_count=1,
+ slots=("first_frame",),
+ required_slots=("first_frame",),
+ ),
+ ),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.UNIFORM,
)
+ offline_training_forward_overrides = LTX2_OFFLINE_FORWARD_OVERRIDES
supports_diffusers_cache = True
trajectory_component_order: ClassVar[Tuple[str, ...]] = LTX2_COMPONENT_ORDER
@@ -270,6 +297,51 @@ def inference_modules(self) -> List[str]:
"""Components needed during inference and training forward."""
return ["transformer", "vae", "audio_vae", "connectors", "vocoder"]
+ def build_condition_state_preparer(self) -> LTX2FirstFrameConditionPreparer:
+ """Declare one-time deterministic first-frame condition encoding."""
+ return LTX2FirstFrameConditionPreparer(self)
+
+ def build_output_state_codec(self) -> LTX2AVOutputCodec:
+ """Declare deterministic on-the-fly encoding for paired AV targets."""
+ return LTX2AVOutputCodec(self, conditioned=True)
+
+ def _validate_encoded_output_geometry(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+ ) -> None:
+ """Validate AV geometry and the pinned first-frame binding."""
+ validate_ltx2_encoded_output_geometry(
+ self,
+ media_batch,
+ condition,
+ encoded,
+ conditioned=True,
+ )
+
+ def _decode_output_state(
+ self,
+ encoded: EncodedOutputState,
+ *,
+ output_type: Literal["pil", "pt", "np"],
+ ) -> Any:
+ """Decode both components through the existing joint AV decoder."""
+ return decode_ltx2_output_state(self, encoded, output_type=output_type)
+
+ def _reduce_flow_matching_objective_values(
+ self,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState] = None,
+ ) -> torch.Tensor:
+ """Sum official video/audio means without changing online reducers."""
+ return reduce_ltx2_flow_matching_objective_values(
+ self,
+ values,
+ state=state,
+ )
+
def _check_inputs(
self,
height: int,
@@ -586,6 +658,7 @@ def encode_video(self, videos, **kwargs):
def preprocess_func(
self,
prompt: List[str],
+ negative_prompt: Optional[Union[str, List[str]]] = None,
images: Optional[List[Union[Image.Image, List[Image.Image]]]] = None,
system_prompt: Optional[str] = None,
prompt_enhancement_seed: int = 10,
@@ -613,6 +686,7 @@ def preprocess_func(
batch = self.encode_prompt(
prompt=prompt,
+ negative_prompt=negative_prompt,
guidance_scale=guidance_scale,
audio_guidance_scale=audio_guidance_scale,
max_sequence_length=max_sequence_length,
@@ -771,6 +845,7 @@ def forward(
compute_log_prob: bool = True,
return_kwargs: List[str] = ["next_latents", "log_prob", "velocity"],
use_cross_timestep: bool = False,
+ preserve_raw_model_velocity: bool = False,
# Component-return mode, owned by ``_forward_state``
_return_components: bool = False,
**kwargs,
@@ -929,6 +1004,15 @@ def forward(
)
video_pred = video_pred.float()
audio_pred = audio_pred.float()
+ raw_video_pred = video_pred
+ raw_audio_pred = audio_pred
+ uses_x0_guidance = (
+ do_cfg
+ or do_stg
+ or do_modality_isolation
+ or guidance_rescale > 0
+ or audio_guidance_rescale > 0
+ )
# --- 2. Convert to x0-space and compute guidance deltas ---
if do_cfg:
@@ -1022,9 +1106,35 @@ def forward(
audio_x0_guided, audio_x0, guidance_rescale=audio_guidance_rescale
)
- # --- 6. Convert back to velocity for scheduler step ---
- video_pred = self.convert_x0_to_velocity(video_latents, video_x0_guided, sigma)
- audio_pred = self.convert_x0_to_velocity(audio_latents, audio_x0_guided, sigma)
+ # --- 6. Convert back only when an x0-space guidance transform ran. ---
+ # The algebraic round trip is numerically unstable near sigma=0 and is
+ # unnecessary for SFT/offline-DPO's neutral-guidance forward.
+ if uses_x0_guidance or not preserve_raw_model_velocity:
+ video_pred = self.convert_x0_to_velocity(
+ video_latents,
+ video_x0_guided,
+ sigma,
+ )
+ audio_pred = self.convert_x0_to_velocity(
+ audio_latents,
+ audio_x0_guided,
+ sigma,
+ )
+ else:
+ video_pred = raw_video_pred
+ audio_pred = raw_audio_pred
+
+ velocity_only = (
+ not compute_log_prob and next_latents is None and tuple(return_kwargs) == ("velocity",)
+ )
+ if velocity_only:
+ if _return_components:
+ return MultiModalStepOutput(
+ velocity=LatentState({"video": video_pred, "audio": audio_pred})
+ )
+ return FlowMatchEulerDiscreteSDESchedulerOutput(
+ velocity=torch.cat([video_pred, audio_pred], dim=1)
+ )
# --- 7. [I2AV] Video scheduler step with frame-slicing ---
if conditioning_mask is not None:
diff --git a/src/flow_factory/models/ltx2/ltx2_t2av.py b/src/flow_factory/models/ltx2/ltx2_t2av.py
index 66751c7fa..13e4e6fed 100644
--- a/src/flow_factory/models/ltx2/ltx2_t2av.py
+++ b/src/flow_factory/models/ltx2/ltx2_t2av.py
@@ -16,12 +16,13 @@
from __future__ import annotations
from dataclasses import dataclass
-from typing import Any, ClassVar, Dict, List, Mapping, Optional, Tuple, Union
+from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Tuple, Union
import torch
from accelerate import Accelerator
from diffusers.pipelines.ltx2.pipeline_ltx2 import LTX2Pipeline, rescale_noise_cfg
+from ...contracts import BatchCapability, GeometrySource, NegativePromptPolicy
from ...hparams import *
from ...samples import (
ComponentTimes,
@@ -46,6 +47,8 @@
create_trajectory_collector,
)
from ..abc import BaseAdapter
+from ..output_state import DecodedMediaBatch, EncodedOutputState
+from ..pipeline_contracts import audio_video_output_contract
from ._common import (
LTX2_COMPONENT_ORDER,
attach_ltx2_state_masks,
@@ -61,6 +64,13 @@
split_ltx2_callback_results,
validate_ltx2_forward_state_inputs,
)
+from ._output import (
+ LTX2_OFFLINE_FORWARD_OVERRIDES,
+ LTX2AVOutputCodec,
+ decode_ltx2_output_state,
+ reduce_ltx2_flow_matching_objective_values,
+ validate_ltx2_encoded_output_geometry,
+)
logger = setup_logger(__name__)
@@ -162,11 +172,13 @@ class LTX2_T2AV_Adapter(BaseAdapter):
log_probs is the joint policy log_prob that drives policy gradient training.
"""
- output_state_codec_unavailable_reason = (
- "LTX2 offline targets require paired video/audio decoding with detected sample rates, "
- "official mel preprocessing, and exact duration alignment; the offline data plane "
- "does not yet provide that lossless audiovisual boundary"
+ component_load_dtype_defaults = {"audio_vae": torch.float32}
+ pipeline_io_contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.UNIFORM,
)
+ offline_training_forward_overrides = LTX2_OFFLINE_FORWARD_OVERRIDES
supports_diffusers_cache = True
trajectory_component_order: ClassVar[Tuple[str, ...]] = LTX2_COMPONENT_ORDER
@@ -265,6 +277,47 @@ def inference_modules(self) -> List[str]:
"""Components needed during inference and training forward."""
return ["transformer", "vae", "audio_vae", "connectors", "vocoder"]
+ def build_output_state_codec(self) -> LTX2AVOutputCodec:
+ """Declare deterministic on-the-fly encoding for paired AV targets."""
+ return LTX2AVOutputCodec(self)
+
+ def _validate_encoded_output_geometry(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+ ) -> None:
+ """Validate target state geometry against the configured LTX2 clocks."""
+ validate_ltx2_encoded_output_geometry(
+ self,
+ media_batch,
+ condition,
+ encoded,
+ conditioned=False,
+ )
+
+ def _decode_output_state(
+ self,
+ encoded: EncodedOutputState,
+ *,
+ output_type: Literal["pil", "pt", "np"],
+ ) -> Any:
+ """Decode both components through the existing joint AV decoder."""
+ return decode_ltx2_output_state(self, encoded, output_type=output_type)
+
+ def _reduce_flow_matching_objective_values(
+ self,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState] = None,
+ ) -> torch.Tensor:
+ """Sum official video/audio means without changing online reducers."""
+ return reduce_ltx2_flow_matching_objective_values(
+ self,
+ values,
+ state=state,
+ )
+
# ============================== Input Validation ==============================
def _check_inputs(
@@ -712,6 +765,7 @@ def forward(
return_kwargs: List[str] = ["next_latents", "log_prob", "velocity"],
# LTX-2.3 compatibility
use_cross_timestep: bool = False,
+ preserve_raw_model_velocity: bool = False,
# Component-return mode, owned by ``_forward_state``
_return_components: bool = False,
**kwargs,
@@ -873,6 +927,15 @@ def forward(
)
video_pred = video_pred.float()
audio_pred = audio_pred.float()
+ raw_video_pred = video_pred
+ raw_audio_pred = audio_pred
+ uses_x0_guidance = (
+ do_cfg
+ or do_stg
+ or do_modality_isolation
+ or guidance_rescale > 0
+ or audio_guidance_rescale > 0
+ )
# --- 2. Convert to x0-space and compute guidance deltas (pipeline L1250-1400) ---
if do_cfg:
@@ -959,9 +1022,35 @@ def forward(
audio_x0_guided, audio_x0, guidance_rescale=audio_guidance_rescale
)
- # --- 7. Convert back to velocity for scheduler step ---
- video_pred = self.convert_x0_to_velocity(video_latents, video_x0_guided, sigma)
- audio_pred = self.convert_x0_to_velocity(audio_latents, audio_x0_guided, sigma)
+ # --- 7. Convert back only when an x0-space guidance transform ran. ---
+ # The algebraic round trip is numerically unstable near sigma=0 and is
+ # unnecessary for SFT/offline-DPO's neutral-guidance forward.
+ if uses_x0_guidance or not preserve_raw_model_velocity:
+ video_pred = self.convert_x0_to_velocity(
+ video_latents,
+ video_x0_guided,
+ sigma,
+ )
+ audio_pred = self.convert_x0_to_velocity(
+ audio_latents,
+ audio_x0_guided,
+ sigma,
+ )
+ else:
+ video_pred = raw_video_pred
+ audio_pred = raw_audio_pred
+
+ velocity_only = (
+ not compute_log_prob and next_latents is None and tuple(return_kwargs) == ("velocity",)
+ )
+ if velocity_only:
+ if _return_components:
+ return MultiModalStepOutput(
+ velocity=LatentState({"video": video_pred, "audio": audio_pred})
+ )
+ return FlowMatchEulerDiscreteSDESchedulerOutput(
+ velocity=torch.cat([video_pred, audio_pred], dim=1)
+ )
# --- 8. Video: SDE scheduler step (with log_prob) ---
video_output = self.scheduler.step(
diff --git a/src/flow_factory/models/minimax_h3/__init__.py b/src/flow_factory/models/minimax_h3/__init__.py
index 2d24d6bbf..bf823ec55 100644
--- a/src/flow_factory/models/minimax_h3/__init__.py
+++ b/src/flow_factory/models/minimax_h3/__init__.py
@@ -37,7 +37,12 @@
MiniMaxH3Ref2VAAdapter,
MiniMaxH3T2VAAdapter,
)
-from .blocks import encode_h3_workflow_inputs, prepare_h3_rollout_state, run_h3_blocks
+from .blocks import (
+ encode_h3_workflow_inputs,
+ prepare_h3_condition_prefixes,
+ prepare_h3_rollout_state,
+ run_h3_blocks,
+)
from .decoding import decode_h3_targets
from .denoise import forward_h3_state, run_h3_joint_transformer, step_h3_components
from .dependency import require_minimax_h3_support
@@ -65,6 +70,7 @@
"model_time_to_framework_sigma",
"pack_audio_latents",
"pack_video_latents",
+ "prepare_h3_condition_prefixes",
"prepare_h3_rollout_state",
"require_minimax_h3_support",
"run_h3_blocks",
diff --git a/src/flow_factory/models/minimax_h3/_condition.py b/src/flow_factory/models/minimax_h3/_condition.py
new file mode 100644
index 000000000..04492d2a3
--- /dev/null
+++ b/src/flow_factory/models/minimax_h3/_condition.py
@@ -0,0 +1,99 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Runtime condition-prefix preparation for conditioned MiniMax H3 workflows."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, ClassVar, Optional, Tuple
+
+import torch
+
+from ..condition_state import PreparedConditionState
+from .blocks import prepare_h3_condition_prefixes
+from .workflow import _normalize_layout
+
+_RUNTIME_OWNED_FIELDS = frozenset(
+ {
+ "condition_latents",
+ "audio_condition_latents",
+ "condition_prefixes",
+ "layout",
+ "position_ids",
+ "token_tags",
+ "video_indices",
+ "audio_indices",
+ "text_indices",
+ "num_condition_video_rows",
+ "num_condition_audio_rows",
+ }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class MiniMaxH3ConditionStatePreparer:
+ """Realize one FL2VA/Ref2VA prefix and share it across offline arms."""
+
+ adapter: Any
+ required_components: ClassVar[Tuple[str, ...]] = ("scheduler",)
+
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ """Noise cached visual conditions once and bind canonical H3 layout.
+
+ Args:
+ condition: B=1 cached H3 condition with clean condition latents.
+ generator: Optional generator consumed in official packed condition order.
+
+ Returns:
+ Input-owned condition realization shared by every offline target arm.
+ """
+ workflow = getattr(self.adapter, "workflow", None)
+ if workflow not in ("fl2va", "ref2va"):
+ raise ValueError(
+ "MiniMax H3 condition preparer requires workflow 'fl2va' or 'ref2va', "
+ f"received {workflow!r}"
+ )
+ if "condition_prefixes" in condition:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} cached condition must not contain "
+ "already-realized condition_prefixes"
+ )
+
+ layout = _normalize_layout(condition)
+ prefixes = prepare_h3_condition_prefixes(
+ self.adapter.pipeline,
+ condition,
+ workflow=workflow,
+ generator=generator,
+ )
+ static_condition = {
+ key: value for key, value in condition.items() if key not in _RUNTIME_OWNED_FIELDS
+ }
+ return PreparedConditionState(
+ condition=static_condition,
+ forward_context={
+ "condition_prefixes": prefixes,
+ "layout": layout,
+ },
+ output_context={"layout": layout},
+ )
+
+
+__all__ = ["MiniMaxH3ConditionStatePreparer"]
diff --git a/src/flow_factory/models/minimax_h3/_output.py b/src/flow_factory/models/minimax_h3/_output.py
index dff208c3f..c5a151f63 100644
--- a/src/flow_factory/models/minimax_h3/_output.py
+++ b/src/flow_factory/models/minimax_h3/_output.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""On-the-fly audiovisual target encoding for MiniMax H3 T2VA."""
+"""On-the-fly audiovisual target encoding for all MiniMax H3 workflows."""
from __future__ import annotations
@@ -75,7 +75,7 @@ class _H3ModelShape:
@dataclass(frozen=True, slots=True)
class MiniMaxH3AVOutputCodec:
- """Encode one configured T2VA video/audio target into packed H3 rows."""
+ """Encode one configured H3 video/audio target into packed target rows."""
adapter: Any
required_components: ClassVar[Tuple[str, ...]] = ("vae", "audio_vae")
@@ -89,29 +89,28 @@ def encode_output_state(
"""Encode the exact ``(video, audio)`` target sequence for one sample.
Diffusers defines H3's conditioning encoder, but not an offline target
- encoder. Video posterior sampling is therefore a framework policy inferred
- from latent-diffusion training objectives. In particular, the condition-only
- float16 rounding step is intentionally not applied to clean targets. Audio
- targets use the posterior mode because the official H3 pipeline never samples
- the audio posterior. Both components are normalized and packed exactly like
- generated target rows.
+ encoder. Both target posteriors therefore use their deterministic modes,
+ matching the released H3 SFT data flow. The condition-only float16 rounding
+ step is intentionally not applied to clean targets. Both components are
+ normalized and packed exactly like generated target rows.
Args:
media_batch: One decoded video/audio pair.
- condition: Cached T2VA condition and its authoritative packed layout.
- generator: Optional generator forwarded only to video posterior sampling.
+ condition: Cached H3 condition and its authoritative packed layout.
+ generator: Accepted for the shared codec protocol; H3 target modes are
+ deterministic and do not consume it.
Returns:
Detached clean video/audio rows plus the output-owned decode context.
"""
if len(media_batch) != 1:
raise ValueError(
- f"MiniMax H3 T2VA output encoding requires B=1, received B={len(media_batch)}"
+ f"MiniMax H3 output encoding requires B=1, received B={len(media_batch)}"
)
candidate = media_batch[0]
if len(candidate) != 2:
raise ValueError(
- "MiniMax H3 T2VA output codec expected exact (video, audio) media, "
+ "MiniMax H3 output codec expected exact (video, audio) media, "
f"received {len(candidate)} items"
)
video_media, audio_media = candidate
@@ -126,11 +125,8 @@ def encode_output_state(
height=geometry["height"],
width=geometry["width"],
)
- video_latents = encode_h3_target_video(
- self.adapter,
- video_pixels,
- generator=generator,
- )
+ del generator
+ video_latents = encode_h3_target_video(self.adapter, video_pixels)
expected_video_shape = (
1,
model_shape.video_latent_channels,
@@ -180,7 +176,7 @@ def encode_output_state(
}
)
validate_target_state(clean_state)
- _validate_h3_t2va_input_layout(condition, clean_state)
+ _validate_h3_input_layout(condition, clean_state)
frame_rate = float(self.adapter.pipeline.fps)
signature = GeometrySignature(
media=(
@@ -207,7 +203,7 @@ def encode_output_state(
def resolve_h3_output_geometry(adapter: Any, condition: Mapping[str, Any]) -> Dict[str, int]:
- """Resolve and validate configured geometry from the cached T2VA condition.
+ """Resolve and validate configured geometry from the cached H3 condition.
Args:
adapter: Active MiniMax H3 adapter with configured pipeline components.
@@ -458,22 +454,17 @@ def prepare_h3_target_video(
def encode_h3_target_video(
adapter: Any,
pixel_values: torch.Tensor,
- *,
- generator: Optional[torch.Generator],
) -> torch.Tensor:
- """Apply the framework's sampled-posterior policy for clean H3 targets.
+ """Encode clean H3 targets through the deterministic video posterior mode.
Diffusers specifies a fixed-seed sample followed by float16 rounding only for
- H3 *conditions*. Offline clean targets instead use the caller's generator and
- retain the sampled posterior in float32 before normalization. This distinction
- is an explicit framework inference from latent training, not an official H3
- target-encoding recipe.
+ H3 *conditions*. Its released pipeline has no offline target encoder, so this
+ path follows the H3 SFT reference data flow: take the posterior mean/mode and
+ retain float32 precision before normalization.
Args:
adapter: Active MiniMax H3 adapter exposing the video VAE.
pixel_values: Normalized-shape pixels ``(1, 3, F, H, W)`` in ``[0, 1]``.
- generator: Optional posterior sampling generator.
-
Returns:
Normalized video latents shaped ``(1, 24, F', H', W')``.
"""
@@ -500,8 +491,7 @@ def encode_h3_target_video(
encoded = vae.encode(normalized_pixels)
latents = retrieve_vae_latents(
encoded,
- sample_mode="sample",
- generator=generator,
+ sample_mode="argmax",
source="MiniMax H3 target video",
).to(torch.float32)
latent_mean, latent_std = _latent_statistics(
@@ -609,11 +599,11 @@ def encode_h3_target_audio(adapter: Any, waveform: torch.Tensor) -> torch.Tensor
return (latents - latent_mean) / latent_std
-def _validate_h3_t2va_input_layout(
+def _validate_h3_input_layout(
condition: Mapping[str, Any],
clean_state: LatentState,
) -> None:
- """Validate the authoritative flat T2VA layout retained in input conditions."""
+ """Validate target rows against an authoritative H3 input layout."""
layout = _normalize_layout(condition)
missing = tuple(field for field in _LAYOUT_FIELDS if field not in layout)
if missing:
@@ -629,7 +619,7 @@ def _validate_h3_t2va_input_layout(
or position_ids.dtype != torch.float64
):
raise ValueError(
- "MiniMax H3 T2VA position_ids expected float64 shape (N,3), "
+ "MiniMax H3 position_ids expected float64 shape (N,3), "
f"received {type(position_ids).__name__}/{getattr(position_ids, 'shape', None)}/"
f"{getattr(position_ids, 'dtype', None)}"
)
@@ -639,37 +629,39 @@ def _validate_h3_t2va_input_layout(
or token_tags.dtype != torch.long
):
raise ValueError(
- "MiniMax H3 T2VA token_tags expected one-dimensional torch.long, "
+ "MiniMax H3 token_tags expected one-dimensional torch.long, "
f"received {type(token_tags).__name__}/{getattr(token_tags, 'shape', None)}/"
f"{getattr(token_tags, 'dtype', None)}"
)
for field, values in zip(("video_indices", "audio_indices", "text_indices"), index_tensors):
if not isinstance(values, torch.Tensor) or values.ndim != 1 or values.dtype != torch.long:
raise ValueError(
- f"MiniMax H3 T2VA {field} expected one-dimensional torch.long, "
+ f"MiniMax H3 {field} expected one-dimensional torch.long, "
f"received {type(values).__name__}/{getattr(values, 'shape', None)}/"
f"{getattr(values, 'dtype', None)}"
)
for component in ("video", "audio"):
count_field = f"num_condition_{component}_rows"
- if layout[count_field] != 0:
+ condition_rows = layout[count_field]
+ if condition_rows < 0:
raise ValueError(
- "MiniMax H3 T2VA offline layout requires no condition rows, "
- f"received {count_field}={layout[count_field]}"
+ f"MiniMax H3 offline layout requires non-negative {count_field}, "
+ f"received {condition_rows}"
)
component_indices = layout[f"{component}_indices"]
- expected_rows = clean_state.components[component].shape[1]
- if component_indices.numel() != expected_rows:
+ target_rows = clean_state.components[component].shape[1]
+ expected_total_rows = condition_rows + target_rows
+ if component_indices.numel() != expected_total_rows:
raise ValueError(
- f"MiniMax H3 T2VA {component} layout expected {expected_rows} target rows, "
- f"received {component_indices.numel()} indices"
+ f"MiniMax H3 {component} layout expected {condition_rows} condition + "
+ f"{target_rows} target rows, received {component_indices.numel()} indices"
)
sequence_length = sum(values.numel() for values in index_tensors)
if position_ids.shape[0] != sequence_length or token_tags.numel() != sequence_length:
raise ValueError(
- "MiniMax H3 T2VA flat layout sequence lengths disagree: "
+ "MiniMax H3 flat layout sequence lengths disagree: "
f"indices={sequence_length}, position_ids={position_ids.shape[0]}, "
f"token_tags={token_tags.numel()}"
)
@@ -679,9 +671,7 @@ def _validate_h3_t2va_input_layout(
*(values.device for values in index_tensors),
}
if len(devices) != 1:
- raise ValueError(
- f"MiniMax H3 T2VA flat layout tensors must share one device, got {devices}"
- )
+ raise ValueError(f"MiniMax H3 flat layout tensors must share one device, got {devices}")
permutation = torch.cat(index_tensors).sort().values
expected_permutation = torch.arange(
sequence_length,
@@ -689,9 +679,7 @@ def _validate_h3_t2va_input_layout(
device=permutation.device,
)
if not torch.equal(permutation, expected_permutation):
- raise ValueError(
- "MiniMax H3 T2VA video/audio/text indices must partition the packed sequence"
- )
+ raise ValueError("MiniMax H3 video/audio/text indices must partition the packed sequence")
def validate_h3_encoded_output_geometry(
@@ -700,7 +688,7 @@ def validate_h3_encoded_output_geometry(
condition: Mapping[str, Any],
encoded: EncodedOutputState,
) -> None:
- """Prove codec geometry and contexts agree with the cached T2VA layout.
+ """Prove codec geometry and contexts agree with the cached H3 layout.
Args:
adapter: Active MiniMax H3 adapter.
@@ -781,10 +769,10 @@ def validate_h3_encoded_output_geometry(
if encoded.forward_context:
raise ValueError(
- "MiniMax H3 T2VA output codec must not duplicate input-owned layout/prefix fields, "
+ "MiniMax H3 output codec must not duplicate input-owned layout/prefix fields, "
f"received keys={tuple(encoded.forward_context)}"
)
- _validate_h3_t2va_input_layout(condition, encoded.clean_state)
+ _validate_h3_input_layout(condition, encoded.clean_state)
def _resolve_h3_model_shape(adapter: Any) -> _H3ModelShape:
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index 6c0a001c3..a9c7e74bc 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -21,7 +21,13 @@
from ...contracts import (
BatchCapability,
GeometrySource,
+ InputMediaBinding,
+ InputMediaOrder,
+ InputMediaRule,
+ MediaFormat,
+ MediaType,
NegativePromptPolicy,
+ RateRequirement,
)
from ...samples import (
ComponentTimes,
@@ -34,9 +40,14 @@
from ..abc import BaseAdapter
from ..checkpointing import CheckpointUnit
from ..output_state import DecodedMediaBatch, EncodedOutputState, OutputStateCodec
-from ..pipeline_contracts import audio_video_output_contract
+from ..pipeline_contracts import (
+ IMAGE_FORMAT,
+ VIDEO_FORMAT_OPTIONAL_FPS,
+ audio_video_output_contract,
+)
from ..runtime import ModularPipelineRuntime
from ._common import apply_forward_process_noise, draw_forward_process_noise
+from ._condition import MiniMaxH3ConditionStatePreparer
from ._output import MiniMaxH3AVOutputCodec, validate_h3_encoded_output_geometry
from .workflow import (
build_h3_component_runtime,
@@ -53,7 +64,12 @@
)
_H3_PREPROCESS_CACHE_FIELDS = frozenset({"height", "width", "num_frames"})
-_H3_PREPROCESS_CACHE_VERSION = "minimax-h3-v1"
+_H3_PREPROCESS_CACHE_VERSION = "minimax-h3-v2"
+_H3_OPTIONAL_AUDIO_REFERENCE_FORMAT = MediaFormat(
+ type=MediaType.AUDIO,
+ fps=RateRequirement.NOT_APPLICABLE,
+ sample_rate=RateRequirement.OPTIONAL,
+)
class _MiniMaxH3WorkflowAdapter:
@@ -132,6 +148,46 @@ def _init_target_module_map(self) -> Dict[str, Union[List[str], None]]:
def preprocess_func(self, **kwargs: Any) -> Dict[str, Any]:
return preprocess_h3_workflow(self, **kwargs)
+ def build_condition_state_preparer(
+ self,
+ ) -> Optional[MiniMaxH3ConditionStatePreparer]:
+ """Declare runtime prefix preparation only for conditioned workflows."""
+ if self.workflow == "t2va":
+ return None
+ return MiniMaxH3ConditionStatePreparer(self)
+
+ def build_output_state_codec(self) -> OutputStateCodec:
+ """Declare the shared audiovisual target codec without loading components."""
+ return MiniMaxH3AVOutputCodec(self)
+
+ def _validate_encoded_output_geometry(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+ ) -> None:
+ """Require encoded AV rows to match this workflow's cached layout."""
+ validate_h3_encoded_output_geometry(self, media_batch, condition, encoded)
+
+ def _decode_output_state(
+ self,
+ encoded: EncodedOutputState,
+ *,
+ output_type: Literal["pil", "pt", "np"],
+ ) -> Any:
+ """Decode both H3 target components through the existing decoder."""
+ geometry = encoded.decode_context.get("geometry")
+ if not isinstance(geometry, Mapping):
+ raise TypeError(
+ "MiniMax H3 decode_context requires a geometry mapping, "
+ f"received {type(geometry).__name__}: {geometry!r}"
+ )
+ return self.decode_latents(
+ encoded.clean_state,
+ geometry=geometry,
+ output_type=output_type,
+ )
+
def build_training_component_times(
self,
primary_timesteps: torch.Tensor,
@@ -157,6 +213,21 @@ def apply_forward_process_noise(
) -> NoisedState:
return apply_forward_process_noise(clean_state, times, noise)
+ def _reduce_flow_matching_objective_values(
+ self,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState] = None,
+ ) -> torch.Tensor:
+ """Sum the official video and audio means for offline flow matching.
+
+ H3's audiovisual objective gives each modality its own mean-squared-error
+ term. This objective-specific hook intentionally leaves the globally
+ element-weighted reducer used by online likelihood objectives unchanged.
+ """
+ component_means = self.reduce_component_latent_values(values, state=state)
+ return component_means["video"] + component_means["audio"]
+
def decode_latents(self, latents: Any, **kwargs: Any) -> Any:
return decode_h3_adapter_latents(self, latents, **kwargs)
@@ -221,50 +292,24 @@ class MiniMaxH3T2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"audio_vae",
]
- def build_output_state_codec(self) -> OutputStateCodec:
- """Declare the configured audiovisual target codec without loading components.
-
- Returns:
- Immutable MiniMax H3 audiovisual output codec declaration.
- """
- return MiniMaxH3AVOutputCodec(self)
-
- def _validate_encoded_output_geometry(
- self,
- media_batch: DecodedMediaBatch,
- condition: Mapping[str, Any],
- encoded: EncodedOutputState,
- ) -> None:
- """Require encoded rows and rate metadata to match cached T2VA geometry."""
- validate_h3_encoded_output_geometry(self, media_batch, condition, encoded)
-
- def _decode_output_state(
- self,
- encoded: EncodedOutputState,
- *,
- output_type: Literal["pil", "pt", "np"],
- ) -> Any:
- """Decode the two-component target state through H3's existing decoder."""
- geometry = encoded.decode_context.get("geometry")
- if not isinstance(geometry, Mapping):
- raise TypeError(
- "MiniMax H3 decode_context requires a geometry mapping, "
- f"received {type(geometry).__name__}: {geometry!r}"
- )
- return self.decode_latents(
- encoded.clean_state,
- geometry=geometry,
- output_type=output_type,
- )
-
class MiniMaxH3FL2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"""Load the workflow-pruned MiniMax H3 first/last-frame partition."""
- output_state_codec_unavailable_reason = (
- "MiniMax H3 conditioned offline forward requires a shared, reproducible "
- "conditioned-prefix binder; paired offline-DPO arms must reuse identical "
- "condition posterior noise"
+ pipeline_io_contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ input_rules=(
+ InputMediaRule(
+ format=IMAGE_FORMAT,
+ min_count=1,
+ max_count=2,
+ slots=("first_frame", "last_frame"),
+ ),
+ ),
+ input_binding=InputMediaBinding.GROUPED_BY_TYPE,
+ input_order=InputMediaOrder.WITHIN_TYPE,
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
)
workflow: ClassVar[str] = "fl2va"
@@ -287,10 +332,24 @@ class MiniMaxH3FL2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
class MiniMaxH3Ref2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"""Load the workflow-pruned MiniMax H3 omni-reference partition."""
- output_state_codec_unavailable_reason = (
- "MiniMax H3 conditioned offline forward requires a shared, reproducible "
- "conditioned-prefix binder; paired offline-DPO arms must reuse identical "
- "condition posterior noise"
+ pipeline_io_contract = audio_video_output_contract(
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ input_rules=(
+ InputMediaRule(format=IMAGE_FORMAT, min_count=0, max_count=9),
+ InputMediaRule(format=VIDEO_FORMAT_OPTIONAL_FPS, min_count=0, max_count=3),
+ InputMediaRule(
+ format=_H3_OPTIONAL_AUDIO_REFERENCE_FORMAT,
+ min_count=0,
+ max_count=3,
+ ),
+ ),
+ input_binding=InputMediaBinding.ORDERED_REFERENCES,
+ input_order=InputMediaOrder.GLOBAL,
+ min_input_media_count=1,
+ max_input_media_count=12,
+ required_any_input_types=(MediaType.IMAGE, MediaType.VIDEO),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
)
workflow: ClassVar[str] = "ref2va"
diff --git a/src/flow_factory/models/minimax_h3/blocks.py b/src/flow_factory/models/minimax_h3/blocks.py
index 0810123cd..908a769e5 100644
--- a/src/flow_factory/models/minimax_h3/blocks.py
+++ b/src/flow_factory/models/minimax_h3/blocks.py
@@ -41,21 +41,14 @@
)
ENCODE_COMMON_FIELDS = (
"prompt_embeds",
- "text_token_tags",
*GEOMETRY_FIELDS,
*LAYOUT_FIELDS,
)
ENCODE_WORKFLOW_FIELDS = {
- "t2va": (*ENCODE_COMMON_FIELDS, "keyframe_anchors"),
- "fl2va": (
- *ENCODE_COMMON_FIELDS,
- "keyframes",
- "keyframe_anchors",
- "condition_latents",
- ),
+ "t2va": ENCODE_COMMON_FIELDS,
+ "fl2va": (*ENCODE_COMMON_FIELDS, "condition_latents"),
"ref2va": (
*ENCODE_COMMON_FIELDS,
- "normalized_references",
"condition_latents",
"audio_condition_latents",
),
@@ -169,14 +162,36 @@ def prepare_h3_rollout_state(
Returns:
Target-only state and immutable video/audio prefixes.
"""
- symbols = require_minimax_h3_support()
_validate_workflow(workflow)
+ symbols = require_minimax_h3_support()
values = dict(cached_values)
values.update(generator=generator, latents=latents, audio_latents=audio_latents)
- block_types = []
- if workflow != "t2va":
- block_types.append(symbols.PrepareConditionLatentsStep)
- block_types.append(symbols.PrepareLatentsStep)
+ video_count = _condition_count(cached_values, workflow, "video")
+ audio_count = _condition_count(cached_values, workflow, "audio")
+ prefixes: Dict[str, torch.Tensor]
+ if workflow == "t2va":
+ if video_count or audio_count:
+ raise ValueError(
+ "MiniMax H3 workflow='t2va' requires zero condition rows, "
+ f"received video={video_count}, audio={audio_count}"
+ )
+ prefixes = {}
+ else:
+ prefixes = prepare_h3_condition_prefixes(
+ pipeline,
+ cached_values,
+ workflow=workflow,
+ generator=generator,
+ )
+ values.update(
+ condition_rows=prefixes["video"][0],
+ num_condition_video_rows=video_count,
+ num_condition_audio_rows=audio_count,
+ )
+ if workflow == "ref2va":
+ values["audio_condition_latents"] = [] if audio_count == 0 else [prefixes["audio"][0]]
+
+ block_types = [symbols.PrepareLatentsStep]
if workflow == "fl2va":
block_types.append(symbols.FL2VAPrepareLatentsStep)
elif workflow == "ref2va":
@@ -190,8 +205,6 @@ def prepare_h3_rollout_state(
)
video_rows = _as_batched_rows(output["latents"], 96, workflow, "latents")
audio_rows = _as_batched_rows(output["audio_latents"], 32, workflow, "audio_latents")
- video_count = _condition_count(cached_values, workflow, "video")
- audio_count = _condition_count(cached_values, workflow, "audio")
_validate_rollout_row_counts(
cached_values,
workflow,
@@ -219,6 +232,100 @@ def prepare_h3_rollout_state(
return targets, prefixes
+def prepare_h3_condition_prefixes(
+ pipeline: Any,
+ cached_values: Mapping[str, Any],
+ *,
+ workflow: str,
+ generator: Optional[torch.Generator] = None,
+) -> Dict[str, torch.Tensor]:
+ """Realize one immutable H3 condition prefix from cached clean latents.
+
+ The helper owns the only runtime condition-noise draw. Online rollout calls it
+ before generated video/audio noise, while offline SFT and DPO call it once per
+ batch and reuse the returned tensors for every target and policy/reference arm.
+
+ Args:
+ pipeline: Adapter-owned modular pipeline with the canonical scheduler.
+ cached_values: Cached layout and clean visual/audio condition latents.
+ workflow: Conditioned H3 workflow, ``fl2va`` or ``ref2va``.
+ generator: Optional request generator consumed in packed visual order.
+
+ Returns:
+ Batched immutable video/audio prefix rows in model packing order.
+ """
+ _validate_workflow(workflow)
+ if workflow == "t2va":
+ raise ValueError("MiniMax H3 T2VA has no runtime condition prefix to realize")
+
+ symbols = require_minimax_h3_support()
+ video_count = _condition_count(cached_values, workflow, "video")
+ audio_count = _condition_count(cached_values, workflow, "audio")
+ condition_latents = _normalize_condition_tensor_list(
+ cached_values.get("condition_latents"),
+ workflow=workflow,
+ field="condition_latents",
+ item_rank=5,
+ )
+ values = dict(cached_values)
+ values.update(
+ generator=generator,
+ num_condition_video_rows=video_count,
+ condition_latents=condition_latents,
+ )
+ output = run_h3_blocks(
+ pipeline,
+ [symbols.PrepareConditionLatentsStep()],
+ values,
+ requested_outputs=("condition_rows",),
+ workflow=workflow,
+ )
+ video_rows = _as_batched_rows(
+ output["condition_rows"],
+ 96,
+ workflow,
+ "condition_rows",
+ )
+ _validate_condition_prefix_rows(
+ cached_values,
+ workflow,
+ "video",
+ prefix_rows=video_rows.shape[1],
+ condition_row_count=video_count,
+ )
+
+ if workflow == "fl2va":
+ if audio_count:
+ raise ValueError(
+ "MiniMax H3 FL2VA requires zero audio condition rows, " f"received {audio_count}"
+ )
+ audio_rows = video_rows.new_empty((1, 0, 32))
+ else:
+ audio_latents = _normalize_condition_tensor_list(
+ cached_values.get("audio_condition_latents", []),
+ workflow=workflow,
+ field="audio_condition_latents",
+ item_rank=2,
+ )
+ if audio_latents:
+ audio_rows = _as_batched_rows(
+ torch.cat([rows.to(video_rows.device) for rows in audio_latents], dim=0),
+ 32,
+ workflow,
+ "audio_condition_latents",
+ )
+ else:
+ audio_rows = video_rows.new_empty((1, 0, 32))
+ _validate_condition_prefix_rows(
+ cached_values,
+ workflow,
+ "audio",
+ prefix_rows=audio_rows.shape[1],
+ condition_row_count=audio_count,
+ )
+ return {"video": video_rows, "audio": audio_rows}
+
+
def _validate_workflow(workflow: str) -> None:
if workflow not in WORKFLOWS:
raise ValueError(f"expected MiniMax H3 workflow in {WORKFLOWS}, received {workflow!r}")
@@ -227,6 +334,20 @@ def _validate_workflow(workflow: str) -> None:
def _condition_count(values: Mapping[str, Any], workflow: str, component: str) -> int:
field = f"num_condition_{component}_rows"
count = values.get(field, 0)
+ if isinstance(count, (list, tuple)):
+ if len(count) != 1:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r} expected B=1, "
+ f"received length={len(count)}"
+ )
+ count = count[0]
+ if isinstance(count, torch.Tensor):
+ if count.numel() != 1:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r} expected one scalar, "
+ f"received shape={tuple(count.shape)}"
+ )
+ count = count.item()
if not isinstance(count, int) or isinstance(count, bool) or count < 0:
raise ValueError(
f"MiniMax H3 workflow={workflow!r} field={field!r} expected non-negative int, "
@@ -244,13 +365,7 @@ def _validate_rollout_row_counts(
condition_row_count: int,
) -> None:
index_field = f"{component}_indices"
- indices = values.get(index_field)
- if not isinstance(indices, torch.Tensor) or indices.ndim != 1:
- raise ValueError(
- f"MiniMax H3 workflow={workflow!r} component={component!r} field={index_field!r} "
- f"actual={getattr(indices, 'shape', type(indices).__name__)} "
- "expected=one-dimensional Tensor"
- )
+ indices = _normalize_layout_indices(values.get(index_field), workflow, index_field)
expected_full_rows = indices.numel()
if full_row_count != expected_full_rows:
raise ValueError(
@@ -267,6 +382,97 @@ def _validate_rollout_row_counts(
)
+def _validate_condition_prefix_rows(
+ values: Mapping[str, Any],
+ workflow: str,
+ component: str,
+ *,
+ prefix_rows: int,
+ condition_row_count: int,
+) -> None:
+ index_field = f"{component}_indices"
+ indices = _normalize_layout_indices(values.get(index_field), workflow, index_field)
+ if prefix_rows != condition_row_count:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} component={component!r} "
+ f"field='condition_row_count' actual={prefix_rows} expected={condition_row_count}"
+ )
+ if condition_row_count > indices.numel():
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} component={component!r} "
+ f"field='condition_row_count' actual={condition_row_count} "
+ f"expected_at_most={indices.numel()}"
+ )
+
+
+def _normalize_layout_indices(
+ value: Any,
+ workflow: str,
+ field: str,
+) -> torch.Tensor:
+ if not isinstance(value, torch.Tensor):
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r} expected Tensor, "
+ f"received {type(value).__name__}"
+ )
+ if value.ndim == 2 and value.shape[0] == 1:
+ value = value[0]
+ if value.ndim != 1:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r} expected shape (N,) or "
+ f"collated (B=1,N), received {tuple(value.shape)}"
+ )
+ return value
+
+
+def _normalize_condition_tensor_list(
+ value: Any,
+ *,
+ workflow: str,
+ field: str,
+ item_rank: int,
+) -> list[torch.Tensor]:
+ """Normalize Arrow/HF B=1 list and stacked-tensor representations."""
+ if isinstance(value, tuple):
+ value = list(value)
+ if isinstance(value, list) and len(value) == 1 and isinstance(value[0], (list, tuple)):
+ value = list(value[0])
+
+ if isinstance(value, list) and len(value) == 1 and isinstance(value[0], torch.Tensor):
+ first = value[0]
+ if first.ndim == item_rank + 1:
+ value = first
+ if isinstance(value, torch.Tensor):
+ if value.ndim == item_rank:
+ tensors = [value]
+ else:
+ if value.ndim == item_rank + 2 and value.shape[0] == 1:
+ value = value[0]
+ if value.ndim != item_rank + 1:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r} expected one "
+ f"rank-{item_rank} tensor or a B=1/stacked container, received "
+ f"shape={tuple(value.shape)}"
+ )
+ tensors = list(torch.unbind(value, dim=0))
+ elif isinstance(value, list):
+ tensors = value
+ else:
+ raise TypeError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r} expected tensor/list, "
+ f"received {type(value).__name__}"
+ )
+
+ for index, tensor in enumerate(tensors):
+ if not isinstance(tensor, torch.Tensor) or tensor.ndim != item_rank:
+ raise ValueError(
+ f"MiniMax H3 workflow={workflow!r} field={field!r}[{index}] expected "
+ f"rank-{item_rank} Tensor, received "
+ f"{type(tensor).__name__}/{getattr(tensor, 'shape', None)}"
+ )
+ return tensors
+
+
def _as_batched_rows(rows: torch.Tensor, width: int, workflow: str, field: str) -> torch.Tensor:
if not isinstance(rows, torch.Tensor):
raise TypeError(
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index abc53f797..c0702da5d 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -171,10 +171,11 @@ def preprocess_h3_workflow(adapter: Any, **kwargs: Any) -> Dict[str, Any]:
"num_frames": kwargs["num_frames"],
}
if adapter.workflow == "fl2va":
- images = _validate_fl2va_condition_images(kwargs, "preprocess")
- values["image"] = images[0]
- if len(images) == 2:
- values["last_image"] = images[1]
+ first_image, last_image = _validate_fl2va_condition_images(kwargs, "preprocess")
+ if first_image is not None:
+ values["image"] = first_image
+ if last_image is not None:
+ values["last_image"] = last_image
elif adapter.workflow == "ref2va":
references = _single_outer_value(kwargs.get("references"), "references", adapter.workflow)
values["references"] = _build_pinned_references(references)
@@ -214,8 +215,18 @@ def infer_h3_workflow(adapter: Any, **kwargs: Any) -> List[Any]:
_validate_public_no_cfg_inputs(adapter.workflow, kwargs, "inference")
_validate_workflow_media_inputs(adapter.workflow, kwargs, "inference")
condition_images = None
+ condition_image_slots = None
if adapter.workflow == "fl2va":
- condition_images = _validate_fl2va_condition_images(kwargs, "inference")
+ first_image, last_image = _validate_fl2va_condition_images(kwargs, "inference")
+ condition_images = tuple(image for image in (first_image, last_image) if image is not None)
+ condition_image_slots = tuple(
+ slot
+ for slot, image in (
+ ("first_frame", first_image),
+ ("last_frame", last_image),
+ )
+ if image is not None
+ )
prompt = kwargs.get("prompt")
prompt_value = _single_outer_value(prompt, "prompt", adapter.workflow)
prompt_embeds = kwargs["prompt_embeds"]
@@ -420,6 +431,11 @@ def infer_h3_workflow(adapter: Any, **kwargs: Any) -> List[Any]:
},
"layout": layout,
"geometry": geometry,
+ **(
+ {}
+ if condition_image_slots is None
+ else {"condition_image_slots": condition_image_slots}
+ ),
},
)
return [sample]
@@ -732,23 +748,63 @@ def _validate_public_no_cfg_inputs(workflow: str, values: Mapping[str, Any], bou
)
-def _validate_fl2va_condition_images(values: Mapping[str, Any], boundary: str) -> Sequence[Any]:
+def _validate_fl2va_condition_images(
+ values: Mapping[str, Any],
+ boundary: str,
+) -> tuple[Any | None, Any | None]:
+ direct_first = values.get("image")
+ direct_last = values.get("last_image")
outer_images = values.get("images")
if outer_images is None:
outer_images = values.get("condition_images")
+ if (direct_first is not None or direct_last is not None) and outer_images is not None:
+ raise ValueError(
+ "MiniMax H3 workflow='fl2va' cannot combine direct image/last_image "
+ "arguments with grouped images"
+ )
+ if direct_first is not None or direct_last is not None:
+ return direct_first, direct_last
+
images = _single_outer_value(outer_images, "images", "fl2va")
if not isinstance(images, (list, tuple)) or not 1 <= len(images) <= 2:
raise ValueError(
f"MiniMax H3 workflow='fl2va' public {boundary} field='images' expected "
f"one or two ordered images, received {images!r}"
)
- return images
+ outer_slots = values.get("image_slots")
+ if outer_slots is None:
+ slots = ("first_frame", "last_frame")[: len(images)]
+ else:
+ slots = _single_outer_value(outer_slots, "image_slots", "fl2va")
+ if not isinstance(slots, (list, tuple)) or len(slots) != len(images):
+ raise ValueError(
+ f"MiniMax H3 workflow='fl2va' public {boundary} field='image_slots' "
+ f"expected {len(images)} slot(s), received {slots!r}"
+ )
+ slots = tuple(slots)
+ if len(set(slots)) != len(slots) or any(
+ slot not in ("first_frame", "last_frame") for slot in slots
+ ):
+ raise ValueError(
+ f"MiniMax H3 workflow='fl2va' public {boundary} field='image_slots' "
+ f"expected unique first_frame/last_frame values, received {slots!r}"
+ )
+ bound = dict(zip(slots, images))
+ return bound.get("first_frame"), bound.get("last_frame")
def _validate_workflow_media_inputs(
workflow: str, values: Mapping[str, Any], boundary: str
) -> None:
- media_fields = ("images", "condition_images", "videos", "audios", "references")
+ media_fields = (
+ "images",
+ "condition_images",
+ "image",
+ "last_image",
+ "videos",
+ "audios",
+ "references",
+ )
present = {
field for field in media_fields if field in values and _media_value_present(values[field])
}
@@ -757,7 +813,14 @@ def _validate_workflow_media_inputs(
f"MiniMax H3 workflow='t2va' {boundary} rejects media fields={tuple(sorted(present))}"
)
if workflow == "ref2va":
- generic = present & {"images", "condition_images", "videos", "audios"}
+ generic = present & {
+ "images",
+ "condition_images",
+ "image",
+ "last_image",
+ "videos",
+ "audios",
+ }
if generic:
raise ValueError(
f"MiniMax H3 workflow='ref2va' {boundary} rejects generic media "
diff --git a/src/flow_factory/models/pipeline_contracts.py b/src/flow_factory/models/pipeline_contracts.py
index 3903811d4..f39bfe6b1 100644
--- a/src/flow_factory/models/pipeline_contracts.py
+++ b/src/flow_factory/models/pipeline_contracts.py
@@ -110,6 +110,8 @@ def video_output_contract(
negative_prompt: NegativePromptPolicy,
input_image_min_count: Optional[int] = None,
input_image_max_count: Optional[int] = None,
+ input_image_slots: Tuple[str, ...] = (),
+ required_input_image_slots: Tuple[str, ...] = (),
output_fps: RateRequirement = RateRequirement.OPTIONAL,
geometry_source: GeometrySource = GeometrySource.OUTPUT_MEDIA,
batch_capability: BatchCapability = BatchCapability.UNIFORM,
@@ -120,6 +122,8 @@ def video_output_contract(
negative_prompt: Whether negative prompts are unsupported, optional, or required.
input_image_min_count: Minimum condition-image count, or ``None`` for no image input.
input_image_max_count: Maximum condition-image count when an image rule is present.
+ input_image_slots: Semantic image argument slots in positional fallback order.
+ required_input_image_slots: Slots that every valid request must fill.
output_fps: Whether target video frame rate metadata is required.
geometry_source: Boundary that determines output geometry.
batch_capability: Whether the adapter accepts uniform batches or one sample only.
@@ -134,6 +138,8 @@ def video_output_contract(
format=IMAGE_FORMAT,
min_count=input_image_min_count,
max_count=input_image_max_count,
+ slots=input_image_slots,
+ required_slots=required_input_image_slots,
),
)
video_format = MediaFormat(
@@ -164,6 +170,9 @@ def audio_video_output_contract(
input_rules: Tuple[InputMediaRule, ...] = (),
input_binding: InputMediaBinding = InputMediaBinding.GROUPED_BY_TYPE,
input_order: InputMediaOrder = InputMediaOrder.INSENSITIVE,
+ min_input_media_count: Optional[int] = None,
+ max_input_media_count: Optional[int] = None,
+ required_any_input_types: Tuple[MediaType, ...] = (),
output_fps: RateRequirement = RateRequirement.REQUIRED,
output_sample_rate: RateRequirement = RateRequirement.REQUIRED,
geometry_source: GeometrySource = GeometrySource.OUTPUT_MEDIA,
@@ -173,15 +182,17 @@ def audio_video_output_contract(
Supplying explicit input rules keeps this constructor neutral to how a model
binds conditions: prompt-only, grouped image conditions, and globally ordered
- heterogeneous references all use the same output contract. Cross-type total
- cardinality constraints remain adapter-owned because ``InputMediaSpec``
- represents per-type bounds only.
+ heterogeneous references all use the same output contract. Aggregate count
+ and required-any-type constraints cover cross-modality request invariants.
Args:
negative_prompt: Whether negative prompts are unsupported, optional, or required.
input_rules: Canonically ordered per-type input-media rules.
input_binding: How input media are projected into model-facing arguments.
input_order: Which input-media ordering carries semantic meaning.
+ min_input_media_count: Optional minimum count across all input modalities.
+ max_input_media_count: Optional maximum count across all input modalities.
+ required_any_input_types: Media types of which at least one must be present.
output_fps: Whether target-video frame-rate metadata is accepted or required.
output_sample_rate: Whether target-audio sample-rate metadata is accepted or required.
geometry_source: Boundary that determines aligned output geometry.
@@ -205,6 +216,9 @@ def audio_video_output_contract(
rules=input_rules,
binding=input_binding,
order=input_order,
+ min_total_count=min_input_media_count,
+ max_total_count=max_input_media_count,
+ required_any_types=required_any_input_types,
),
negative_prompt=negative_prompt,
output_media=OutputMediaSequence(items=(video_format, audio_format)),
diff --git a/src/flow_factory/models/wan/_conditioning.py b/src/flow_factory/models/wan/_conditioning.py
new file mode 100644
index 000000000..71556ab1d
--- /dev/null
+++ b/src/flow_factory/models/wan/_conditioning.py
@@ -0,0 +1,527 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Wan image-to-video condition realization shared by offline and rollout paths."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, ClassVar, Optional, Tuple
+
+import numpy as np
+import torch
+from PIL import Image
+
+from ...utils.image import is_image, is_image_batch, is_multi_image_batch, standardize_image_batch
+from ..condition_state import PreparedConditionState
+from ..configured_image_output import retrieve_vae_latents
+from ._output import configured_wan_video_output_geometry, normalize_wan_video_latents
+
+WanConditionImageRows = Tuple[Tuple[Image.Image, ...], ...]
+
+
+@dataclass(frozen=True, slots=True)
+class WanI2VConditionTensors:
+ """Hold the transformer condition channels and optional expanded-time mask."""
+
+ condition: torch.Tensor
+ first_frame_mask: Optional[torch.Tensor]
+
+
+def normalize_wan_i2v_image_rows(
+ images: Any,
+ *,
+ expected_batch_size: Optional[int] = None,
+) -> WanConditionImageRows:
+ """Normalize first/optional-last images while preserving within-sample order."""
+ if expected_batch_size is not None and (
+ type(expected_batch_size) is not int or expected_batch_size <= 0
+ ):
+ raise ValueError(
+ "Wan I2V expected_batch_size must be a positive integer or None, "
+ f"received {expected_batch_size!r}"
+ )
+
+ rows: list[tuple[Image.Image, ...]]
+ if is_multi_image_batch(images):
+ rows = []
+ if isinstance(images, (torch.Tensor, np.ndarray)):
+ iterable = list(images)
+ else:
+ iterable = images
+ for row_index, row in enumerate(iterable):
+ if isinstance(row, list) and not row:
+ raise ValueError(f"Wan I2V sample {row_index} requires at least one input image")
+ standardized = standardize_image_batch(row, output_type="pil")
+ rows.append(tuple(standardized))
+ elif is_image(images):
+ rows = [(standardize_image_batch(images, output_type="pil")[0],)]
+ elif is_image_batch(images):
+ standardized = tuple(standardize_image_batch(images, output_type="pil"))
+ if expected_batch_size == 1 and len(standardized) <= 2:
+ rows = [standardized]
+ else:
+ rows = [(image,) for image in standardized]
+ else:
+ raise TypeError(
+ "Wan I2V images must be one image, an image batch, or a multi-image batch, "
+ f"received {type(images).__name__}"
+ )
+
+ if expected_batch_size is not None and len(rows) != expected_batch_size:
+ raise ValueError(
+ "Wan I2V condition-image batch size mismatch: "
+ f"expected {expected_batch_size}, received {len(rows)}"
+ )
+ for row_index, row in enumerate(rows):
+ if len(row) not in (1, 2):
+ raise ValueError(
+ "Wan I2V requires one first-frame image and at most one optional "
+ f"last-frame image per sample, received {len(row)} at sample {row_index}"
+ )
+ return tuple(rows)
+
+
+def append_wan_i2v_last_images(
+ rows: WanConditionImageRows,
+ last_images: Any,
+) -> WanConditionImageRows:
+ """Bind a legacy separate last-image argument to normalized first-frame rows."""
+ if last_images is None:
+ return rows
+ if any(len(row) != 1 for row in rows):
+ raise ValueError(
+ "Wan I2V last_image cannot be combined with image rows that already contain "
+ "an optional last frame"
+ )
+ normalized_last = normalize_wan_i2v_image_rows(
+ last_images,
+ expected_batch_size=len(rows),
+ )
+ if any(len(row) != 1 for row in normalized_last):
+ raise ValueError("Wan I2V last_image must provide exactly one image per sample")
+ return tuple((row[0], last_row[0]) for row, last_row in zip(rows, normalized_last))
+
+
+def preprocess_wan_i2v_image_rows(
+ adapter: Any,
+ rows: WanConditionImageRows,
+ *,
+ height: int,
+ width: int,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """Apply the official Wan video processor to ordered first/last images."""
+ first_images = [row[0] for row in rows]
+ first = adapter.pipeline.video_processor.preprocess(
+ first_images,
+ height=height,
+ width=width,
+ )
+ if not isinstance(first, torch.Tensor):
+ raise TypeError(
+ "Wan video_processor.preprocess must return torch.Tensor, "
+ f"received {type(first).__name__}"
+ )
+ first = first.to(device=adapter.device, dtype=torch.float32)
+
+ has_last = [len(row) == 2 for row in rows]
+ if any(has_last) and not all(has_last):
+ raise ValueError("Wan I2V batches cannot mix first-only and first/last conditions")
+ if not any(has_last):
+ return first, None
+
+ last = adapter.pipeline.video_processor.preprocess(
+ [row[1] for row in rows],
+ height=height,
+ width=width,
+ )
+ if not isinstance(last, torch.Tensor):
+ raise TypeError(
+ "Wan video_processor.preprocess must return torch.Tensor, "
+ f"received {type(last).__name__}"
+ )
+ return first, last.to(device=adapter.device, dtype=torch.float32)
+
+
+def restore_wan_i2v_condition_pixels(
+ condition_images: Any,
+ *,
+ batch_size: int,
+ height: int,
+ width: int,
+ device: torch.device,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """Restore Arrow-cached per-sample first/last pixel tensors for VAE encoding."""
+ rows: list[torch.Tensor]
+ if isinstance(condition_images, torch.Tensor):
+ if condition_images.ndim == 5:
+ if condition_images.shape[0] != batch_size:
+ raise ValueError(
+ "Wan cached condition_images batch mismatch: "
+ f"expected {batch_size}, received {condition_images.shape[0]}"
+ )
+ rows = list(condition_images.unbind(0))
+ elif condition_images.ndim == 4 and batch_size == 1:
+ rows = [condition_images]
+ elif condition_images.ndim == 4 and condition_images.shape[0] == batch_size:
+ rows = [value.unsqueeze(0) for value in condition_images.unbind(0)]
+ else:
+ raise ValueError(
+ "Wan cached condition_images tensor must be BNCHW or one sample's NCHW, "
+ f"received shape {tuple(condition_images.shape)}"
+ )
+ elif isinstance(condition_images, Sequence) and not isinstance(condition_images, (str, bytes)):
+ if len(condition_images) != batch_size:
+ raise ValueError(
+ "Wan cached condition_images sequence batch mismatch: "
+ f"expected {batch_size}, received {len(condition_images)}"
+ )
+ rows = []
+ for index, value in enumerate(condition_images):
+ if not isinstance(value, torch.Tensor):
+ raise TypeError(
+ "Wan cached condition_images entries must be tensors, "
+ f"received {type(value).__name__} at sample {index}"
+ )
+ rows.append(value.unsqueeze(0) if value.ndim == 3 else value)
+ else:
+ raise TypeError(
+ "Wan cached condition_images must be a tensor or per-sample tensor sequence, "
+ f"received {type(condition_images).__name__}"
+ )
+
+ expected_tail = (3, height, width)
+ image_counts = []
+ for index, row in enumerate(rows):
+ if row.ndim != 4 or row.shape[0] not in (1, 2) or tuple(row.shape[1:]) != expected_tail:
+ raise ValueError(
+ "Wan cached condition_images must contain one first and optional last "
+ f"pixel tensor shaped {expected_tail}; received {tuple(row.shape)} "
+ f"at sample {index}"
+ )
+ image_counts.append(row.shape[0])
+ if any(count != image_counts[0] for count in image_counts):
+ raise ValueError("Wan I2V batches cannot mix first-only and first/last conditions")
+
+ first = torch.stack([row[0] for row in rows], dim=0).to(
+ device=device,
+ dtype=torch.float32,
+ )
+ if image_counts[0] == 1:
+ return first, None
+ last = torch.stack([row[1] for row in rows], dim=0).to(
+ device=device,
+ dtype=torch.float32,
+ )
+ return first, last
+
+
+def prepare_wan_i2v_condition_tensors(
+ adapter: Any,
+ image: torch.Tensor,
+ *,
+ height: int,
+ width: int,
+ num_frames: int,
+ dtype: torch.dtype,
+ device: torch.device,
+ last_image: Optional[torch.Tensor] = None,
+) -> WanI2VConditionTensors:
+ """Encode ordered input frames with posterior mode and build Wan condition channels."""
+ if not isinstance(dtype, torch.dtype) or not dtype.is_floating_point:
+ raise TypeError(f"Wan I2V condition dtype must be floating, received {dtype!r}")
+ if not isinstance(image, torch.Tensor) or image.ndim != 4:
+ raise ValueError(
+ "Wan I2V first-frame pixels must be rank-4 BCHW, "
+ f"received {type(image).__name__} with shape {getattr(image, 'shape', None)}"
+ )
+ batch_size = image.shape[0]
+ expected_pixels = (batch_size, 3, height, width)
+ if tuple(image.shape) != expected_pixels:
+ raise ValueError(
+ f"Wan I2V first-frame pixels must have shape {expected_pixels}, "
+ f"received {tuple(image.shape)}"
+ )
+ if last_image is not None:
+ if adapter.pipeline.config.expand_timesteps:
+ raise ValueError(
+ "Wan I2V expand_timesteps does not support an optional last-frame image; "
+ "Diffusers would otherwise ignore it"
+ )
+ if num_frames < 2:
+ raise ValueError(
+ "Wan I2V optional last-frame conditioning requires num_frames >= 2, "
+ f"received {num_frames}"
+ )
+ if not isinstance(last_image, torch.Tensor) or tuple(last_image.shape) != expected_pixels:
+ raise ValueError(
+ f"Wan I2V last-frame pixels must have shape {expected_pixels}, "
+ f"received {type(last_image).__name__} with shape "
+ f"{getattr(last_image, 'shape', None)}"
+ )
+
+ temporal_scale = adapter.pipeline.vae_scale_factor_temporal
+ spatial_scale = adapter.pipeline.vae_scale_factor_spatial
+ if (num_frames - 1) % temporal_scale:
+ raise ValueError(
+ "Wan I2V num_frames must satisfy "
+ f"(num_frames - 1) % {temporal_scale} == 0, received {num_frames}"
+ )
+ num_latent_frames = (num_frames - 1) // temporal_scale + 1
+ latent_height = height // spatial_scale
+ latent_width = width // spatial_scale
+
+ image = image.unsqueeze(2)
+ if adapter.pipeline.config.expand_timesteps:
+ video_condition = image
+ elif last_image is None:
+ video_condition = torch.cat(
+ [
+ image,
+ image.new_zeros(batch_size, image.shape[1], num_frames - 1, height, width),
+ ],
+ dim=2,
+ )
+ else:
+ video_condition = torch.cat(
+ [
+ image,
+ image.new_zeros(batch_size, image.shape[1], num_frames - 2, height, width),
+ last_image.unsqueeze(2),
+ ],
+ dim=2,
+ )
+
+ vae = adapter.vae
+ vae_dtype = getattr(vae, "dtype", None)
+ if not isinstance(vae_dtype, torch.dtype) or not vae_dtype.is_floating_point:
+ raise TypeError(
+ "Wan I2V condition encoder expected VAE to expose a floating dtype, "
+ f"received {vae_dtype!r}"
+ )
+ encoded = vae.encode(video_condition.to(device=device, dtype=vae_dtype))
+ latent_condition = retrieve_vae_latents(
+ encoded,
+ sample_mode="argmax",
+ source="Wan input video condition",
+ ).to(device=device, dtype=dtype)
+ latent_condition = normalize_wan_video_latents(adapter, latent_condition)
+
+ expected_latents = (
+ batch_size,
+ vae.config.z_dim,
+ num_latent_frames,
+ latent_height,
+ latent_width,
+ )
+ if tuple(latent_condition.shape) != expected_latents:
+ raise ValueError(
+ "Wan I2V condition latent geometry mismatch: "
+ f"expected {expected_latents}, received {tuple(latent_condition.shape)}"
+ )
+
+ if adapter.pipeline.config.expand_timesteps:
+ first_frame_mask = torch.ones(
+ batch_size,
+ 1,
+ num_latent_frames,
+ latent_height,
+ latent_width,
+ dtype=dtype,
+ device=device,
+ )
+ first_frame_mask[:, :, 0] = 0
+ return WanI2VConditionTensors(latent_condition, first_frame_mask)
+
+ mask = torch.ones(
+ batch_size,
+ 1,
+ num_frames,
+ latent_height,
+ latent_width,
+ dtype=dtype,
+ device=device,
+ )
+ if last_image is None:
+ mask[:, :, 1:] = 0
+ else:
+ mask[:, :, 1:-1] = 0
+ first = torch.repeat_interleave(mask[:, :, :1], dim=2, repeats=temporal_scale)
+ mask = torch.cat([first, mask[:, :, 1:]], dim=2)
+ mask = mask.view(
+ batch_size,
+ -1,
+ temporal_scale,
+ latent_height,
+ latent_width,
+ ).transpose(1, 2)
+ condition = torch.cat([mask, latent_condition], dim=1)
+ return WanI2VConditionTensors(condition, None)
+
+
+def normalize_wan_image_embeds(image_embeds: Any, *, batch_size: int) -> torch.Tensor:
+ """Pack cached per-sample CLIP embeddings into Diffusers' forward layout."""
+ if isinstance(image_embeds, torch.Tensor):
+ if image_embeds.ndim == 2:
+ if batch_size != 1:
+ raise ValueError(
+ "unbatched Wan image_embeds require batch_size=1, " f"received {batch_size}"
+ )
+ return image_embeds.unsqueeze(0)
+ if image_embeds.ndim == 3:
+ return image_embeds
+ if image_embeds.ndim == 4:
+ if image_embeds.shape[0] != batch_size:
+ raise ValueError(
+ "Wan image_embeds batch mismatch: "
+ f"expected {batch_size}, received {image_embeds.shape[0]}"
+ )
+ return image_embeds.flatten(0, 1)
+ raise ValueError(
+ "Wan image_embeds must be rank 2, 3, or 4, "
+ f"received shape {tuple(image_embeds.shape)}"
+ )
+ if isinstance(image_embeds, Sequence) and not isinstance(image_embeds, (str, bytes)):
+ if len(image_embeds) != batch_size:
+ raise ValueError(
+ "Wan image_embeds sequence batch mismatch: "
+ f"expected {batch_size}, received {len(image_embeds)}"
+ )
+ packed = []
+ for index, value in enumerate(image_embeds):
+ if not isinstance(value, torch.Tensor) or value.ndim not in (2, 3):
+ raise ValueError(
+ "Wan image_embeds sequence entries must be rank-2 or rank-3 tensors, "
+ f"received {type(value).__name__} with shape "
+ f"{getattr(value, 'shape', None)} at sample {index}"
+ )
+ packed.append(value.unsqueeze(0) if value.ndim == 2 else value)
+ return torch.cat(packed, dim=0)
+ raise TypeError(
+ "Wan image_embeds must be a tensor or per-sample tensor sequence, "
+ f"received {type(image_embeds).__name__}"
+ )
+
+
+def split_wan_image_embeds(
+ image_embeds: torch.Tensor,
+ image_counts: Sequence[int],
+) -> Tuple[torch.Tensor, ...]:
+ """Split Diffusers' packed CLIP layout back into per-sample replay tensors."""
+ if not isinstance(image_embeds, torch.Tensor) or image_embeds.ndim != 3:
+ raise ValueError(
+ "packed Wan image_embeds must be rank-3, "
+ f"received {type(image_embeds).__name__} with shape "
+ f"{getattr(image_embeds, 'shape', None)}"
+ )
+ counts = tuple(image_counts)
+ if not counts or any(type(count) is not int or count not in (1, 2) for count in counts):
+ raise ValueError(
+ "Wan image_counts must contain one or two images per sample, " f"received {counts!r}"
+ )
+ if sum(counts) != image_embeds.shape[0]:
+ raise ValueError(
+ "packed Wan image_embeds count mismatch: "
+ f"expected {sum(counts)}, received {image_embeds.shape[0]}"
+ )
+ per_sample = []
+ offset = 0
+ for count in counts:
+ value = image_embeds[offset : offset + count]
+ per_sample.append(value[0] if count == 1 else value)
+ offset += count
+ return tuple(per_sample)
+
+
+@dataclass(frozen=True, slots=True)
+class WanI2VConditionStatePreparer:
+ """Realize input-owned Wan VAE conditions once per offline batch."""
+
+ adapter: Any
+ required_components: ClassVar[Tuple[str, ...]] = ("vae",)
+
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ """Bind cached images to configured geometry with deterministic VAE mode."""
+ del generator
+ if "condition_images" not in condition:
+ raise ValueError(
+ "Wan I2V cached condition is missing preprocessed 'condition_images'; "
+ "rebuild the input-condition cache with this adapter version"
+ )
+ height, width, num_frames, _ = configured_wan_video_output_geometry(self.adapter)
+ first, last = restore_wan_i2v_condition_pixels(
+ condition["condition_images"],
+ batch_size=1,
+ height=height,
+ width=width,
+ device=self.adapter.device,
+ )
+ realized = prepare_wan_i2v_condition_tensors(
+ self.adapter,
+ first,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ dtype=torch.float32,
+ device=self.adapter.device,
+ last_image=last,
+ )
+
+ cached_condition = dict(condition)
+ cached_condition.pop("condition_images")
+ cached_condition.pop("images", None)
+ forward_context: dict[str, Any] = {"latent_condition": realized.condition}
+ image_embeds = cached_condition.pop("image_embeds", None)
+ if image_embeds is not None:
+ packed_image_embeds = normalize_wan_image_embeds(
+ image_embeds,
+ batch_size=first.shape[0],
+ )
+ expected_image_embeds = first.shape[0] * (2 if last is not None else 1)
+ if packed_image_embeds.shape[0] != expected_image_embeds:
+ raise ValueError(
+ "Wan cached image_embeds count disagrees with condition_images: "
+ f"expected {expected_image_embeds}, received "
+ f"{packed_image_embeds.shape[0]}"
+ )
+ forward_context["image_embeds"] = packed_image_embeds
+
+ output_context: dict[str, Any] = {}
+ if realized.first_frame_mask is not None:
+ forward_context["first_frame_mask"] = realized.first_frame_mask
+ output_context["first_frame_mask"] = realized.first_frame_mask
+ return PreparedConditionState(
+ condition=cached_condition,
+ forward_context=forward_context,
+ output_context=output_context,
+ )
+
+
+__all__ = [
+ "WanConditionImageRows",
+ "WanI2VConditionStatePreparer",
+ "WanI2VConditionTensors",
+ "append_wan_i2v_last_images",
+ "normalize_wan_i2v_image_rows",
+ "normalize_wan_image_embeds",
+ "prepare_wan_i2v_condition_tensors",
+ "preprocess_wan_i2v_image_rows",
+ "restore_wan_i2v_condition_pixels",
+ "split_wan_image_embeds",
+]
diff --git a/src/flow_factory/models/wan/_output.py b/src/flow_factory/models/wan/_output.py
index 887e61232..430c0de78 100644
--- a/src/flow_factory/models/wan/_output.py
+++ b/src/flow_factory/models/wan/_output.py
@@ -12,12 +12,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""On-the-fly clean-video encoding for Wan text-to-video adapters."""
+"""Shared on-the-fly clean-video encoding for Wan video adapters."""
from __future__ import annotations
+import math
from collections.abc import Mapping
from dataclasses import dataclass
+from numbers import Real
from typing import Any, ClassVar, Optional, Tuple
import numpy as np
@@ -34,11 +36,172 @@
)
+def configured_wan_video_output_geometry(adapter: Any) -> Tuple[int, int, int, float]:
+ """Return configured Wan geometry after exact latent-grid validation."""
+ geometry = []
+ for name in ("height", "width", "num_frames"):
+ value = getattr(adapter.training_args, name, None)
+ if type(value) is not int or value <= 0:
+ raise ValueError(
+ f"Wan output geometry requires positive integer train.{name}, "
+ f"received {value!r}"
+ )
+ geometry.append(value)
+ frame_rate = getattr(adapter.training_args, "frame_rate", None)
+ if isinstance(frame_rate, bool) or not isinstance(frame_rate, Real):
+ raise TypeError(
+ "Wan output geometry requires finite positive train.frame_rate, "
+ f"received {type(frame_rate).__name__}: {frame_rate!r}"
+ )
+ frame_rate = float(frame_rate)
+ if not math.isfinite(frame_rate) or frame_rate <= 0:
+ raise ValueError(
+ "Wan output geometry requires finite positive train.frame_rate, "
+ f"received {frame_rate!r}"
+ )
+
+ height, width, num_frames = geometry
+ temporal_scale = adapter.pipeline.vae_scale_factor_temporal
+ spatial_scale = adapter.pipeline.vae_scale_factor_spatial
+ if (num_frames - 1) % temporal_scale:
+ raise ValueError(
+ "Wan output num_frames must satisfy "
+ f"(num_frames - 1) % {temporal_scale} == 0, received {num_frames}"
+ )
+ transformer = (
+ adapter.pipeline.transformer
+ if adapter.pipeline.transformer is not None
+ else adapter.pipeline.transformer_2
+ )
+ if transformer is None:
+ raise RuntimeError("Wan output geometry requires one materialized transformer")
+ patch_size = transformer.config.patch_size
+ height_multiple = spatial_scale * patch_size[1]
+ width_multiple = spatial_scale * patch_size[2]
+ if height % height_multiple or width % width_multiple:
+ raise ValueError(
+ "Wan output height/width must be divisible by transformer latent-grid "
+ f"multiples {(height_multiple, width_multiple)}, received {(height, width)}"
+ )
+ return height, width, num_frames, frame_rate
+
+
+def resample_wan_output_video(
+ video: np.ndarray,
+ *,
+ source_fps: Optional[float],
+ target_frames: int,
+ target_fps: float,
+) -> np.ndarray:
+ """Select deterministic nearest-time frames for configured target cadence."""
+ if video.dtype != np.uint8 or video.ndim != 4 or video.shape[-1] != 3:
+ raise ValueError(
+ "Wan decoded target video must be uint8 RGB shaped (F,H,W,3), "
+ f"received dtype={video.dtype}, shape={tuple(video.shape)}"
+ )
+ if video.shape[0] < 1:
+ raise ValueError("Wan decoded target video must contain at least one frame")
+ if isinstance(source_fps, bool) or not isinstance(source_fps, Real):
+ raise TypeError(
+ "Wan target video requires source fps metadata, "
+ f"received {type(source_fps).__name__}: {source_fps!r}"
+ )
+ source_fps = float(source_fps)
+ if not math.isfinite(source_fps) or source_fps <= 0:
+ raise ValueError(f"Wan target video requires positive finite fps, got {source_fps!r}")
+ indices = np.rint(np.arange(target_frames, dtype=np.float64) * source_fps / target_fps).astype(
+ np.int64
+ )
+ if indices[-1] >= video.shape[0]:
+ required_duration = (target_frames - 1) / target_fps
+ available_duration = (video.shape[0] - 1) / source_fps
+ raise ValueError(
+ "Wan target video is too short for configured temporal geometry: "
+ f"requires {required_duration:.6f}s, has {available_duration:.6f}s"
+ )
+ return np.ascontiguousarray(video[indices])
+
+
+def normalize_wan_video_latents(adapter: Any, latents: torch.Tensor) -> torch.Tensor:
+ """Apply the exact inverse of Wan's existing decode normalization."""
+ if not isinstance(latents, torch.Tensor) or latents.ndim != 5:
+ raise ValueError(
+ "Wan VAE video latents must be rank-5 BCFHW, "
+ f"received {type(latents).__name__} with shape "
+ f"{getattr(latents, 'shape', None)}"
+ )
+ config = adapter.vae.config
+ z_dim = config.z_dim
+ if latents.shape[1] != z_dim:
+ raise ValueError(
+ f"Wan VAE video latent channels must equal z_dim={z_dim}, "
+ f"received {latents.shape[1]}"
+ )
+ latents_mean = torch.as_tensor(
+ config.latents_mean,
+ device=latents.device,
+ dtype=latents.dtype,
+ ).view(1, z_dim, 1, 1, 1)
+ inverse_std = (
+ torch.as_tensor(
+ config.latents_std,
+ device=latents.device,
+ dtype=latents.dtype,
+ )
+ .reciprocal()
+ .view(1, z_dim, 1, 1, 1)
+ )
+ return (latents - latents_mean) * inverse_std
+
+
+def validate_wan_encoded_output_geometry(
+ adapter: Any,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+) -> None:
+ """Require encoded signatures and decode metadata to match train geometry."""
+ del condition
+ height, width, num_frames, frame_rate = configured_wan_video_output_geometry(adapter)
+ if len(encoded.geometry_signatures) != len(media_batch):
+ raise ValueError(
+ "Wan output codec must return one geometry signature per sample, "
+ f"received {len(encoded.geometry_signatures)} for {len(media_batch)}"
+ )
+ for sample_index, signature in enumerate(encoded.geometry_signatures):
+ geometry = signature.media[0]
+ received = (
+ geometry.height,
+ geometry.width,
+ geometry.frames,
+ geometry.fps,
+ )
+ expected = (height, width, num_frames, frame_rate)
+ if received != expected:
+ raise ValueError(
+ "Wan encoded output geometry disagrees with configured geometry for "
+ f"sample {sample_index}: expected {expected}, received {received}"
+ )
+ expected_context = {
+ "height": height,
+ "width": width,
+ "num_frames": num_frames,
+ "frame_rate": frame_rate,
+ }
+ for name, expected in expected_context.items():
+ if encoded.decode_context.get(name) != expected:
+ raise ValueError(
+ f"Wan decode_context {name!r} must equal {expected!r}, "
+ f"received {encoded.decode_context.get(name)!r}"
+ )
+
+
@dataclass(frozen=True, slots=True)
class WanVideoOutputCodec:
"""Encode configured Wan target videos without retaining pixels or latents."""
adapter: Any
+ bind_condition_active_mask: bool = False
required_components: ClassVar[Tuple[str, ...]] = ("vae",)
def encode_output_state(
@@ -48,8 +211,28 @@ def encode_output_state(
generator: Optional[torch.Generator] = None,
) -> EncodedOutputState:
"""Preprocess, VAE-sample, and normalize one video per sample."""
- del condition
- height, width, num_frames, frame_rate = self.adapter._configured_video_output_geometry()
+ height, width, num_frames, frame_rate = configured_wan_video_output_geometry(self.adapter)
+ first_frame_mask = None
+ if self.bind_condition_active_mask:
+ first_frame_mask = condition.get("first_frame_mask")
+ expand_timesteps = bool(self.adapter.pipeline.config.expand_timesteps)
+ if expand_timesteps and first_frame_mask is None:
+ raise ValueError(
+ "Wan I2V expand_timesteps target encoding requires first_frame_mask "
+ "from the prepared input condition"
+ )
+ if not expand_timesteps and first_frame_mask is not None:
+ raise ValueError(
+ "Wan I2V non-expanded target encoding must not receive first_frame_mask"
+ )
+ if first_frame_mask is not None:
+ if not isinstance(first_frame_mask, torch.Tensor):
+ raise TypeError(
+ "Wan I2V first_frame_mask must be torch.Tensor, "
+ f"received {type(first_frame_mask).__name__}"
+ )
+ if not torch.all((first_frame_mask == 0) | (first_frame_mask == 1)):
+ raise ValueError("Wan I2V first_frame_mask must contain only zero and one")
videos = []
for sample_index, candidate in enumerate(media_batch):
if len(candidate) != 1:
@@ -65,7 +248,7 @@ def encode_output_state(
f"received {type(payload).__name__} for sample {sample_index}"
)
videos.append(
- self.adapter._resample_output_video(
+ resample_wan_output_video(
payload,
source_fps=media.fps,
target_frames=num_frames,
@@ -105,7 +288,7 @@ def encode_output_state(
generator=generator,
source="Wan target video",
)
- latents = self.adapter._normalize_output_video_latents(latents)
+ latents = normalize_wan_video_latents(self.adapter, latents)
temporal_scale = self.adapter.pipeline.vae_scale_factor_temporal
spatial_scale = self.adapter.pipeline.vae_scale_factor_spatial
@@ -133,8 +316,13 @@ def encode_output_state(
),
)
)
+ active_masks = None
+ if first_frame_mask is not None:
+ active_mask = first_frame_mask.to(device=latents.device, dtype=torch.bool)
+ active_masks = {"latent": active_mask}
+
return EncodedOutputState(
- clean_state=LatentState({"latent": latents}),
+ clean_state=LatentState({"latent": latents}, active_masks=active_masks),
forward_context={},
decode_context={
"height": height,
@@ -146,4 +334,10 @@ def encode_output_state(
)
-__all__ = ["WanVideoOutputCodec"]
+__all__ = [
+ "WanVideoOutputCodec",
+ "configured_wan_video_output_geometry",
+ "normalize_wan_video_latents",
+ "resample_wan_output_video",
+ "validate_wan_encoded_output_geometry",
+]
diff --git a/src/flow_factory/models/wan/wan2_i2v.py b/src/flow_factory/models/wan/wan2_i2v.py
index b437e79b0..297d97c59 100644
--- a/src/flow_factory/models/wan/wan2_i2v.py
+++ b/src/flow_factory/models/wan/wan2_i2v.py
@@ -15,11 +15,9 @@
# src/flow_factory/models/wan/wan2_i2v.py
from __future__ import annotations
-import logging
-import os
-from collections import defaultdict
from dataclasses import dataclass
-from typing import Any, ClassVar, Dict, Iterable, List, Literal, Optional, Tuple, Union
+from types import MappingProxyType
+from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Tuple, Union
import numpy as np
import torch
@@ -27,30 +25,49 @@
from diffusers.pipelines.wan.pipeline_wan_i2v import WanImageToVideoPipeline, prompt_clean
from diffusers.utils.torch_utils import randn_tensor
from peft import PeftModel
-from PIL import Image
+from ...contracts import (
+ BatchCapability,
+ GeometrySource,
+ NegativePromptPolicy,
+ PipelineIOContract,
+ RateRequirement,
+)
from ...hparams import *
from ...samples import I2VSample
from ...scheduler import UniPCMultistepSDEScheduler, UniPCMultistepSDESchedulerOutput
-from ...utils.base import filter_kwargs
from ...utils.image import (
ImageBatch,
ImageSingle,
MultiImageBatch,
- is_image,
- is_image_batch,
- is_multi_image_batch,
- standardize_image_batch,
)
from ...utils.logger_utils import setup_logger
from ...utils.trajectory_collector import (
- CallbackCollector,
- TrajectoryCollector,
TrajectoryIndicesType,
create_callback_collector,
create_trajectory_collector,
)
from ..abc import BaseAdapter
+from ..condition_state import ConditionStatePreparer
+from ..output_state import DecodedMediaBatch, EncodedOutputState, OutputStateCodec
+from ..pipeline_contracts import video_output_contract
+from ._conditioning import (
+ WanI2VConditionStatePreparer,
+ append_wan_i2v_last_images,
+ normalize_wan_i2v_image_rows,
+ normalize_wan_image_embeds,
+ prepare_wan_i2v_condition_tensors,
+ preprocess_wan_i2v_image_rows,
+ restore_wan_i2v_condition_pixels,
+ split_wan_image_embeds,
+)
+from ._output import (
+ WanVideoOutputCodec,
+ configured_wan_video_output_geometry,
+ normalize_wan_video_latents,
+ resample_wan_output_video,
+ validate_wan_encoded_output_geometry,
+)
logger = setup_logger(__name__)
@@ -58,33 +75,18 @@
@dataclass
class WanI2VSample(I2VSample):
# Class var
- _shared_fields: ClassVar[frozenset[str]] = frozenset({"first_frame_mask"})
+ _shared_fields: ClassVar[frozenset[str]] = frozenset()
# Obj var
image_embeds: Optional[torch.FloatTensor] = None
- condition: Optional[torch.FloatTensor] = None
+ latent_condition: Optional[torch.FloatTensor] = None
first_frame_mask: Optional[torch.FloatTensor] = None
-def retrieve_latents(
- encoder_output: torch.Tensor,
- generator: Optional[torch.Generator] = None,
- sample_mode: str = "sample",
-):
- if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
- return encoder_output.latent_dist.sample(generator)
- elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
- return encoder_output.latent_dist.mode()
- elif hasattr(encoder_output, "latents"):
- return encoder_output.latents
- else:
- raise AttributeError("Could not access latents of provided encoder_output")
-
-
class Wan2_I2V_Adapter(BaseAdapter):
- output_state_codec_unavailable_reason = (
- "Wan I2V target encoding must bind an output-geometry-dependent first-frame VAE "
- "condition (and Wan 2.2 first-frame mask); the current condition cache retains only "
- "CLIP image embeddings, so this target/condition binder is not yet implemented"
+ preprocess_cache_fields = frozenset({"height", "width"})
+ preprocess_cache_version = "wan-i2v-condition-pixels-v1"
+ offline_training_forward_overrides = MappingProxyType(
+ {"guidance_scale": 1.0, "guidance_scale_2": 1.0}
)
# Wan2.2 trains both transformer and transformer_2 but uses only one per
# timestep (boundary_ratio), so under DDP the other's trainable params get no
@@ -97,12 +99,21 @@ class Wan2_I2V_Adapter(BaseAdapter):
"vae": torch.float32,
"image_encoder": torch.float32,
}
+ pipeline_io_contract = video_output_contract(
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ input_image_min_count=1,
+ input_image_max_count=2,
+ input_image_slots=("first_frame", "last_frame"),
+ required_input_image_slots=("first_frame",),
+ output_fps=RateRequirement.REQUIRED,
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
def __init__(self, config: Arguments, accelerator: Accelerator):
super().__init__(config, accelerator)
self.pipeline: WanImageToVideoPipeline
self.scheduler: UniPCMultistepSDEScheduler
- self._has_warned_multi_image = False
def load_pipeline(self) -> WanImageToVideoPipeline:
return self._load_diffusers_pipeline(
@@ -110,6 +121,21 @@ def load_pipeline(self) -> WanImageToVideoPipeline:
self.model_args.model_name_or_path,
)
+ def _resolve_pipeline_io_contract(self) -> PipelineIOContract:
+ """Narrow expand-timestep checkpoints to their first-frame-only semantics."""
+ if not self.pipeline.config.expand_timesteps:
+ return type(self).pipeline_io_contract
+ return video_output_contract(
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ input_image_min_count=1,
+ input_image_max_count=1,
+ input_image_slots=("first_frame",),
+ required_input_image_slots=("first_frame",),
+ output_fps=RateRequirement.REQUIRED,
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
+
@property
def default_target_modules(self) -> List[str]:
"""Default LoRA target modules for Wan transformer."""
@@ -143,7 +169,7 @@ def inference_modules(self) -> List[str]:
@property
def preprocessing_modules(self) -> List[str]:
"""Modules that are requires for preprocessing"""
- return ["text_encoders", "vae", "image_encoder"]
+ return ["text_encoders", "image_encoder"]
def apply_lora(
self,
@@ -298,59 +324,112 @@ def encode_image(
self,
images: Union[ImageSingle, ImageBatch, MultiImageBatch],
device: Optional[torch.device] = None,
- ) -> Union[None, Dict[str, torch.Tensor]]:
- images = self._standardize_image_input(
- images,
- output_type="pil",
- )
-
- if not is_image_batch(images):
+ height: Optional[int] = None,
+ width: Optional[int] = None,
+ ) -> Dict[str, Union[List[torch.Tensor], torch.Tensor]]:
+ rows = normalize_wan_i2v_image_rows(images)
+ height = height if height is not None else getattr(self.training_args, "height", None)
+ width = width if width is not None else getattr(self.training_args, "width", None)
+ if type(height) is not int or height <= 0 or type(width) is not int or width <= 0:
raise ValueError(
- f"Invalid image input type: {type(images)}. "
- f"Must be a PIL Image, numpy array, torch tensor, or a list of these types."
+ "Wan I2V image preprocessing requires positive integer train.height/width, "
+ f"received {(height, width)}"
)
+ condition_images = []
+ for row_index, row in enumerate(rows):
+ pixel_values = self.pipeline.video_processor.preprocess(
+ list(row),
+ height=height,
+ width=width,
+ )
+ expected_shape = (len(row), 3, height, width)
+ if (
+ not isinstance(pixel_values, torch.Tensor)
+ or tuple(pixel_values.shape) != expected_shape
+ ):
+ raise ValueError(
+ "Wan video_processor.preprocess changed condition-image geometry at "
+ f"sample {row_index}: expected {expected_shape}, received "
+ f"{getattr(pixel_values, 'shape', None)}"
+ )
+ condition_images.append(pixel_values.detach().to(device="cpu", dtype=torch.float32))
+ results: Dict[str, Union[List[torch.Tensor], torch.Tensor]] = {
+ "condition_images": condition_images
+ }
# only Wan 2.1 I2V transformer accepts image_embeds, else None directly
if (
self.pipeline.transformer is not None
and self.pipeline.transformer.config.image_dim is not None
):
- batch_size = len(images)
device = device or self.image_encoder.device
- images = self.pipeline.image_processor(images=images, return_tensors="pt").to(device)
- image_embeds = self.pipeline.image_encoder(**images, output_hidden_states=True)
- return {
- "image_embeds": image_embeds.hidden_states[-2],
- }
- else:
- return None
-
- def _standardize_image_input(
- self,
- images: Union[ImageSingle, ImageBatch, MultiImageBatch],
- output_type: Literal["pil", "pt", "np"] = "pil",
- ):
- """
- Standardize image input to desired output type.
- """
- if isinstance(images, Image.Image):
- images = [images]
- elif is_multi_image_batch(images):
- # A list of list of images
- if any(len(batch) > 1 for batch in images) and not self._has_warned_multi_image:
- self._has_warned_multi_image = True
- logger.warning(
- "Multiple condition images are not supported for Wan2_I2V. Only the first image of each batch will be used."
+ counts = [len(row) for row in rows]
+ flattened = [image for row in rows for image in row]
+ processor_output = self.pipeline.image_processor(
+ images=flattened,
+ return_tensors="pt",
+ ).to(device)
+ image_embeds = self.pipeline.image_encoder(
+ **processor_output,
+ output_hidden_states=True,
+ ).hidden_states[-2]
+ if not isinstance(image_embeds, torch.Tensor) or image_embeds.shape[0] != sum(counts):
+ raise ValueError(
+ "Wan image encoder must preserve the flattened condition-image count, "
+ f"expected {sum(counts)}, received "
+ f"{getattr(image_embeds, 'shape', None)}"
)
+ if all(count == 1 for count in counts):
+ results["image_embeds"] = image_embeds
+ return results
+ per_sample = []
+ offset = 0
+ for count in counts:
+ per_sample.append(image_embeds[offset : offset + count])
+ offset += count
+ results["image_embeds"] = per_sample
+ return results
- images = [batch[0] for batch in images]
+ def build_condition_state_preparer(self) -> ConditionStatePreparer:
+ """Declare on-the-fly input-frame VAE conditioning."""
+ return WanI2VConditionStatePreparer(self)
+
+ def build_output_state_codec(self) -> OutputStateCodec:
+ """Declare the shared Wan target-video codec with I2V active-mask binding."""
+ return WanVideoOutputCodec(self, bind_condition_active_mask=True)
+
+ def _configured_video_output_geometry(self) -> Tuple[int, int, int, float]:
+ """Return configured Wan geometry after exact latent-grid validation."""
+ return configured_wan_video_output_geometry(self)
+
+ @staticmethod
+ def _resample_output_video(
+ video: np.ndarray,
+ *,
+ source_fps: Optional[float],
+ target_frames: int,
+ target_fps: float,
+ ) -> np.ndarray:
+ """Select deterministic nearest-time frames for configured target cadence."""
+ return resample_wan_output_video(
+ video,
+ source_fps=source_fps,
+ target_frames=target_frames,
+ target_fps=target_fps,
+ )
- images = standardize_image_batch(images, output_type=output_type)
- return images
+ def _normalize_output_video_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ """Apply Wan's VAE latent normalization."""
+ return normalize_wan_video_latents(self, latents)
- # ------------------------ Video Encoding ------------------------
- def encode_video(self, videos: Union[np.ndarray, torch.Tensor, List[Image.Image]]):
- pass
+ def _validate_encoded_output_geometry(
+ self,
+ media_batch: DecodedMediaBatch,
+ condition: Mapping[str, Any],
+ encoded: EncodedOutputState,
+ ) -> None:
+ """Require encoded target metadata to match configured video geometry."""
+ validate_wan_encoded_output_geometry(self, media_batch, condition, encoded)
# ------------------------ Latent Decoding ------------------------
def decode_latents(
@@ -388,8 +467,10 @@ def prepare_latents(
last_image: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
- Modified from `diffusers: WanImageToVideoPipeline` with batch_size bug fixed
+ Prepare rollout noise and reuse the offline condition realization helper.
"""
+ dtype = dtype or torch.float32
+ device = device or self.device
num_latent_frames = (num_frames - 1) // self.pipeline.vae_scale_factor_temporal + 1
latent_height = height // self.pipeline.vae_scale_factor_spatial
latent_width = width // self.pipeline.vae_scale_factor_spatial
@@ -406,74 +487,19 @@ def prepare_latents(
else:
latents = latents.to(device=device, dtype=dtype)
- image = image.unsqueeze(2) # [batch_size, channels, 1, height, width]
-
- if self.pipeline.config.expand_timesteps:
- video_condition = image
-
- elif last_image is None:
- video_condition = torch.cat(
- [
- image,
- image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, height, width),
- ],
- dim=2,
- )
- else:
- last_image = last_image.unsqueeze(2)
- video_condition = torch.cat(
- [
- image,
- image.new_zeros(image.shape[0], image.shape[1], num_frames - 2, height, width),
- last_image,
- ],
- dim=2,
- )
- video_condition = video_condition.to(device=device, dtype=self.pipeline.vae.dtype)
-
- latents_mean = (
- torch.tensor(self.pipeline.vae.config.latents_mean)
- .view(1, self.pipeline.vae.config.z_dim, 1, 1, 1)
- .to(latents.device, latents.dtype)
- )
- latents_std = 1.0 / torch.tensor(self.pipeline.vae.config.latents_std).view(
- 1, self.pipeline.vae.config.z_dim, 1, 1, 1
- ).to(latents.device, latents.dtype)
-
- latent_condition = retrieve_latents(
- self.pipeline.vae.encode(video_condition), sample_mode="argmax"
- )
- if latent_condition.shape[0] == 1 and batch_size > 1:
- latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
-
- latent_condition = latent_condition.to(dtype)
- latent_condition = (latent_condition - latents_mean) * latents_std
-
- if self.pipeline.config.expand_timesteps:
- first_frame_mask = torch.ones(
- 1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device
- )
- first_frame_mask[:, :, 0] = 0
- return latents, latent_condition, first_frame_mask
-
- mask_lat_size = torch.ones(batch_size, 1, num_frames, latent_height, latent_width)
-
- if last_image is None:
- mask_lat_size[:, :, 1:] = 0
- else:
- mask_lat_size[:, :, 1:-1] = 0
- first_frame_mask = mask_lat_size[:, :, 0:1]
- first_frame_mask = torch.repeat_interleave(
- first_frame_mask, dim=2, repeats=self.pipeline.vae_scale_factor_temporal
- )
- mask_lat_size = torch.concat([first_frame_mask, mask_lat_size[:, :, 1:, :]], dim=2)
- mask_lat_size = mask_lat_size.view(
- batch_size, -1, self.pipeline.vae_scale_factor_temporal, latent_height, latent_width
+ realized = prepare_wan_i2v_condition_tensors(
+ self,
+ image,
+ height=height,
+ width=width,
+ num_frames=num_frames,
+ dtype=dtype,
+ device=device,
+ last_image=last_image,
)
- mask_lat_size = mask_lat_size.transpose(1, 2)
- mask_lat_size = mask_lat_size.to(latent_condition.device)
-
- return latents, torch.concat([mask_lat_size, latent_condition], dim=1)
+ if realized.first_frame_mask is not None:
+ return latents, realized.condition, realized.first_frame_mask
+ return latents, realized.condition
# ======================== Inference ========================
@torch.no_grad()
@@ -499,7 +525,7 @@ def inference(
# Encoded Image
image_embeds: Optional[torch.Tensor] = None,
condition_images: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
- last_image: Optional[torch.Tensor] = None, # Not supported yet
+ last_image: Optional[Union[ImageSingle, ImageBatch]] = None,
# Other args
compute_log_prob: bool = False,
attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -543,8 +569,6 @@ def inference(
)
height, width = calc_height, calc_width
- images = self._standardize_image_input(images, output_type="pil")
-
# 2. Encode prompt
if prompt_embeds is None:
encoded = self.encode_prompt(
@@ -573,6 +597,17 @@ def inference(
if negative_prompt_embeds is not None:
negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
+ condition_rows = normalize_wan_i2v_image_rows(
+ images,
+ expected_batch_size=batch_size,
+ )
+ condition_rows = append_wan_i2v_last_images(condition_rows, last_image)
+ if self.pipeline.config.expand_timesteps and any(len(row) == 2 for row in condition_rows):
+ raise ValueError(
+ "Wan I2V expand_timesteps does not support an optional last-frame image; "
+ "Diffusers would otherwise ignore it"
+ )
+
# 3. Set scheduler
self.scheduler.set_timesteps(num_inference_steps, device=device)
timesteps = self.scheduler.timesteps
@@ -584,28 +619,50 @@ def inference(
and self.pipeline.transformer.config.image_dim is not None
):
if image_embeds is None:
- image_to_encode = images if last_image is None else [images, last_image]
- image_encoded = self.encode_image(image_to_encode, device)
+ image_encoded = self.encode_image(
+ condition_rows,
+ device,
+ height=height,
+ width=width,
+ )
image_embeds = image_encoded["image_embeds"]
-
- image_embeds = (
- image_embeds.to(device=device, dtype=transformer_dtype)
- if image_embeds is not None
- else None
- )
+ if condition_images is None:
+ condition_images = image_encoded["condition_images"]
+ if image_embeds is not None:
+ image_embeds = normalize_wan_image_embeds(
+ image_embeds,
+ batch_size=batch_size,
+ ).to(device=device, dtype=transformer_dtype)
+ per_sample_image_embeds = split_wan_image_embeds(
+ image_embeds,
+ [len(row) for row in condition_rows],
+ )
+ else:
+ per_sample_image_embeds = (None,) * batch_size
# 5. Prepare latent variables
num_channels_latents = self.pipeline.vae.config.z_dim
- images = self.pipeline.video_processor.preprocess(images, height=height, width=width).to(
- device, dtype=torch.float32
- )
- if last_image is not None:
- last_image = self.pipeline.video_processor.preprocess(
- last_image, height=height, width=width
- ).to(device, dtype=torch.float32)
+ if condition_images is None:
+ images, last_image_pixels = preprocess_wan_i2v_image_rows(
+ self,
+ condition_rows,
+ height=height,
+ width=width,
+ )
+ else:
+ images, last_image_pixels = restore_wan_i2v_condition_pixels(
+ condition_images,
+ batch_size=batch_size,
+ height=height,
+ width=width,
+ device=device,
+ )
+ expected_last = any(len(row) == 2 for row in condition_rows)
+ if expected_last != (last_image_pixels is not None):
+ raise ValueError(
+ "Wan cached condition_images count disagrees with ordered raw input images"
+ )
- # Inside the following function, preparing `latents_condition` requires `latents_mean` and `latents_std`,
- # which depend on `latents` initialized at runtime. Therefore, this part is kept inside inference function and not moved to preprocess_func.
latents_outputs = self.prepare_latents(
image=images,
batch_size=batch_size,
@@ -617,10 +674,10 @@ def inference(
device=device,
generator=generator,
latents=None,
- last_image=last_image,
+ last_image=last_image_pixels,
)
if self.pipeline.config.expand_timesteps:
- # wan 2.2 5b i2v use firt_frame_mask to mask timesteps
+ # Wan 2.2 5B I2V uses first_frame_mask to expand timesteps.
latents, condition, first_frame_mask = latents_outputs
else:
latents, condition = latents_outputs
@@ -657,7 +714,7 @@ def inference(
guidance_scale=guidance_scale,
guidance_scale_2=guidance_scale_2,
image_embeds=image_embeds,
- condition=condition,
+ latent_condition=condition,
first_frame_mask=first_frame_mask,
attention_kwargs=attention_kwargs,
compute_log_prob=current_compute_log_prob,
@@ -712,10 +769,10 @@ def inference(
height=height,
width=width,
# Conditions
- condition_images=images[b],
- condition=condition[b],
- first_frame_mask=first_frame_mask, # Possibly None
- image_embeds=image_embeds[b] if image_embeds is not None else None,
+ condition_images=condition_rows[b],
+ latent_condition=condition[b],
+ first_frame_mask=(first_frame_mask[b] if first_frame_mask is not None else None),
+ image_embeds=per_sample_image_embeds[b],
# Prompt info
prompt=prompt[b] if isinstance(prompt, list) else prompt,
prompt_ids=prompt_ids[b] if prompt_ids is not None else None,
@@ -755,7 +812,7 @@ def forward(
guidance_scale_2: Optional[float] = None,
# Optional for I2V
image_embeds: Optional[torch.Tensor] = None,
- condition: Optional[torch.Tensor] = None,
+ latent_condition: Optional[torch.Tensor] = None,
first_frame_mask: Optional[torch.Tensor] = None,
boundary_timestep: Optional[float] = None,
# Next timestep info
@@ -780,7 +837,7 @@ def forward(
Args:
t: Current timestep tensor.
latents: Current latent representations.
- condition: Condition latents (first frame encoded).
+ latent_condition: Condition latents (first/optional-last frames encoded).
prompt_embeds: Text prompt embeddings.
negative_prompt_embeds: Optional negative prompt embeddings (for CFG).
guidance_scale: CFG scale for transformer (wan2.1 / wan2.2 high-noise).
@@ -809,6 +866,17 @@ def forward(
else self.pipeline.transformer_2.dtype
)
device = latents.device
+ if latent_condition is None:
+ raise ValueError("Wan I2V forward requires realized VAE condition channels")
+ if self.pipeline.config.expand_timesteps and first_frame_mask is None:
+ raise ValueError("Wan I2V expand_timesteps forward requires first_frame_mask")
+ if not self.pipeline.config.expand_timesteps and first_frame_mask is not None:
+ raise ValueError("Wan I2V non-expanded forward must not receive first_frame_mask")
+ if image_embeds is not None:
+ image_embeds = normalize_wan_image_embeds(
+ image_embeds,
+ batch_size=batch_size,
+ ).to(device=device, dtype=dtype)
# Determine boundary timestep
if boundary_timestep is None and self.pipeline.config.boundary_ratio is not None:
@@ -840,13 +908,15 @@ def forward(
# Prepare latent model input based on wan version
if first_frame_mask is not None:
# wan2.2: expand timesteps with mask
- latent_model_input = (1 - first_frame_mask) * condition + first_frame_mask * latents
+ latent_model_input = (
+ 1 - first_frame_mask
+ ) * latent_condition + first_frame_mask * latents
latent_model_input = latent_model_input.to(dtype)
temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten()
timestep = temp_ts.unsqueeze(0).expand(batch_size, -1)
else:
# wan2.1: concatenate condition
- latent_model_input = torch.cat([latents, condition], dim=1).to(dtype)
+ latent_model_input = torch.cat([latents, latent_condition], dim=1).to(dtype)
timestep = t.expand(batch_size)
# Conditional forward pass
@@ -873,6 +943,9 @@ def forward(
)[0]
velocity = velocity_uncond + current_guidance_scale * (velocity - velocity_uncond)
+ if not compute_log_prob and next_latents is None and tuple(return_kwargs) == ("velocity",):
+ return UniPCMultistepSDESchedulerOutput(velocity=velocity)
+
# Scheduler step
output = self.scheduler.step(
velocity=velocity,
diff --git a/src/flow_factory/models/wan/wan2_t2v.py b/src/flow_factory/models/wan/wan2_t2v.py
index 8039e0124..908d0b383 100644
--- a/src/flow_factory/models/wan/wan2_t2v.py
+++ b/src/flow_factory/models/wan/wan2_t2v.py
@@ -15,12 +15,7 @@
# src/flow_factory/models/wan/wan2_t2v.py
from __future__ import annotations
-import logging
-import math
-import os
-from collections import defaultdict
from dataclasses import dataclass
-from numbers import Real
from types import MappingProxyType
from typing import Any, ClassVar, Dict, List, Literal, Mapping, Optional, Tuple, Union
@@ -31,15 +26,12 @@
from peft import PeftModel
from PIL import Image
-from ...contracts import GeometrySource, NegativePromptPolicy, RateRequirement
+from ...contracts import BatchCapability, GeometrySource, NegativePromptPolicy, RateRequirement
from ...hparams import *
from ...samples import T2VSample
from ...scheduler import UniPCMultistepSDEScheduler, UniPCMultistepSDESchedulerOutput
-from ...utils.base import filter_kwargs
from ...utils.logger_utils import setup_logger
from ...utils.trajectory_collector import (
- CallbackCollector,
- TrajectoryCollector,
TrajectoryIndicesType,
create_callback_collector,
create_trajectory_collector,
@@ -47,7 +39,13 @@
from ..abc import BaseAdapter
from ..output_state import DecodedMediaBatch, EncodedOutputState, OutputStateCodec
from ..pipeline_contracts import video_output_contract
-from ._output import WanVideoOutputCodec
+from ._output import (
+ WanVideoOutputCodec,
+ configured_wan_video_output_geometry,
+ normalize_wan_video_latents,
+ resample_wan_output_video,
+ validate_wan_encoded_output_geometry,
+)
logger = setup_logger(__name__)
@@ -76,6 +74,7 @@ class Wan2_T2V_Adapter(BaseAdapter):
negative_prompt=NegativePromptPolicy.OPTIONAL,
output_fps=RateRequirement.REQUIRED,
geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
)
def __init__(self, config: Arguments, accelerator: Accelerator):
@@ -265,52 +264,7 @@ def build_output_state_codec(self) -> OutputStateCodec:
def _configured_video_output_geometry(self) -> Tuple[int, int, int, float]:
"""Return configured Wan geometry after exact latent-grid validation."""
- geometry = []
- for name in ("height", "width", "num_frames"):
- value = getattr(self.training_args, name, None)
- if type(value) is not int or value <= 0:
- raise ValueError(
- f"Wan output geometry requires positive integer train.{name}, "
- f"received {value!r}"
- )
- geometry.append(value)
- frame_rate = getattr(self.training_args, "frame_rate", None)
- if isinstance(frame_rate, bool) or not isinstance(frame_rate, Real):
- raise TypeError(
- "Wan output geometry requires finite positive train.frame_rate, "
- f"received {type(frame_rate).__name__}: {frame_rate!r}"
- )
- frame_rate = float(frame_rate)
- if not math.isfinite(frame_rate) or frame_rate <= 0:
- raise ValueError(
- "Wan output geometry requires finite positive train.frame_rate, "
- f"received {frame_rate!r}"
- )
-
- height, width, num_frames = geometry
- temporal_scale = self.pipeline.vae_scale_factor_temporal
- spatial_scale = self.pipeline.vae_scale_factor_spatial
- if (num_frames - 1) % temporal_scale:
- raise ValueError(
- "Wan output num_frames must satisfy "
- f"(num_frames - 1) % {temporal_scale} == 0, received {num_frames}"
- )
- transformer = (
- self.pipeline.transformer
- if self.pipeline.transformer is not None
- else self.pipeline.transformer_2
- )
- if transformer is None:
- raise RuntimeError("Wan output geometry requires one materialized transformer")
- patch_size = transformer.config.patch_size
- height_multiple = spatial_scale * patch_size[1]
- width_multiple = spatial_scale * patch_size[2]
- if height % height_multiple or width % width_multiple:
- raise ValueError(
- "Wan output height/width must be divisible by transformer latent-grid "
- f"multiples {(height_multiple, width_multiple)}, received {(height, width)}"
- )
- return height, width, num_frames, frame_rate
+ return configured_wan_video_output_geometry(self)
@staticmethod
def _resample_output_video(
@@ -321,63 +275,16 @@ def _resample_output_video(
target_fps: float,
) -> np.ndarray:
"""Select deterministic nearest-time frames for configured target cadence."""
- if video.dtype != np.uint8 or video.ndim != 4 or video.shape[-1] != 3:
- raise ValueError(
- "Wan decoded target video must be uint8 RGB shaped (F,H,W,3), "
- f"received dtype={video.dtype}, shape={tuple(video.shape)}"
- )
- if video.shape[0] < 1:
- raise ValueError("Wan decoded target video must contain at least one frame")
- if isinstance(source_fps, bool) or not isinstance(source_fps, Real):
- raise TypeError(
- "Wan target video requires source fps metadata, "
- f"received {type(source_fps).__name__}: {source_fps!r}"
- )
- source_fps = float(source_fps)
- if not math.isfinite(source_fps) or source_fps <= 0:
- raise ValueError(f"Wan target video requires positive finite fps, got {source_fps!r}")
- indices = np.rint(
- np.arange(target_frames, dtype=np.float64) * source_fps / target_fps
- ).astype(np.int64)
- if indices[-1] >= video.shape[0]:
- required_duration = (target_frames - 1) / target_fps
- available_duration = (video.shape[0] - 1) / source_fps
- raise ValueError(
- "Wan target video is too short for configured temporal geometry: "
- f"requires {required_duration:.6f}s, has {available_duration:.6f}s"
- )
- return np.ascontiguousarray(video[indices])
+ return resample_wan_output_video(
+ video,
+ source_fps=source_fps,
+ target_frames=target_frames,
+ target_fps=target_fps,
+ )
def _normalize_output_video_latents(self, latents: torch.Tensor) -> torch.Tensor:
"""Apply the exact inverse of Wan's existing decode normalization."""
- if not isinstance(latents, torch.Tensor) or latents.ndim != 5:
- raise ValueError(
- "Wan VAE target latents must be rank-5 BCFHW, "
- f"received {type(latents).__name__} with shape "
- f"{getattr(latents, 'shape', None)}"
- )
- config = self.vae.config
- z_dim = config.z_dim
- if latents.shape[1] != z_dim:
- raise ValueError(
- f"Wan VAE target latent channels must equal z_dim={z_dim}, "
- f"received {latents.shape[1]}"
- )
- latents_mean = torch.as_tensor(
- config.latents_mean,
- device=latents.device,
- dtype=latents.dtype,
- ).view(1, z_dim, 1, 1, 1)
- inverse_std = (
- torch.as_tensor(
- config.latents_std,
- device=latents.device,
- dtype=latents.dtype,
- )
- .reciprocal()
- .view(1, z_dim, 1, 1, 1)
- )
- return (latents - latents_mean) * inverse_std
+ return normalize_wan_video_latents(self, latents)
def _validate_encoded_output_geometry(
self,
@@ -386,39 +293,7 @@ def _validate_encoded_output_geometry(
encoded: EncodedOutputState,
) -> None:
"""Require encoded signatures and decode metadata to match train geometry."""
- del condition
- height, width, num_frames, frame_rate = self._configured_video_output_geometry()
- if len(encoded.geometry_signatures) != len(media_batch):
- raise ValueError(
- "Wan output codec must return one geometry signature per sample, "
- f"received {len(encoded.geometry_signatures)} for {len(media_batch)}"
- )
- for sample_index, signature in enumerate(encoded.geometry_signatures):
- geometry = signature.media[0]
- received = (
- geometry.height,
- geometry.width,
- geometry.frames,
- geometry.fps,
- )
- expected = (height, width, num_frames, frame_rate)
- if received != expected:
- raise ValueError(
- "Wan encoded output geometry disagrees with configured geometry for "
- f"sample {sample_index}: expected {expected}, received {received}"
- )
- expected_context = {
- "height": height,
- "width": width,
- "num_frames": num_frames,
- "frame_rate": frame_rate,
- }
- for name, expected in expected_context.items():
- if encoded.decode_context.get(name) != expected:
- raise ValueError(
- f"Wan decode_context {name!r} must equal {expected!r}, "
- f"received {encoded.decode_context.get(name)!r}"
- )
+ validate_wan_encoded_output_geometry(self, media_batch, condition, encoded)
def decode_latents(
self, latents: torch.Tensor, output_type: Literal["pt", "pil", "np"] = "pil"
@@ -779,6 +654,9 @@ def forward(
)[0]
velocity = velocity_uncond + current_guidance_scale * (velocity - velocity_uncond)
+ if not compute_log_prob and next_latents is None and tuple(return_kwargs) == ("velocity",):
+ return UniPCMultistepSDESchedulerOutput(velocity=velocity)
+
# 5. Scheduler step
output = self.scheduler.step(
velocity=velocity,
diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py
index f82186a2d..d84e9b08d 100644
--- a/src/flow_factory/trainers/abc.py
+++ b/src/flow_factory/trainers/abc.py
@@ -1121,6 +1121,7 @@ def _load_inference_components(self, trainable_module_names: List[str]):
ONLINE_EXECUTION_CONTRACT,
)
if execution_contract.acquisition is AcquisitionMode.DATASET:
+ modules_to_load.extend(self.adapter.condition_state_encoding_modules)
modules_to_load.extend(self.adapter.output_state_encoding_modules)
if not self.config.data_args.enable_preprocess:
diff --git a/src/flow_factory/trainers/common/flow_matching.py b/src/flow_factory/trainers/common/flow_matching.py
index 236b9dea7..18ad79d57 100644
--- a/src/flow_factory/trainers/common/flow_matching.py
+++ b/src/flow_factory/trainers/common/flow_matching.py
@@ -164,16 +164,20 @@ def flow_matching_per_sample_loss(
)
squared_errors[name] = (predicted.float() - target.float()).square()
- reduced = adapter.reduce_latent_values(squared_errors, state=noised.state)
+ reduced = adapter.reduce_flow_matching_objective_values(
+ squared_errors,
+ state=noised.state,
+ )
if not isinstance(reduced, torch.Tensor):
raise TypeError(
- "adapter.reduce_latent_values must return torch.Tensor, "
+ "adapter.reduce_flow_matching_objective_values must return torch.Tensor, "
f"received {type(reduced).__name__}"
)
batch_size = next(iter(squared_errors.values())).shape[0]
if reduced.shape != (batch_size,):
raise ValueError(
- "adapter.reduce_latent_values must return one value per sample with shape "
+ "adapter.reduce_flow_matching_objective_values must return one value per sample "
+ "with shape "
f"({batch_size},), received {tuple(reduced.shape)}"
)
if reduced.dtype is not torch.float32:
diff --git a/src/flow_factory/trainers/common/offline_batch.py b/src/flow_factory/trainers/common/offline_batch.py
index 4b25d3e2d..a87712640 100644
--- a/src/flow_factory/trainers/common/offline_batch.py
+++ b/src/flow_factory/trainers/common/offline_batch.py
@@ -21,6 +21,7 @@
import torch
from ...contracts import NON_MODEL_CONDITION_KEYS
+from ...models.condition_state import PreparedConditionState
def move_condition_to_device(
@@ -79,6 +80,30 @@ def bind_output_forward_context(
return {**condition, **forward_context}
+def bind_prepared_condition_output(
+ prepared: PreparedConditionState,
+ output_forward_context: Mapping[str, Any],
+) -> Dict[str, Any]:
+ """Bind one realized input condition to candidate-specific output fields.
+
+ Args:
+ prepared: Input-owned realization shared by every candidate for a request.
+ output_forward_context: Candidate-specific output-derived model fields.
+
+ Returns:
+ A new complete model-forward mapping.
+ """
+ if not isinstance(prepared, PreparedConditionState):
+ raise TypeError(
+ "expected PreparedConditionState for prepared offline condition, "
+ f"received {type(prepared).__name__}: {prepared!r}"
+ )
+ return bind_output_forward_context(
+ prepared.model_forward_condition(),
+ output_forward_context,
+ )
+
+
def _move_condition_value(
value: Any,
device: torch.device,
@@ -136,4 +161,8 @@ def _reject_non_model_keys(value: Mapping[str, Any], identifier: str) -> None:
)
-__all__ = ["bind_output_forward_context", "move_condition_to_device"]
+__all__ = [
+ "bind_output_forward_context",
+ "bind_prepared_condition_output",
+ "move_condition_to_device",
+]
diff --git a/src/flow_factory/trainers/common/runtime_identity.py b/src/flow_factory/trainers/common/runtime_identity.py
index 47f498ba0..356df6f9d 100644
--- a/src/flow_factory/trainers/common/runtime_identity.py
+++ b/src/flow_factory/trainers/common/runtime_identity.py
@@ -128,6 +128,11 @@ def build_default_execution_identity_payload(trainer: Any) -> dict[str, Any]:
execution_contract = type(trainer).execution_contract
acquisition = getattr(execution_contract, "acquisition", None)
feedback = getattr(execution_contract, "feedback", None)
+ pipeline_io_contract = getattr(
+ trainer.adapter,
+ "effective_pipeline_io_contract",
+ None,
+ )
return {
"contract": {
"acquisition": _enum_identity_value(
@@ -140,6 +145,7 @@ def build_default_execution_identity_payload(trainer: Any) -> dict[str, Any]:
),
"paradigm": getattr(type(trainer), "paradigm", None),
},
+ "pipeline_io_contract": pipeline_io_contract,
"training": training,
"scheduler": scheduler,
"realized_scheduler_group": _scheduler_group_schema(trainer.adapter),
diff --git a/src/flow_factory/trainers/offline/offline_dpo.py b/src/flow_factory/trainers/offline/offline_dpo.py
index 31d94d011..1786f1cb7 100644
--- a/src/flow_factory/trainers/offline/offline_dpo.py
+++ b/src/flow_factory/trainers/offline/offline_dpo.py
@@ -34,7 +34,7 @@
validate_preference_component_times,
validate_preference_output_states,
)
-from ..common.offline_batch import bind_output_forward_context, move_condition_to_device
+from ..common.offline_batch import bind_prepared_condition_output, move_condition_to_device
from ..forward_process import forward_velocity_state
MetricAccumulator = Dict[str, List[torch.Tensor]]
@@ -53,7 +53,7 @@ def _build_train_dataloader(self) -> Tuple[DataLoader, Dict[str, DataLoader]]:
accelerator=self.accelerator,
preprocess_func=self.adapter.preprocess_func,
supervision_type="preference",
- pipeline_io_contract=self.adapter.pipeline_io_contract,
+ pipeline_io_contract=self.adapter.effective_pipeline_io_contract,
)
return dataloader, {}
@@ -63,12 +63,19 @@ def optimize_batch(self, batch: Any) -> None:
self.adapter.train()
condition = move_condition_to_device(batch.condition, self.accelerator.device)
- chosen = self.adapter.encode_output_state(preference.chosen_media, condition)
- rejected = self.adapter.encode_output_state(preference.rejected_media, condition)
+ prepared_condition = self.adapter.prepare_condition_state(condition)
+ chosen = self.adapter.encode_output_state(preference.chosen_media, prepared_condition)
+ rejected = self.adapter.encode_output_state(preference.rejected_media, prepared_condition)
validate_preference_output_states(chosen, rejected)
- chosen_batch = bind_output_forward_context(condition, chosen.forward_context)
- rejected_batch = bind_output_forward_context(condition, rejected.forward_context)
+ chosen_batch = bind_prepared_condition_output(
+ prepared_condition,
+ chosen.forward_context,
+ )
+ rejected_batch = bind_prepared_condition_output(
+ prepared_condition,
+ rejected.forward_context,
+ )
all_timesteps = sample_offline_timesteps(
self.training_args,
batch_size=len(preference.chosen_media),
diff --git a/src/flow_factory/trainers/offline/sft.py b/src/flow_factory/trainers/offline/sft.py
index 9e4885fd5..039966fa6 100644
--- a/src/flow_factory/trainers/offline/sft.py
+++ b/src/flow_factory/trainers/offline/sft.py
@@ -31,7 +31,7 @@
flow_matching_per_sample_loss,
sample_offline_timesteps,
)
-from ..common.offline_batch import bind_output_forward_context, move_condition_to_device
+from ..common.offline_batch import bind_prepared_condition_output, move_condition_to_device
from ..forward_process import forward_velocity_state
@@ -48,7 +48,7 @@ def _build_train_dataloader(self) -> Tuple[DataLoader, Dict[str, DataLoader]]:
accelerator=self.accelerator,
preprocess_func=self.adapter.preprocess_func,
supervision_type="demonstration",
- pipeline_io_contract=self.adapter.pipeline_io_contract,
+ pipeline_io_contract=self.adapter.effective_pipeline_io_contract,
)
return dataloader, {}
@@ -69,8 +69,12 @@ def optimize_batch(self, batch: Any) -> None:
# microstep explicitly restores training mode before policy execution.
self.adapter.train()
- encoded = self.adapter.encode_output_state(output.target_media, condition)
- model_batch = bind_output_forward_context(condition, encoded.forward_context)
+ prepared_condition = self.adapter.prepare_condition_state(condition)
+ encoded = self.adapter.encode_output_state(output.target_media, prepared_condition)
+ model_batch = bind_prepared_condition_output(
+ prepared_condition,
+ encoded.forward_context,
+ )
all_timesteps = sample_offline_timesteps(
self.training_args,
batch_size=len(output.target_media),
diff --git a/tests/contracts/test_pipeline_io_contract.py b/tests/contracts/test_pipeline_io_contract.py
index bb8684bb4..73ddabd43 100644
--- a/tests/contracts/test_pipeline_io_contract.py
+++ b/tests/contracts/test_pipeline_io_contract.py
@@ -35,6 +35,7 @@
OutputMediaSequence,
PipelineIOContract,
RateRequirement,
+ resolve_pipeline_input_media_slots,
validate_pipeline_model_input,
validate_pipeline_output_candidate,
)
@@ -153,6 +154,7 @@ class _InputMediaFixture:
type: str
fps: float | None = None
sample_rate: int | None = None
+ slot: str | None = None
@dataclass
@@ -266,6 +268,195 @@ def test_model_input_validation_enforces_counts_and_required_rates() -> None:
)
+def test_semantic_slots_preserve_positional_shorthand_and_support_sparse_binding() -> None:
+ contract = PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=(
+ InputMediaRule(
+ format=IMAGE_FORMAT,
+ min_count=1,
+ max_count=2,
+ slots=("first_frame", "last_frame"),
+ ),
+ ),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.WITHIN_TYPE,
+ ),
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ output_media=OutputMediaSequence(items=(IMAGE_FORMAT,)),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
+
+ positional = _ModelInputFixture(
+ prompt="both",
+ media=(_InputMediaFixture("image"), _InputMediaFixture("image")),
+ )
+ assert resolve_pipeline_input_media_slots(positional, contract) == (
+ "first_frame",
+ "last_frame",
+ )
+
+ last_only = _ModelInputFixture(
+ prompt="end here",
+ media=(_InputMediaFixture("image", slot="last_frame"),),
+ )
+ assert resolve_pipeline_input_media_slots(last_only, contract) == ("last_frame",)
+
+ mixed = _ModelInputFixture(
+ prompt="explicit last first in the manifest",
+ media=(
+ _InputMediaFixture("image", slot="last_frame"),
+ _InputMediaFixture("image"),
+ ),
+ )
+ assert resolve_pipeline_input_media_slots(mixed, contract) == (
+ "last_frame",
+ "first_frame",
+ )
+
+
+def test_multi_slot_rules_require_within_type_positional_semantics() -> None:
+ """A contract cannot claim order-insensitivity while using positional fallback."""
+ with pytest.raises(ValueError, match="multi-slot.*within_type ordering"):
+ InputMediaSpec(
+ rules=(
+ InputMediaRule(
+ format=IMAGE_FORMAT,
+ min_count=1,
+ max_count=2,
+ slots=("first_frame", "last_frame"),
+ ),
+ ),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.INSENSITIVE,
+ )
+
+
+def test_semantic_slots_reject_unknown_duplicate_and_missing_required_bindings() -> None:
+ rule = InputMediaRule(
+ format=IMAGE_FORMAT,
+ min_count=1,
+ max_count=2,
+ slots=("first_frame", "last_frame"),
+ required_slots=("first_frame",),
+ )
+ contract = PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=(rule,),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.WITHIN_TYPE,
+ ),
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ output_media=OutputMediaSequence(items=(IMAGE_FORMAT,)),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
+
+ with pytest.raises(ValueError, match="requires input media slots.*first_frame"):
+ validate_pipeline_model_input(
+ _ModelInputFixture(
+ prompt="last only",
+ media=(_InputMediaFixture("image", slot="last_frame"),),
+ ),
+ contract,
+ )
+ with pytest.raises(ValueError, match="slot 'middle_frame' is not accepted"):
+ validate_pipeline_model_input(
+ _ModelInputFixture(
+ prompt="unknown",
+ media=(_InputMediaFixture("image", slot="middle_frame"),),
+ ),
+ contract,
+ )
+ with pytest.raises(ValueError, match="assigned more than once"):
+ validate_pipeline_model_input(
+ _ModelInputFixture(
+ prompt="duplicate",
+ media=(
+ _InputMediaFixture("image", slot="first_frame"),
+ _InputMediaFixture("image", slot="first_frame"),
+ ),
+ ),
+ contract,
+ )
+
+
+def test_aggregate_input_constraints_cover_cross_modality_invariants() -> None:
+ video = MediaFormat(
+ type=MediaType.VIDEO,
+ fps=RateRequirement.OPTIONAL,
+ sample_rate=RateRequirement.NOT_APPLICABLE,
+ )
+ audio = MediaFormat(
+ type=MediaType.AUDIO,
+ fps=RateRequirement.NOT_APPLICABLE,
+ sample_rate=RateRequirement.OPTIONAL,
+ )
+ contract = PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=(
+ InputMediaRule(IMAGE_FORMAT, min_count=0, max_count=9),
+ InputMediaRule(video, min_count=0, max_count=3),
+ InputMediaRule(audio, min_count=0, max_count=3),
+ ),
+ binding=InputMediaBinding.ORDERED_REFERENCES,
+ order=InputMediaOrder.GLOBAL,
+ min_total_count=1,
+ max_total_count=12,
+ required_any_types=(MediaType.IMAGE, MediaType.VIDEO),
+ ),
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ output_media=OutputMediaSequence(items=(IMAGE_FORMAT,)),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
+
+ with pytest.raises(ValueError, match="at least 1 input media item"):
+ validate_pipeline_model_input(_ModelInputFixture(prompt="empty"), contract)
+ with pytest.raises(ValueError, match="whose type is in.*image.*video"):
+ validate_pipeline_model_input(
+ _ModelInputFixture(
+ prompt="audio only",
+ media=(_InputMediaFixture("audio", sample_rate=16000),),
+ ),
+ contract,
+ )
+ validate_pipeline_model_input(
+ _ModelInputFixture(
+ prompt="valid",
+ media=(
+ _InputMediaFixture("audio", sample_rate=16000),
+ _InputMediaFixture("image"),
+ ),
+ ),
+ contract,
+ )
+
+
+@pytest.mark.parametrize(
+ ("min_total_count", "max_total_count", "match"),
+ (
+ (None, 1, "max_total_count cannot be smaller.*per-type minimums"),
+ (3, None, "min_total_count cannot exceed.*per-type maximums"),
+ ),
+)
+def test_aggregate_input_constraints_reject_impossible_rule_combinations(
+ min_total_count: int | None,
+ max_total_count: int | None,
+ match: str,
+) -> None:
+ """Unsatisfiable aggregate and per-type bounds fail at declaration time."""
+ with pytest.raises(ValueError, match=match):
+ InputMediaSpec(
+ rules=(InputMediaRule(IMAGE_FORMAT, min_count=2, max_count=2),),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.WITHIN_TYPE,
+ min_total_count=min_total_count,
+ max_total_count=max_total_count,
+ )
+
+
@pytest.mark.parametrize(
"kwargs,match",
[
@@ -470,7 +661,7 @@ def test_pipeline_contract_rejects_raw_policy_values_and_algorithm_shape_fields(
def test_input_media_geometry_requires_a_guaranteed_input() -> None:
"""A conditional geometry source cannot rely on an optional-only input layout."""
- with pytest.raises(ValueError, match="requires at least one input media rule"):
+ with pytest.raises(ValueError, match="constraints that guarantee at least one"):
PipelineIOContract(
input_media=InputMediaSpec(
rules=(InputMediaRule(format=IMAGE_FORMAT, min_count=0, max_count=1),),
@@ -482,3 +673,59 @@ def test_input_media_geometry_requires_a_guaranteed_input() -> None:
geometry_source=GeometrySource.INPUT_MEDIA,
batch_capability=BatchCapability.UNIFORM,
)
+
+
+@pytest.mark.parametrize(
+ "aggregate_fields",
+ [
+ {"min_total_count": 1},
+ {"required_any_types": (MediaType.IMAGE,)},
+ ],
+)
+def test_input_media_geometry_accepts_aggregate_nonempty_guarantees(
+ aggregate_fields: dict[str, object],
+) -> None:
+ contract = PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=(InputMediaRule(format=IMAGE_FORMAT, min_count=0, max_count=1),),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.INSENSITIVE,
+ **aggregate_fields,
+ ),
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ output_media=OutputMediaSequence(items=(IMAGE_FORMAT,)),
+ geometry_source=GeometrySource.INPUT_MEDIA,
+ batch_capability=BatchCapability.UNIFORM,
+ )
+
+ assert contract.geometry_source is GeometrySource.INPUT_MEDIA
+
+
+def test_required_any_types_requires_canonical_media_type_order() -> None:
+ video = MediaFormat(
+ type=MediaType.VIDEO,
+ fps=RateRequirement.OPTIONAL,
+ sample_rate=RateRequirement.NOT_APPLICABLE,
+ )
+
+ with pytest.raises(ValueError, match="required_any_types must use canonical type order"):
+ InputMediaSpec(
+ rules=(
+ InputMediaRule(format=IMAGE_FORMAT, min_count=0, max_count=1),
+ InputMediaRule(format=video, min_count=0, max_count=1),
+ ),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.INSENSITIVE,
+ required_any_types=(MediaType.VIDEO, MediaType.IMAGE),
+ )
+
+
+def test_required_slots_require_declared_slot_order() -> None:
+ with pytest.raises(ValueError, match="required input media slots must use declared slot order"):
+ InputMediaRule(
+ format=IMAGE_FORMAT,
+ min_count=2,
+ max_count=2,
+ slots=("first_frame", "last_frame"),
+ required_slots=("last_frame", "first_frame"),
+ )
diff --git a/tests/data_utils/test_offline_condition_cache.py b/tests/data_utils/test_offline_condition_cache.py
index 926c6fb6f..9f227046c 100644
--- a/tests/data_utils/test_offline_condition_cache.py
+++ b/tests/data_utils/test_offline_condition_cache.py
@@ -15,6 +15,7 @@
import gc
import json
import weakref
+from dataclasses import replace
from pathlib import Path
from typing import Any, Dict, List
@@ -25,6 +26,20 @@
from datasets import Image as HFImage
from PIL import Image
+from flow_factory.contracts import (
+ BatchCapability,
+ GeometrySource,
+ InputMediaBinding,
+ InputMediaOrder,
+ InputMediaRule,
+ InputMediaSpec,
+ MediaFormat,
+ MediaType,
+ NegativePromptPolicy,
+ OutputMediaSequence,
+ PipelineIOContract,
+ RateRequirement,
+)
from flow_factory.data_utils.dataset import GeneralDataset, _cross_chunk_schema_probe_batches
from flow_factory.data_utils.offline_condition_cache import (
build_offline_condition_cache,
@@ -33,10 +48,47 @@
)
from flow_factory.data_utils.offline_dataset import (
OFFLINE_CONDITION_ID_COLUMN,
+ OfflineDataset,
compute_offline_condition_id,
)
+from flow_factory.data_utils.offline_loader import build_offline_dataloader
from flow_factory.data_utils.schema import NormalizedDatasetRecord, normalize_v2_record
+_IMAGE_FORMAT = MediaFormat(
+ type=MediaType.IMAGE,
+ fps=RateRequirement.NOT_APPLICABLE,
+ sample_rate=RateRequirement.NOT_APPLICABLE,
+)
+_SLOTTED_IMAGE_CONTRACT = PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=(
+ InputMediaRule(
+ format=_IMAGE_FORMAT,
+ min_count=1,
+ max_count=2,
+ slots=("first_frame", "last_frame"),
+ ),
+ ),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.WITHIN_TYPE,
+ ),
+ negative_prompt=NegativePromptPolicy.UNSUPPORTED,
+ output_media=OutputMediaSequence(items=(_IMAGE_FORMAT,)),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+)
+_OPTIONAL_IMAGE_CONTRACT = PipelineIOContract(
+ input_media=InputMediaSpec(
+ rules=(InputMediaRule(format=_IMAGE_FORMAT, min_count=0, max_count=1),),
+ binding=InputMediaBinding.GROUPED_BY_TYPE,
+ order=InputMediaOrder.INSENSITIVE,
+ ),
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ output_media=OutputMediaSequence(items=(_IMAGE_FORMAT,)),
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.UNIFORM,
+)
+
class CountingPreprocessor:
def __init__(self) -> None:
@@ -114,6 +166,37 @@ def preprocess(
}
+class OptionalSourcePreprocessor:
+ """Expose stable output keys for optional condition columns across sources."""
+
+ def preprocess(
+ self,
+ prompt: List[str],
+ negative_prompt: List[str],
+ images: List[List[Image.Image]],
+ ) -> Dict[str, torch.Tensor]:
+ return {
+ "prompt_lengths": torch.tensor([len(value) for value in prompt]),
+ "negative_prompt_lengths": torch.tensor([len(value) for value in negative_prompt]),
+ "image_counts": torch.tensor([len(value) for value in images]),
+ }
+
+
+class SingleSamplePreprocessor:
+ def __init__(self) -> None:
+ self.batch_sizes: List[int] = []
+
+ def preprocess(
+ self,
+ prompt: List[str],
+ images: List[List[Image.Image]],
+ **kwargs: Any,
+ ) -> Dict[str, torch.Tensor]:
+ del images, kwargs
+ self.batch_sizes.append(len(prompt))
+ return {"encoded": torch.ones(len(prompt), 1)}
+
+
class ChunkBoundedDatasetSpy:
def __init__(
self,
@@ -324,6 +407,33 @@ def test_projection_contains_only_input_and_identity_columns(tmp_path: Path) ->
assert "private-target.png" not in repr(projected[0])
+def test_projection_canonicalizes_explicit_and_positional_semantic_slots(
+ tmp_path: Path,
+) -> None:
+ Image.new("RGB", (2, 2), color=(1, 2, 3)).save(tmp_path / "first.png")
+ Image.new("RGB", (2, 2), color=(4, 5, 6)).save(tmp_path / "last.png")
+ record = _demonstration_record(
+ tmp_path,
+ input_media=[
+ {"type": "image", "path": "last.png", "slot": "last_frame"},
+ {"type": "image", "path": "first.png"},
+ ],
+ )
+
+ projected = project_offline_condition_dataset(
+ [record],
+ source_name="slotted",
+ ordered_references=False,
+ pipeline_io_contract=_SLOTTED_IMAGE_CONTRACT,
+ )
+
+ assert projected[0]["images"] == [
+ str(tmp_path / "first.png"),
+ str(tmp_path / "last.png"),
+ ]
+ assert projected[0]["image_slots"] == ["first_frame", "last_frame"]
+
+
def test_projection_normalizes_missing_optional_negative_prompts_in_mixed_batch(
tmp_path: Path,
) -> None:
@@ -340,6 +450,69 @@ def test_projection_normalizes_missing_optional_negative_prompts_in_mixed_batch(
assert projected["negative_prompt"] == ["", "bad quality"]
+def test_optional_condition_columns_are_stable_across_offline_sources(
+ tmp_path: Path,
+) -> None:
+ """ConcatDataset batches can mix an all-empty source with a populated source."""
+ Image.new("RGB", (2, 2), color=(1, 2, 3)).save(tmp_path / "reference.png")
+ Image.new("RGB", (2, 2), color=(4, 5, 6)).save(tmp_path / "target.png")
+ empty_record = _demonstration_record(tmp_path)
+ populated_record = _demonstration_record(
+ tmp_path,
+ input_media=[{"type": "image", "path": "reference.png"}],
+ negative_prompt="bad quality",
+ )
+ preprocessor = OptionalSourcePreprocessor()
+
+ empty_cache = build_offline_condition_cache(
+ [empty_record],
+ source_name="empty-source",
+ dataset_dir=tmp_path,
+ cache_dir=tmp_path / "cache",
+ preprocess_func=preprocessor.preprocess,
+ pipeline_io_contract=_OPTIONAL_IMAGE_CONTRACT,
+ preprocessing_batch_size=1,
+ )
+ populated_cache = build_offline_condition_cache(
+ [populated_record],
+ source_name="populated-source",
+ dataset_dir=tmp_path,
+ cache_dir=tmp_path / "cache",
+ preprocess_func=preprocessor.preprocess,
+ pipeline_io_contract=_OPTIONAL_IMAGE_CONTRACT,
+ preprocessing_batch_size=1,
+ )
+ assert set(empty_cache.column_names) == set(populated_cache.column_names)
+
+ empty_dataset = OfflineDataset(
+ [empty_record],
+ empty_cache,
+ source_name="empty-source",
+ source_id=0,
+ supervision_type="demonstration",
+ )
+ populated_dataset = OfflineDataset(
+ [populated_record],
+ populated_cache,
+ source_name="populated-source",
+ source_id=1,
+ supervision_type="demonstration",
+ )
+ loader = build_offline_dataloader(
+ [empty_dataset, populated_dataset],
+ source_weights=[1, 1],
+ batch_size=2,
+ num_replicas=1,
+ rank=0,
+ gradient_accumulation_steps=1,
+ shuffle=False,
+ )
+
+ batch = next(iter(loader))
+ assert batch.condition["image_counts"].tolist() == [0, 1]
+ assert batch.condition["negative_prompt_lengths"].tolist() == [0, 11]
+
+
def test_preference_arms_never_enter_the_condition_projection(tmp_path: Path) -> None:
projected = project_offline_condition_dataset(
[_preference_record(tmp_path)],
@@ -359,6 +532,31 @@ def test_condition_source_hash_is_order_sensitive() -> None:
)
+def test_condition_source_hash_includes_effective_slot_projection() -> None:
+ reversed_rule = replace(
+ _SLOTTED_IMAGE_CONTRACT.input_media.rules[0],
+ slots=("last_frame", "first_frame"),
+ )
+ reversed_contract = replace(
+ _SLOTTED_IMAGE_CONTRACT,
+ input_media=replace(
+ _SLOTTED_IMAGE_CONTRACT.input_media,
+ rules=(reversed_rule,),
+ ),
+ )
+
+ baseline = compute_offline_condition_source_hash(
+ ["same-input"],
+ pipeline_io_contract=_SLOTTED_IMAGE_CONTRACT,
+ )
+ changed = compute_offline_condition_source_hash(
+ ["same-input"],
+ pipeline_io_contract=reversed_contract,
+ )
+
+ assert changed != baseline
+
+
def test_schema_probe_scans_only_bounded_pending_column_chunks() -> None:
dataset = ChunkBoundedDatasetSpy(
[
@@ -523,6 +721,103 @@ def datasets_3_3_compatible_map(
assert cache[row]["image_latent_ids"].shape == (2, 4)
+def test_distributed_condition_cache_uses_one_global_optional_media_schema(
+ tmp_path: Path,
+) -> None:
+ """Disjoint empty/populated ranks consolidate with identical Arrow features."""
+ Image.new("RGB", (2, 2), color=(1, 2, 3)).save(tmp_path / "reference.png")
+ records = [
+ _demonstration_record(tmp_path),
+ _demonstration_record(tmp_path),
+ _demonstration_record(
+ tmp_path,
+ input_media=[{"type": "image", "path": "reference.png"}],
+ ),
+ _demonstration_record(
+ tmp_path,
+ input_media=[{"type": "image", "path": "reference.png"}],
+ ),
+ ]
+ raw_dataset = project_offline_condition_dataset(
+ records,
+ source_name="distributed-optional-images",
+ ordered_references=False,
+ pipeline_io_contract=_OPTIONAL_IMAGE_CONTRACT,
+ )
+ condition_ids = tuple(raw_dataset[OFFLINE_CONDITION_ID_COLUMN])
+ source_hash = compute_offline_condition_source_hash(
+ condition_ids,
+ pipeline_io_contract=_OPTIONAL_IMAGE_CONTRACT,
+ )
+ preprocessor = OptionalImagePreprocessor()
+ merged_cache_path = GeneralDataset.compute_cache_path(
+ dataset_dir=str(tmp_path),
+ split="train",
+ cache_dir=str(tmp_path / "cache"),
+ max_dataset_size=None,
+ preprocess_func=preprocessor.preprocess,
+ preprocess_kwargs={},
+ source_hash_override=source_hash,
+ )
+
+ for rank in range(2):
+ GeneralDataset(
+ dataset_dir=str(tmp_path),
+ split="train",
+ cache_dir=str(tmp_path / "cache"),
+ force_reprocess=True,
+ preprocessing_batch_size=2,
+ preprocess_func=preprocessor.preprocess,
+ num_shards=2,
+ shard_index=rank,
+ image_dir=str(tmp_path),
+ video_dir=str(tmp_path),
+ audio_dir=str(tmp_path),
+ target_arrow_path=GeneralDataset.build_part_arrow_path(
+ merged_cache_path,
+ rank,
+ 2,
+ ),
+ raw_dataset=raw_dataset,
+ source_hash_override=source_hash,
+ passthrough_columns=(OFFLINE_CONDITION_ID_COLUMN,),
+ )
+
+ GeneralDataset.consolidate_parts(merged_cache_path, 2, split="train")
+ merged = GeneralDataset.load_merged(merged_cache_path).processed_dataset
+
+ assert isinstance(merged.features["condition_images"].feature, HFImage)
+ assert merged[0]["condition_images"] == []
+ assert len(merged[2]["condition_images"]) == 1
+
+
+def test_standalone_condition_cache_honors_single_sample_contract(
+ tmp_path: Path,
+) -> None:
+ Image.new("RGB", (2, 2), color=(1, 2, 3)).save(tmp_path / "first.png")
+ records = [
+ _demonstration_record(
+ tmp_path,
+ input_media=[{"type": "image", "path": "first.png", "slot": "first_frame"}],
+ )
+ for _ in range(2)
+ ]
+ preprocessor = SingleSamplePreprocessor()
+
+ build_offline_condition_cache(
+ records,
+ source_name="single-sample",
+ dataset_dir=tmp_path,
+ cache_dir=tmp_path / "cache",
+ preprocess_func=preprocessor.preprocess,
+ pipeline_io_contract=_SLOTTED_IMAGE_CONTRACT,
+ force_reprocess=True,
+ preprocessing_batch_size=8,
+ )
+
+ assert preprocessor.batch_sizes == [1, 1]
+
+
def test_ordered_heterogeneous_references_cross_arrow_as_canonical_json(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
diff --git a/tests/data_utils/test_offline_train_data.py b/tests/data_utils/test_offline_train_data.py
index 42982bdea..de53c7310 100644
--- a/tests/data_utils/test_offline_train_data.py
+++ b/tests/data_utils/test_offline_train_data.py
@@ -786,7 +786,12 @@ def fake_create_or_load_dataset(**kwargs: Any) -> SimpleNamespace:
assert call["enable_distributed"] is True
assert call["preprocess_parallelism"] == "global"
raw_dataset = call["base_kwargs"]["raw_dataset"]
- assert set(raw_dataset.column_names) == {"prompt", OFFLINE_CONDITION_ID_COLUMN}
+ assert set(raw_dataset.column_names) == {
+ "prompt",
+ "negative_prompt",
+ OFFLINE_CONDITION_ID_COLUMN,
+ }
+ assert raw_dataset["negative_prompt"] == ["", "", "", ""]
assert "target-0.png" not in repr(raw_dataset[0])
assert "revision" not in repr(raw_dataset[0])
assert isinstance(loader.sampler, DistributedSampler)
diff --git a/tests/data_utils/test_schema.py b/tests/data_utils/test_schema.py
index b8c909f2c..9c09f5a30 100644
--- a/tests/data_utils/test_schema.py
+++ b/tests/data_utils/test_schema.py
@@ -109,6 +109,31 @@ def test_demonstration_normalization_preserves_media_order_and_resolves_paths(
assert normalized.supervision.target.media[0].path == str(absolute_target)
+def test_input_media_slot_is_normalized_but_output_media_rejects_slots(
+ tmp_path: Path,
+) -> None:
+ raw = _demonstration_record(
+ input={
+ "prompt": "End on this frame.",
+ "media": [
+ {
+ "type": "image",
+ "path": "ending.png",
+ "slot": "last_frame",
+ }
+ ],
+ }
+ )
+
+ normalized = normalize_v2_record(raw, dataset_dir=tmp_path)
+
+ assert normalized.model_input.media[0].slot == "last_frame"
+
+ raw["supervision"]["target"]["media"][0]["slot"] = "last_frame"
+ with pytest.raises(ValidationError, match="slot is input-only"):
+ DatasetRecordV2.model_validate(raw)
+
+
def test_preference_normalization_keeps_both_arms_under_one_input(tmp_path: Path) -> None:
raw = _demonstration_record(
supervision={
diff --git a/tests/docs/test_minimax_h3_docs.py b/tests/docs/test_minimax_h3_docs.py
index 2902cb745..bbfc174a4 100644
--- a/tests/docs/test_minimax_h3_docs.py
+++ b/tests/docs/test_minimax_h3_docs.py
@@ -71,7 +71,8 @@ def test_examples_readme_links_h3_and_separates_validation_levels() -> None:
relative_link = f"../{root_link}"
assert relative_link in text
assert (ROOT / root_link).is_file()
- assert "Schema/API validated only" in text
+ assert "schema/API and local offline-path validated" in text
+ assert "GPU validation plan" in text
assert "hardware" in text
assert "reward" in text
assert "61 GB" in text
@@ -80,6 +81,44 @@ def test_examples_readme_links_h3_and_separates_validation_levels() -> None:
assert "NonCommercial" in text
+def test_gpu_validation_plan_declares_the_complete_smoke_matrix() -> None:
+ text = _text("guidance/gpu_validation.md")
+
+ assert "10 x 3 x 4 = 120 jobs" in text
+ for mode in (
+ "sd35-t2i",
+ "bagel-mri2i",
+ "wan-t2v",
+ "wan-i2v-first",
+ "wan-flf2v",
+ "ltx2-t2av",
+ "ltx2-i2av",
+ "h3-t2va",
+ "h3-fl2va",
+ "h3-ref2va",
+ ):
+ assert f"`{mode}`" in text
+ for backend in ("ddp", "zero2", "fsdp2"):
+ assert f"`{backend}`" in text
+ for algorithm in ("grpo", "sft", "offline-dpo", "tdm"):
+ assert f"`{algorithm}`" in text
+ assert "exactly two rank-local dataloader batches" in text
+ assert "two training epochs" in text
+ assert "eval.eval_freq: 0" in text
+ assert "DistributedSampler" in text
+
+
+def test_install_docs_use_the_released_diffusers_runtime() -> None:
+ readme = _text("README.md")
+ dockerfile = _text("docker/docker-cuda/Dockerfile")
+ docker_readme = _text("docker/README.md")
+
+ assert "diffusers>=0.40.0" in readme
+ assert "pip install -e ./diffusers" not in readme
+ assert "pip install -e ./diffusers" not in dockerfile
+ assert "submodule (required)" not in docker_readme
+
+
def test_new_model_guide_documents_component_runtime_boundaries() -> None:
text = _text("guidance/new_model.md")
for required in (
diff --git a/tests/models/ltx2/test_ltx2_adapter_init.py b/tests/models/ltx2/test_ltx2_adapter_init.py
index f1c3b2587..ad362e40b 100644
--- a/tests/models/ltx2/test_ltx2_adapter_init.py
+++ b/tests/models/ltx2/test_ltx2_adapter_init.py
@@ -51,11 +51,18 @@ class PipelineStub:
def __init__(self) -> None:
self.scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)
self.transformer = TransformerStub()
+ self.vae = torch.nn.Identity()
+ self.audio_vae = torch.nn.Identity()
@property
def components(self) -> Dict[str, Any]:
"""Return the eager component declaration the classic runtime reads."""
- return {"scheduler": self.scheduler, "transformer": self.transformer}
+ return {
+ "scheduler": self.scheduler,
+ "transformer": self.transformer,
+ "vae": self.vae,
+ "audio_vae": self.audio_vae,
+ }
def _config() -> SimpleNamespace:
@@ -195,3 +202,23 @@ def test_constructor_dispatches_lifecycle_calls_in_component_order(cls: type) ->
adapter.set_trajectory_seed(21)
assert seeded == [("video", 21), ("audio", 21)]
+
+
+def test_i2av_preprocess_forwards_cached_negative_prompt_to_text_encoder() -> None:
+ adapter = object.__new__(LTX2_I2AV_Adapter)
+ captured: Dict[str, Any] = {}
+
+ def encode_prompt(**kwargs: Any) -> Dict[str, Any]:
+ captured.update(kwargs)
+ return {"connector_prompt_embeds": torch.zeros(1, 1, 1)}
+
+ adapter.encode_prompt = encode_prompt
+
+ adapter.preprocess_func(
+ prompt=["describe"],
+ negative_prompt=["avoid blur"],
+ images=None,
+ guidance_scale=4.0,
+ )
+
+ assert captured["negative_prompt"] == ["avoid blur"]
diff --git a/tests/models/ltx2/test_ltx2_component_hooks.py b/tests/models/ltx2/test_ltx2_component_hooks.py
index 2dbabbb04..f5cf3bf78 100644
--- a/tests/models/ltx2/test_ltx2_component_hooks.py
+++ b/tests/models/ltx2/test_ltx2_component_hooks.py
@@ -758,3 +758,32 @@ def test_legacy_concatenated_forward_stays_numerically_unchanged(cls: type) -> N
assert torch.allclose(output.next_latents, torch.cat([expected_video, expected_audio], dim=1))
assert torch.allclose(output.velocity, torch.cat([video_velocity, audio_velocity], dim=1))
assert output.log_prob.shape == (BATCH_SIZE,)
+
+
+@pytest.mark.parametrize("cls", [LTX2_T2AV_Adapter, LTX2_I2AV_Adapter])
+def test_neutral_guidance_preserves_raw_velocity_near_clean_endpoint(cls: type) -> None:
+ adapter = _adapter(cls)
+ video = _video_latents()
+ audio = _audio_latents()
+ extra = {"conditioning_mask": _conditioning_mask()} if cls is LTX2_I2AV_Adapter else {}
+
+ output = adapter.forward(
+ t=torch.full((BATCH_SIZE,), 0.06),
+ t_next=torch.zeros(BATCH_SIZE),
+ latents=torch.cat([video, audio], dim=1),
+ video_seq_len=VIDEO_SEQ_LEN,
+ compute_log_prob=False,
+ return_kwargs=["velocity"],
+ preserve_raw_model_velocity=True,
+ **extra,
+ **_forward_kwargs(),
+ )
+
+ expected_video = video * 0.5 + 1.0
+ expected_audio = audio * -0.25 + 2.0
+ assert torch.equal(
+ output.velocity,
+ torch.cat([expected_video, expected_audio], dim=1),
+ )
+ assert adapter.scheduler.steps == []
+ assert adapter.audio_scheduler.steps == []
diff --git a/tests/models/ltx2/test_ltx2_output_codec.py b/tests/models/ltx2/test_ltx2_output_codec.py
new file mode 100644
index 000000000..8fa0df2d4
--- /dev/null
+++ b/tests/models/ltx2/test_ltx2_output_codec.py
@@ -0,0 +1,526 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU tests for LTX2's config-driven offline audiovisual boundary."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import Any, Optional
+
+import numpy as np
+import pytest
+import torch
+import torchaudio
+
+from flow_factory.contracts import BatchCapability, GeometrySource, MediaType
+from flow_factory.models.ltx2._output import (
+ LTX2AVOutputCodec,
+ LTX2FirstFrameConditionPreparer,
+ decode_ltx2_output_state,
+ ltx2_log_mel_spectrogram,
+ resolve_ltx2_output_geometry,
+ validate_ltx2_encoded_output_geometry,
+)
+from flow_factory.models.ltx2.ltx2_i2av import LTX2_I2AV_Adapter
+from flow_factory.models.ltx2.ltx2_t2av import LTX2_T2AV_Adapter
+from flow_factory.models.output_state import validate_encoded_output_state
+
+BATCH_SIZE = 2
+HEIGHT = 8
+WIDTH = 8
+NUM_FRAMES = 3
+FRAME_RATE = 4.0
+VIDEO_LATENT_CHANNELS = 2
+LATENT_FRAMES = 2
+LATENT_HEIGHT = 4
+LATENT_WIDTH = 4
+VIDEO_SEQUENCE_LENGTH = LATENT_FRAMES * LATENT_HEIGHT * LATENT_WIDTH
+AUDIO_SAMPLE_RATE = 8000
+AUDIO_HOP_LENGTH = 80
+AUDIO_MEL_BINS = 8
+AUDIO_LATENT_CHANNELS = 2
+AUDIO_LATENT_MEL_BINS = 4
+AUDIO_FEATURE_DIM = AUDIO_LATENT_CHANNELS * AUDIO_LATENT_MEL_BINS
+AUDIO_TARGET_SAMPLES = 6000
+AUDIO_LATENT_FRAMES = 38
+
+
+@dataclass(frozen=True)
+class _DecodedMedia:
+ type: str
+ payload: Any
+ fps: Optional[float] = None
+ sample_rate: Optional[int] = None
+
+
+class _ModeOnlyPosterior:
+ """Posterior fake proving target and condition codecs never sample."""
+
+ def __init__(self, value: torch.Tensor) -> None:
+ self.value = value
+ self.mode_calls = 0
+ self.sample_calls = 0
+
+ def mode(self) -> torch.Tensor:
+ self.mode_calls += 1
+ return self.value
+
+ def sample(self, generator: Optional[torch.Generator] = None) -> torch.Tensor:
+ del generator
+ self.sample_calls += 1
+ raise AssertionError("LTX2 offline targets must use posterior mode")
+
+
+class _VideoVAEFake:
+ dtype = torch.float32
+
+ def __init__(self) -> None:
+ self.config = SimpleNamespace(latent_channels=VIDEO_LATENT_CHANNELS, scaling_factor=2.0)
+ self.latents_mean = torch.tensor([0.5, -0.5])
+ self.latents_std = torch.tensor([2.0, 4.0])
+ self.posteriors: list[_ModeOnlyPosterior] = []
+
+ def encode(self, pixels: torch.Tensor) -> Any:
+ if pixels.shape[2] == 1:
+ latents = pixels[:, :VIDEO_LATENT_CHANNELS, :, ::2, ::2]
+ else:
+ latents = pixels[:, :VIDEO_LATENT_CHANNELS, ::2, ::2, ::2]
+ posterior = _ModeOnlyPosterior(latents + 0.25)
+ self.posteriors.append(posterior)
+ return SimpleNamespace(latent_dist=posterior)
+
+
+class _AudioVAEFake:
+ dtype = torch.float32
+
+ def __init__(self) -> None:
+ self.config = SimpleNamespace(
+ sample_rate=AUDIO_SAMPLE_RATE,
+ mel_hop_length=AUDIO_HOP_LENGTH,
+ mel_bins=AUDIO_MEL_BINS,
+ in_channels=2,
+ latent_channels=AUDIO_LATENT_CHANNELS,
+ )
+ self.latents_mean = torch.linspace(-0.4, 0.3, AUDIO_FEATURE_DIM)
+ self.latents_std = torch.linspace(1.0, 1.7, AUDIO_FEATURE_DIM)
+ self.posteriors: list[_ModeOnlyPosterior] = []
+
+ def encode(self, log_mel: torch.Tensor) -> Any:
+ latents = log_mel[:, :AUDIO_LATENT_CHANNELS, ::2, ::2]
+ posterior = _ModeOnlyPosterior(latents)
+ self.posteriors.append(posterior)
+ return SimpleNamespace(latent_dist=posterior)
+
+
+class _VideoProcessorFake:
+ def preprocess_video(
+ self,
+ videos: list[np.ndarray],
+ *,
+ height: int,
+ width: int,
+ ) -> torch.Tensor:
+ assert (height, width) == (HEIGHT, WIDTH)
+ stacked = np.stack(videos)
+ return torch.from_numpy(stacked).permute(0, 4, 1, 2, 3).float().div(127.5).sub(1.0)
+
+
+class _PipelineFake:
+ vae_spatial_compression_ratio = 2
+ vae_temporal_compression_ratio = 2
+ transformer_spatial_patch_size = 1
+ transformer_temporal_patch_size = 1
+ audio_vae_temporal_compression_ratio = 2
+ audio_vae_mel_compression_ratio = 2
+ audio_sampling_rate = AUDIO_SAMPLE_RATE
+ audio_hop_length = AUDIO_HOP_LENGTH
+
+ def __init__(self) -> None:
+ self.video_processor = _VideoProcessorFake()
+
+ @staticmethod
+ def _normalize_latents(
+ latents: torch.Tensor,
+ latents_mean: torch.Tensor,
+ latents_std: torch.Tensor,
+ scaling_factor: float = 1.0,
+ ) -> torch.Tensor:
+ mean = latents_mean.view(1, -1, 1, 1, 1).to(latents)
+ std = latents_std.view(1, -1, 1, 1, 1).to(latents)
+ return (latents - mean) * scaling_factor / std
+
+ @staticmethod
+ def _pack_latents(
+ latents: torch.Tensor,
+ patch_size: int = 1,
+ patch_size_t: int = 1,
+ ) -> torch.Tensor:
+ batch, channels, frames, height, width = latents.shape
+ latents = latents.reshape(
+ batch,
+ -1,
+ frames // patch_size_t,
+ patch_size_t,
+ height // patch_size,
+ patch_size,
+ width // patch_size,
+ patch_size,
+ )
+ return latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
+
+ @staticmethod
+ def _unpack_latents(
+ latents: torch.Tensor,
+ num_frames: int,
+ height: int,
+ width: int,
+ patch_size: int = 1,
+ patch_size_t: int = 1,
+ ) -> torch.Tensor:
+ batch = latents.shape[0]
+ latents = latents.reshape(
+ batch,
+ num_frames,
+ height,
+ width,
+ -1,
+ patch_size_t,
+ patch_size,
+ patch_size,
+ )
+ return latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
+
+ @staticmethod
+ def _pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
+ return latents.transpose(1, 2).flatten(2, 3)
+
+ @staticmethod
+ def _normalize_audio_latents(
+ latents: torch.Tensor,
+ latents_mean: torch.Tensor,
+ latents_std: torch.Tensor,
+ ) -> torch.Tensor:
+ return (latents - latents_mean.to(latents)) / latents_std.to(latents)
+
+
+class _AdapterFake:
+ def __init__(self) -> None:
+ self.device = torch.device("cpu")
+ self.training_args = SimpleNamespace(
+ height=HEIGHT,
+ width=WIDTH,
+ num_frames=NUM_FRAMES,
+ frame_rate=FRAME_RATE,
+ )
+ self.pipeline = _PipelineFake()
+ self.components = {
+ "vae": _VideoVAEFake(),
+ "audio_vae": _AudioVAEFake(),
+ "transformer": SimpleNamespace(
+ config=SimpleNamespace(
+ in_channels=VIDEO_LATENT_CHANNELS,
+ audio_in_channels=AUDIO_FEATURE_DIM,
+ audio_patch_size=1,
+ audio_patch_size_t=1,
+ # LTX-2.3-only topology flags do not alter the AV target contract.
+ gated_attn=True,
+ cross_attn_mod=True,
+ )
+ ),
+ }
+ self.decode_calls: list[tuple[Any, ...]] = []
+
+ def get_component(self, name: str) -> Any:
+ return self.components[name]
+
+ def get_component_config(self, name: str) -> Any:
+ return self.components[name].config
+
+ def decode_latents(
+ self,
+ video_latents: torch.Tensor,
+ audio_latents: torch.Tensor,
+ **kwargs: Any,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ self.decode_calls.append((video_latents, audio_latents, kwargs))
+ return video_latents, audio_latents
+
+
+def _media_batch() -> tuple[tuple[_DecodedMedia, _DecodedMedia], ...]:
+ samples = []
+ time = torch.arange(AUDIO_TARGET_SAMPLES, dtype=torch.float32) / AUDIO_SAMPLE_RATE
+ for index in range(BATCH_SIZE):
+ frames = np.full(
+ (NUM_FRAMES, HEIGHT, WIDTH, 3),
+ fill_value=32 + index * 48,
+ dtype=np.uint8,
+ )
+ waveform = torch.sin(2 * torch.pi * (220 + index * 30) * time).unsqueeze(0)
+ samples.append(
+ (
+ _DecodedMedia(type="video", payload=frames, fps=FRAME_RATE),
+ _DecodedMedia(
+ type="audio",
+ payload=waveform,
+ sample_rate=AUDIO_SAMPLE_RATE,
+ ),
+ )
+ )
+ return tuple(samples)
+
+
+def _condition_images() -> torch.Tensor:
+ values = torch.linspace(-0.8, 0.8, BATCH_SIZE * 3 * HEIGHT * WIDTH)
+ return values.reshape(BATCH_SIZE, 3, HEIGHT, WIDTH)
+
+
+def test_official_lightricks_log_mel_frontend_matches_torchaudio() -> None:
+ waveforms = torch.linspace(-1.0, 1.0, 2 * 2048).reshape(1, 2, 2048)
+
+ actual = ltx2_log_mel_spectrogram(
+ waveforms,
+ sample_rate=AUDIO_SAMPLE_RATE,
+ hop_length=AUDIO_HOP_LENGTH,
+ mel_bins=AUDIO_MEL_BINS,
+ )
+ frontend = torchaudio.transforms.MelSpectrogram(
+ sample_rate=AUDIO_SAMPLE_RATE,
+ n_fft=1024,
+ win_length=1024,
+ hop_length=AUDIO_HOP_LENGTH,
+ f_min=0.0,
+ f_max=AUDIO_SAMPLE_RATE / 2,
+ n_mels=AUDIO_MEL_BINS,
+ window_fn=torch.hann_window,
+ center=True,
+ pad_mode="reflect",
+ power=1.0,
+ norm="slaney",
+ mel_scale="slaney",
+ )
+ expected = frontend(waveforms).clamp_min(1e-5).log().permute(0, 1, 3, 2)
+
+ assert torch.equal(actual, expected)
+
+
+def test_t2av_codec_encodes_config_driven_joint_mode_state() -> None:
+ adapter = _AdapterFake()
+ geometry = resolve_ltx2_output_geometry(adapter, conditioned=False)
+
+ with torch.no_grad():
+ encoded = LTX2AVOutputCodec(adapter).encode_output_state(_media_batch(), {})
+
+ assert geometry.video.sequence_length == VIDEO_SEQUENCE_LENGTH
+ assert geometry.audio.latent_frames == AUDIO_LATENT_FRAMES
+ assert geometry.audio.target_samples == AUDIO_TARGET_SAMPLES
+ assert encoded.clean_state.component_names == ("video", "audio")
+ assert encoded.clean_state.components["video"].shape == (
+ BATCH_SIZE,
+ VIDEO_SEQUENCE_LENGTH,
+ VIDEO_LATENT_CHANNELS,
+ )
+ assert encoded.clean_state.components["audio"].shape == (
+ BATCH_SIZE,
+ AUDIO_LATENT_FRAMES,
+ AUDIO_FEATURE_DIM,
+ )
+ assert encoded.clean_state.active_masks is None
+ assert encoded.forward_context["video_seq_len"] == VIDEO_SEQUENCE_LENGTH
+ assert encoded.forward_context["audio_num_frames"] == AUDIO_LATENT_FRAMES
+ assert encoded.geometry_signatures[0].media[0].type is MediaType.VIDEO
+ assert encoded.geometry_signatures[0].media[1].type is MediaType.AUDIO
+ assert encoded.geometry_signatures[0].media[1].samples == AUDIO_TARGET_SAMPLES
+ assert all(posterior.mode_calls == 1 for posterior in adapter.components["vae"].posteriors)
+ assert all(
+ posterior.mode_calls == 1 for posterior in adapter.components["audio_vae"].posteriors
+ )
+
+ validate_encoded_output_state(
+ encoded,
+ contract=LTX2_T2AV_Adapter.pipeline_io_contract,
+ expected_component_order=("video", "audio"),
+ expected_batch_size=BATCH_SIZE,
+ device="cpu",
+ )
+ validate_ltx2_encoded_output_geometry(
+ adapter,
+ _media_batch(),
+ {},
+ encoded,
+ conditioned=False,
+ )
+
+
+def test_i2av_preparer_binds_and_masks_the_exact_first_latent_frame() -> None:
+ adapter = _AdapterFake()
+ cached_condition = {
+ "connector_prompt_embeds": torch.zeros(BATCH_SIZE, 2, 3),
+ "condition_images": _condition_images(),
+ }
+
+ with torch.no_grad():
+ prepared = LTX2FirstFrameConditionPreparer(adapter).prepare_condition_state(
+ cached_condition
+ )
+ encoded = LTX2AVOutputCodec(adapter, conditioned=True).encode_output_state(
+ _media_batch(),
+ prepared.output_codec_condition(),
+ )
+
+ assert "condition_images" not in prepared.condition
+ assert prepared.forward_context == {}
+ condition_latents = prepared.output_context["condition_video_latents"]
+ assert condition_latents.shape == (
+ BATCH_SIZE,
+ VIDEO_LATENT_CHANNELS,
+ 1,
+ LATENT_HEIGHT,
+ LATENT_WIDTH,
+ )
+ conditioning_mask = encoded.forward_context["conditioning_mask"]
+ first_frame_tokens = LATENT_HEIGHT * LATENT_WIDTH
+ assert torch.equal(
+ conditioning_mask[:, :first_frame_tokens],
+ torch.ones(BATCH_SIZE, first_frame_tokens),
+ )
+ assert not bool(conditioning_mask[:, first_frame_tokens:].any())
+ assert torch.equal(
+ encoded.clean_state.active_masks["video"].squeeze(-1),
+ ~conditioning_mask.bool(),
+ )
+ assert bool(encoded.clean_state.active_masks["audio"].all())
+
+ unpacked = adapter.pipeline._unpack_latents(
+ encoded.clean_state.components["video"],
+ LATENT_FRAMES,
+ LATENT_HEIGHT,
+ LATENT_WIDTH,
+ )
+ expected_first_frame = adapter.pipeline._normalize_latents(
+ condition_latents,
+ adapter.components["vae"].latents_mean,
+ adapter.components["vae"].latents_std,
+ adapter.components["vae"].config.scaling_factor,
+ )
+ assert torch.equal(unpacked[:, :, :1], expected_first_frame)
+
+ validate_encoded_output_state(
+ encoded,
+ contract=LTX2_I2AV_Adapter.pipeline_io_contract,
+ expected_component_order=("video", "audio"),
+ expected_batch_size=BATCH_SIZE,
+ device="cpu",
+ )
+ validate_ltx2_encoded_output_geometry(
+ adapter,
+ _media_batch(),
+ prepared.output_codec_condition(),
+ encoded,
+ conditioned=True,
+ )
+
+
+def test_shared_decoder_routes_both_components_and_configured_geometry() -> None:
+ adapter = _AdapterFake()
+ with torch.no_grad():
+ encoded = LTX2AVOutputCodec(adapter).encode_output_state(_media_batch(), {})
+
+ decoded = decode_ltx2_output_state(adapter, encoded, output_type="pt")
+
+ assert decoded == (
+ encoded.clean_state.components["video"],
+ encoded.clean_state.components["audio"],
+ )
+ assert adapter.decode_calls[0][2] == {
+ "height": HEIGHT,
+ "width": WIDTH,
+ "num_frames": NUM_FRAMES,
+ "frame_rate": FRAME_RATE,
+ "output_type": "pt",
+ }
+
+
+def test_adapters_declare_complete_offline_av_capability() -> None:
+ LTX2_T2AV_Adapter.validate_offline_output_capability()
+ LTX2_I2AV_Adapter.validate_offline_output_capability()
+
+ assert LTX2_T2AV_Adapter.component_load_dtype_defaults == {"audio_vae": torch.float32}
+ assert LTX2_I2AV_Adapter.component_load_dtype_defaults == {"audio_vae": torch.float32}
+ assert LTX2_T2AV_Adapter.pipeline_io_contract.geometry_source is GeometrySource.CONFIGURED
+ assert LTX2_I2AV_Adapter.pipeline_io_contract.geometry_source is GeometrySource.CONFIGURED
+ assert LTX2_T2AV_Adapter.pipeline_io_contract.batch_capability is BatchCapability.UNIFORM
+ assert LTX2_I2AV_Adapter.pipeline_io_contract.batch_capability is BatchCapability.UNIFORM
+ assert LTX2_T2AV_Adapter.pipeline_io_contract.input_media.rules == ()
+ assert LTX2_I2AV_Adapter.pipeline_io_contract.input_media.rules[0].min_count == 1
+ assert LTX2_I2AV_Adapter.pipeline_io_contract.input_media.rules[0].max_count == 1
+ assert LTX2_T2AV_Adapter.output_state_codec_unavailable_reason is None
+ assert LTX2_I2AV_Adapter.output_state_codec_unavailable_reason is None
+ assert LTX2_T2AV_Adapter.offline_training_forward_overrides == {
+ "guidance_scale": 1.0,
+ "audio_guidance_scale": 1.0,
+ "guidance_rescale": 0.0,
+ "audio_guidance_rescale": 0.0,
+ "stg_scale": 0.0,
+ "audio_stg_scale": 0.0,
+ "spatio_temporal_guidance_blocks": None,
+ "modality_scale": 1.0,
+ "audio_modality_scale": 1.0,
+ "preserve_raw_model_velocity": True,
+ }
+
+ with pytest.raises(TypeError):
+ LTX2_T2AV_Adapter.offline_training_forward_overrides["guidance_scale"] = 2.0
+
+
+@pytest.mark.parametrize(
+ ("adapter_class", "conditioned"),
+ [(LTX2_T2AV_Adapter, False), (LTX2_I2AV_Adapter, True)],
+)
+def test_offline_flow_objective_sums_modalities_without_changing_joint_reducer(
+ adapter_class: type,
+ conditioned: bool,
+) -> None:
+ codec_adapter = _AdapterFake()
+ condition: Any = {}
+ if conditioned:
+ condition = LTX2FirstFrameConditionPreparer(codec_adapter).prepare_condition_state(
+ {"condition_images": _condition_images()}
+ )
+ condition = condition.output_codec_condition()
+ with torch.no_grad():
+ encoded = LTX2AVOutputCodec(codec_adapter, conditioned=conditioned).encode_output_state(
+ _media_batch(),
+ condition,
+ )
+ values = {
+ "video": torch.full_like(encoded.clean_state.components["video"], 2.0),
+ "audio": torch.full_like(encoded.clean_state.components["audio"], 5.0),
+ }
+ if conditioned:
+ values["video"] = values["video"].masked_fill(
+ ~encoded.clean_state.active_masks["video"],
+ 1000.0,
+ )
+ adapter = object.__new__(adapter_class)
+
+ offline = adapter.reduce_flow_matching_objective_values(
+ values,
+ state=encoded.clean_state,
+ )
+ joint = adapter.reduce_latent_values(values, state=encoded.clean_state)
+
+ assert torch.equal(offline, torch.full((BATCH_SIZE,), 7.0))
+ assert not torch.equal(joint, offline)
diff --git a/tests/models/minimax_h3/test_condition_state.py b/tests/models/minimax_h3/test_condition_state.py
new file mode 100644
index 000000000..57d708283
--- /dev/null
+++ b/tests/models/minimax_h3/test_condition_state.py
@@ -0,0 +1,116 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for one-shot MiniMax H3 offline condition realization."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from flow_factory.models.minimax_h3._condition import MiniMaxH3ConditionStatePreparer
+from flow_factory.models.minimax_h3.adapters import (
+ MiniMaxH3FL2VAAdapter,
+ MiniMaxH3Ref2VAAdapter,
+ MiniMaxH3T2VAAdapter,
+)
+
+
+def _cached_condition() -> dict[str, object]:
+ return {
+ "prompt_embeds": torch.zeros(1, 2, 4),
+ "position_ids": torch.zeros(1, 9, 3, dtype=torch.float64),
+ "token_tags": torch.zeros(1, 9, dtype=torch.long),
+ "text_indices": torch.tensor([[0, 1]], dtype=torch.long),
+ "audio_indices": torch.tensor([[2, 3, 4]], dtype=torch.long),
+ "video_indices": torch.tensor([[5, 6, 7, 8]], dtype=torch.long),
+ "num_condition_video_rows": torch.tensor([2]),
+ "num_condition_audio_rows": torch.tensor([1]),
+ "condition_latents": [[torch.ones(1, 24, 1, 2, 2)]],
+ "audio_condition_latents": [[torch.ones(1, 32)]],
+ "height": torch.tensor([32]),
+ "width": torch.tensor([32]),
+ "num_frames": torch.tensor([22]),
+ "reference_manifest": "ordered-manifest",
+ }
+
+
+def test_conditioned_adapters_declare_preparer_but_t2va_keeps_identity() -> None:
+ t2va = object.__new__(MiniMaxH3T2VAAdapter)
+ fl2va = object.__new__(MiniMaxH3FL2VAAdapter)
+ ref2va = object.__new__(MiniMaxH3Ref2VAAdapter)
+
+ assert t2va.build_condition_state_preparer() is None
+ assert isinstance(fl2va.build_condition_state_preparer(), MiniMaxH3ConditionStatePreparer)
+ assert isinstance(ref2va.build_condition_state_preparer(), MiniMaxH3ConditionStatePreparer)
+ assert MiniMaxH3ConditionStatePreparer.required_components == ("scheduler",)
+ assert {
+ MiniMaxH3T2VAAdapter.preprocess_cache_version,
+ MiniMaxH3FL2VAAdapter.preprocess_cache_version,
+ MiniMaxH3Ref2VAAdapter.preprocess_cache_version,
+ } == {"minimax-h3-v2"}
+
+
+def test_ref2va_preparer_realizes_prefix_once_and_routes_owned_contexts(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter = SimpleNamespace(workflow="ref2va", pipeline=object())
+ preparer = MiniMaxH3ConditionStatePreparer(adapter)
+ condition = _cached_condition()
+ prefixes = {
+ "video": torch.randn(1, 2, 96),
+ "audio": torch.randn(1, 1, 32),
+ }
+ calls: list[object] = []
+
+ def prepare(pipeline, cached, *, workflow, generator):
+ calls.append((pipeline, cached, workflow, generator))
+ return prefixes
+
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3._condition.prepare_h3_condition_prefixes",
+ prepare,
+ )
+ generator = torch.Generator().manual_seed(11)
+
+ prepared = preparer.prepare_condition_state(condition, generator)
+
+ assert calls == [(adapter.pipeline, condition, "ref2va", generator)]
+ assert "condition_latents" not in prepared.condition
+ assert "audio_condition_latents" not in prepared.condition
+ assert "position_ids" not in prepared.condition
+ assert "num_condition_video_rows" not in prepared.condition
+ assert prepared.condition["reference_manifest"] == "ordered-manifest"
+ forward = prepared.model_forward_condition()
+ codec = prepared.output_codec_condition()
+ assert forward["condition_prefixes"] is prefixes
+ assert forward["layout"]["video_indices"].shape == (4,)
+ assert forward["layout"]["num_condition_video_rows"] == 2
+ assert codec["layout"] is forward["layout"]
+ assert "condition_prefixes" not in codec
+ assert prepared.model_forward_condition()["condition_prefixes"]["video"] is prefixes["video"]
+
+
+def test_preparer_rejects_wrong_workflow_and_stale_runtime_prefix() -> None:
+ with pytest.raises(ValueError, match="requires workflow"):
+ MiniMaxH3ConditionStatePreparer(
+ SimpleNamespace(workflow="t2va", pipeline=object())
+ ).prepare_condition_state(_cached_condition())
+
+ condition = _cached_condition()
+ condition["condition_prefixes"] = {"video": torch.empty(1, 0, 96)}
+ with pytest.raises(ValueError, match="must not contain already-realized"):
+ MiniMaxH3ConditionStatePreparer(
+ SimpleNamespace(workflow="fl2va", pipeline=object())
+ ).prepare_condition_state(condition)
diff --git a/tests/models/minimax_h3/test_diffusers_api_contract.py b/tests/models/minimax_h3/test_diffusers_api_contract.py
index d8c8e37ca..9aeca1a8d 100644
--- a/tests/models/minimax_h3/test_diffusers_api_contract.py
+++ b/tests/models/minimax_h3/test_diffusers_api_contract.py
@@ -22,6 +22,7 @@
from PIL import Image
from flow_factory.models.minimax_h3 import dependency
+from flow_factory.models.minimax_h3.blocks import prepare_h3_rollout_state
from flow_factory.models.runtime import ModularPipelineRuntime
EXECUTED_NO_WEIGHT_BLOCKS = frozenset(
@@ -68,9 +69,9 @@ def test_h3_capable_environment_uses_supported_diffusers_release(
del h3_symbols
installed = metadata.version("diffusers")
required = dependency.MINIMAX_H3_DIFFUSERS_MIN_VERSION
- assert Version(installed) >= Version(required), (
- f"MiniMax H3 requires diffusers>={required}, received {installed}"
- )
+ assert Version(installed) >= Version(
+ required
+ ), f"MiniMax H3 requires diffusers>={required}, received {installed}"
def test_real_h3_symbols_preserve_workflows_and_callable_surfaces(h3_symbols) -> None:
@@ -147,9 +148,7 @@ def test_real_h3_blocks_execute_no_weight_pipeline_state_transitions(h3_symbols)
"num_frames": 124,
}
)
- returned_pipeline, t2va_state = h3_symbols.NoKeyframeAnchorsStep()(
- t2va_pipeline, t2va_state
- )
+ returned_pipeline, t2va_state = h3_symbols.NoKeyframeAnchorsStep()(t2va_pipeline, t2va_state)
assert returned_pipeline is t2va_pipeline
assert t2va_state.values["keyframe_anchors"] == ()
executed.add("NoKeyframeAnchorsStep")
@@ -178,9 +177,7 @@ def test_real_h3_blocks_execute_no_weight_pipeline_state_transitions(h3_symbols)
"condition_latents": [torch.zeros(1, 24, 1, 2, 2)],
}
)
- _, condition_state = h3_symbols.PrepareConditionLatentsStep()(
- t2va_pipeline, condition_state
- )
+ _, condition_state = h3_symbols.PrepareConditionLatentsStep()(t2va_pipeline, condition_state)
assert condition_state.values["condition_rows"].shape == (1, 96)
executed.add("PrepareConditionLatentsStep")
@@ -209,6 +206,18 @@ def test_real_h3_blocks_execute_no_weight_pipeline_state_transitions(h3_symbols)
_, resize_state = h3_symbols.ResizeStep()(fl2va_pipeline, resize_state)
assert resize_state.values["keyframe_anchors"] == ("first",)
assert resize_state.values["keyframes"][0].size == (32, 32)
+
+ last_only_state = h3_symbols.PipelineState(
+ values={
+ "image": None,
+ "last_image": Image.new("RGB", (16, 16)),
+ "height": 32,
+ "width": 32,
+ }
+ )
+ _, last_only_state = h3_symbols.ResizeStep()(fl2va_pipeline, last_only_state)
+ assert last_only_state.values["keyframe_anchors"] == ("last",)
+ assert last_only_state.values["keyframes"][0].size == (32, 32)
executed.add("ResizeStep")
fl2va_state = h3_symbols.PipelineState(
@@ -269,9 +278,7 @@ def test_real_h3_blocks_execute_no_weight_pipeline_state_transitions(h3_symbols)
"audio_latents": torch.zeros(2, 32),
}
)
- _, ref_latents_state = h3_symbols.Ref2VAPrepareLatentsStep()(
- ref2va_pipeline, ref_latents_state
- )
+ _, ref_latents_state = h3_symbols.Ref2VAPrepareLatentsStep()(ref2va_pipeline, ref_latents_state)
assert ref_latents_state.values["latents"].shape == (3, 96)
assert ref_latents_state.values["audio_latents"].shape == (3, 32)
executed.add("Ref2VAPrepareLatentsStep")
@@ -282,6 +289,129 @@ def test_real_h3_blocks_execute_no_weight_pipeline_state_transitions(h3_symbols)
)
+def test_shared_rollout_prefix_helper_matches_official_fl2va_chain_and_rng(
+ h3_symbols,
+) -> None:
+ pipeline = h3_symbols.MiniMaxH3ModularPipeline.from_config({"workflow": "fl2va"})
+ scheduler_class = pipeline.get_component_spec("scheduler").type_hint
+ pipeline.register_components(
+ scheduler=scheduler_class(shift=12.0),
+ audio_scheduler=scheduler_class(shift=3.0),
+ )
+ condition = torch.arange(96, dtype=torch.float32).reshape(1, 24, 1, 2, 2) / 100
+ cached = {
+ "num_latent_frames": 1,
+ "latent_height": 2,
+ "latent_width": 2,
+ "num_audio_latents": 1,
+ "num_condition_video_rows": 1,
+ "num_condition_audio_rows": 0,
+ "condition_latents": [condition],
+ "video_indices": torch.tensor([0, 1]),
+ "audio_indices": torch.tensor([2, 3]),
+ }
+ framework_generator = torch.Generator().manual_seed(123)
+ oracle_generator = torch.Generator().manual_seed(123)
+
+ targets, prefixes = prepare_h3_rollout_state(
+ pipeline,
+ cached,
+ workflow="fl2va",
+ generator=framework_generator,
+ )
+
+ oracle_state = h3_symbols.PipelineState(
+ values={
+ **cached,
+ "generator": oracle_generator,
+ "latents": None,
+ "audio_latents": None,
+ }
+ )
+ for block_type in (
+ h3_symbols.PrepareConditionLatentsStep,
+ h3_symbols.PrepareLatentsStep,
+ h3_symbols.FL2VAPrepareLatentsStep,
+ ):
+ _, oracle_state = block_type()(pipeline, oracle_state)
+
+ torch.testing.assert_close(prefixes["video"][0], oracle_state.values["latents"][:1])
+ torch.testing.assert_close(targets.components["video"][0], oracle_state.values["latents"][1:])
+ torch.testing.assert_close(
+ targets.components["audio"][0],
+ oracle_state.values["audio_latents"],
+ )
+ assert prefixes["audio"].shape == (1, 0, 32)
+ torch.testing.assert_close(
+ torch.rand((), generator=framework_generator),
+ torch.rand((), generator=oracle_generator),
+ )
+
+
+def test_shared_rollout_prefix_helper_matches_official_ref2va_audio_order(
+ h3_symbols,
+) -> None:
+ pipeline = h3_symbols.MiniMaxH3ModularPipeline.from_config({"workflow": "ref2va"})
+ scheduler_class = pipeline.get_component_spec("scheduler").type_hint
+ pipeline.register_components(
+ scheduler=scheduler_class(shift=12.0),
+ audio_scheduler=scheduler_class(shift=3.0),
+ )
+ audio_conditions = [
+ torch.full((1, 32), 4.0),
+ torch.stack([torch.full((32,), 5.0), torch.full((32,), 6.0)]),
+ ]
+ cached = {
+ "num_latent_frames": 1,
+ "latent_height": 2,
+ "latent_width": 2,
+ "num_audio_latents": 1,
+ "num_condition_video_rows": 1,
+ "num_condition_audio_rows": 3,
+ "condition_latents": [torch.ones(1, 24, 1, 2, 2)],
+ "audio_condition_latents": audio_conditions,
+ "video_indices": torch.tensor([0, 1]),
+ "audio_indices": torch.tensor([2, 3, 4, 5, 6]),
+ }
+ framework_generator = torch.Generator().manual_seed(321)
+ oracle_generator = torch.Generator().manual_seed(321)
+
+ targets, prefixes = prepare_h3_rollout_state(
+ pipeline,
+ cached,
+ workflow="ref2va",
+ generator=framework_generator,
+ )
+
+ oracle_state = h3_symbols.PipelineState(
+ values={
+ **cached,
+ "generator": oracle_generator,
+ "latents": None,
+ "audio_latents": None,
+ }
+ )
+ for block_type in (
+ h3_symbols.PrepareConditionLatentsStep,
+ h3_symbols.PrepareLatentsStep,
+ h3_symbols.Ref2VAPrepareLatentsStep,
+ ):
+ _, oracle_state = block_type()(pipeline, oracle_state)
+
+ torch.testing.assert_close(prefixes["video"][0], oracle_state.values["latents"][:1])
+ torch.testing.assert_close(prefixes["audio"][0], oracle_state.values["audio_latents"][:3])
+ torch.testing.assert_close(targets.components["video"][0], oracle_state.values["latents"][1:])
+ torch.testing.assert_close(
+ targets.components["audio"][0],
+ oracle_state.values["audio_latents"][3:],
+ )
+ torch.testing.assert_close(prefixes["audio"][0, :, 0], torch.tensor([4.0, 5.0, 6.0]))
+ torch.testing.assert_close(
+ torch.rand((), generator=framework_generator),
+ torch.rand((), generator=oracle_generator),
+ )
+
+
@pytest.mark.parametrize(
("workflow", "target_component", "absent_target"),
[
diff --git a/tests/models/minimax_h3/test_modular_core.py b/tests/models/minimax_h3/test_modular_core.py
index 70fa1bf98..160b02397 100644
--- a/tests/models/minimax_h3/test_modular_core.py
+++ b/tests/models/minimax_h3/test_modular_core.py
@@ -18,6 +18,7 @@
import pytest
import torch
+from flow_factory.models.minimax_h3 import blocks
from flow_factory.models.minimax_h3._common import MINIMAX_H3_COMPONENT_ORDER
from flow_factory.samples import ComponentTimes, LatentState
from flow_factory.scheduler import SDESchedulerOutput
@@ -257,8 +258,6 @@ def _fake_pipeline():
def test_block_executor_propagates_shared_state_and_selects_outputs(monkeypatch):
- from flow_factory.models.minimax_h3 import blocks
-
monkeypatch.setattr(blocks, "require_minimax_h3_support", lambda: FakeSymbols())
pipeline = SimpleNamespace(calls=[])
first = RecordingBlock("first", {"tensor": torch.tensor([1.0])})
@@ -274,8 +273,6 @@ def test_block_executor_propagates_shared_state_and_selects_outputs(monkeypatch)
def test_block_executor_reports_missing_output_and_preserves_exception(monkeypatch):
- from flow_factory.models.minimax_h3 import blocks
-
monkeypatch.setattr(blocks, "require_minimax_h3_support", lambda: FakeSymbols())
with pytest.raises(ValueError, match="workflow='t2va'.*field='missing'"):
blocks.run_h3_blocks(
@@ -302,7 +299,7 @@ def test_block_executor_reports_missing_output_and_preserves_exception(monkeypat
(
"t2va",
["TextEncoderStep", "NoKeyframeAnchorsStep", "PrepareLayoutStep"],
- {"prompt_embeds", "keyframe_anchors", "video_indices", "audio_indices"},
+ {"prompt_embeds", "video_indices", "audio_indices"},
),
(
"fl2va",
@@ -312,7 +309,7 @@ def test_block_executor_reports_missing_output_and_preserves_exception(monkeypat
"KeyframeEncoderStep",
"PrepareLayoutStep",
],
- {"prompt_embeds", "keyframes", "condition_latents", "video_indices"},
+ {"prompt_embeds", "condition_latents", "video_indices"},
),
(
"ref2va",
@@ -324,7 +321,6 @@ def test_block_executor_reports_missing_output_and_preserves_exception(monkeypat
],
{
"prompt_embeds",
- "normalized_references",
"condition_latents",
"audio_condition_latents",
},
@@ -334,8 +330,6 @@ def test_block_executor_reports_missing_output_and_preserves_exception(monkeypat
def test_encoding_workflows_use_pinned_block_order_and_cache_contract(
monkeypatch, workflow, expected_order, required_fields
):
- from flow_factory.models.minimax_h3 import blocks
-
monkeypatch.setattr(blocks, "require_minimax_h3_support", lambda: FakeSymbols())
pipeline = _fake_pipeline()
cache = blocks.encode_h3_workflow_inputs(pipeline, {"prompt": "test"}, workflow=workflow)
@@ -345,6 +339,12 @@ def test_encoding_workflows_use_pinned_block_order_and_cache_contract(
assert set(cache) == set(blocks.ENCODE_WORKFLOW_FIELDS[workflow])
assert "generator" not in cache
assert "text_encoder" not in cache
+ assert not {
+ "keyframes",
+ "keyframe_anchors",
+ "normalized_references",
+ "text_token_tags",
+ }.intersection(cache)
@pytest.mark.parametrize(
@@ -376,8 +376,6 @@ def test_encoding_workflows_use_pinned_block_order_and_cache_contract(
def test_rollout_uses_condition_video_audio_rng_order_and_exact_split(
monkeypatch, workflow, video_conditions, audio_conditions, expected_order
):
- from flow_factory.models.minimax_h3 import blocks
-
monkeypatch.setattr(blocks, "require_minimax_h3_support", lambda: FakeSymbols())
pipeline = _fake_pipeline()
cache = _fake_layout_cache(video_conditions, audio_conditions)
@@ -399,9 +397,34 @@ def test_rollout_uses_condition_video_audio_rng_order_and_exact_split(
assert prefixes["audio"].shape[1] == audio_conditions
-def test_rollout_rejects_full_rows_mismatching_layout_before_split(monkeypatch):
- from flow_factory.models.minimax_h3 import blocks
+def test_condition_prefix_helper_accepts_collated_b1_tensor_containers(monkeypatch):
+ monkeypatch.setattr(blocks, "require_minimax_h3_support", lambda: FakeSymbols())
+ pipeline = _fake_pipeline()
+ cache = _fake_layout_cache(video_conditions=1, audio_conditions=2)
+ cache["video_indices"] = cache["video_indices"].unsqueeze(0)
+ cache["audio_indices"] = cache["audio_indices"].unsqueeze(0)
+ cache["num_condition_video_rows"] = torch.tensor([1])
+ cache["num_condition_audio_rows"] = [2]
+ cache["condition_latents"] = torch.ones(1, 1, 1, 24, 1, 2, 2)
+ cache["audio_condition_latents"] = torch.stack(
+ [torch.full((1, 32), 4.0), torch.full((1, 32), 5.0)]
+ ).unsqueeze(0)
+
+ prefixes = blocks.prepare_h3_condition_prefixes(
+ pipeline,
+ cache,
+ workflow="ref2va",
+ generator=torch.Generator().manual_seed(7),
+ )
+
+ assert pipeline.calls == ["PrepareConditionLatentsStep"]
+ assert [name for name, _ in pipeline.draws] == ["condition"]
+ assert prefixes["video"].shape == (1, 1, 96)
+ assert prefixes["audio"].shape == (1, 2, 32)
+ torch.testing.assert_close(prefixes["audio"][0, :, 0], torch.tensor([4.0, 5.0]))
+
+def test_rollout_rejects_full_rows_mismatching_layout_before_split(monkeypatch):
monkeypatch.setattr(blocks, "require_minimax_h3_support", lambda: FakeSymbols())
cache = _fake_layout_cache(video_conditions=1)
cache["video_indices"] = torch.arange(4)
@@ -423,8 +446,6 @@ def test_rollout_rejects_full_rows_mismatching_layout_before_split(monkeypatch):
def test_rollout_rejects_zero_target_rows_with_exact_diagnostic(monkeypatch):
- from flow_factory.models.minimax_h3 import blocks
-
class EmptyTargetPrepareLatentsStep(FakeWorkflowBlock):
def __call__(self, pipeline, state):
pipeline.calls.append(type(self).__name__)
diff --git a/tests/models/minimax_h3/test_output_codec.py b/tests/models/minimax_h3/test_output_codec.py
index fceacbff9..cd0f92255 100644
--- a/tests/models/minimax_h3/test_output_codec.py
+++ b/tests/models/minimax_h3/test_output_codec.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Tests for MiniMax H3 T2VA offline audiovisual target encoding."""
+"""Tests for MiniMax H3 offline audiovisual target encoding."""
from types import SimpleNamespace
from typing import Any, Optional
@@ -24,11 +24,15 @@
from flow_factory.contracts import (
BatchCapability,
GeometrySource,
+ InputMediaBinding,
+ InputMediaOrder,
MediaType,
NegativePromptPolicy,
RateRequirement,
+ validate_pipeline_model_input,
)
from flow_factory.data_utils.offline_dataset import DecodedMedia
+from flow_factory.data_utils.schema import MediaAsset, NormalizedModelInput
from flow_factory.models.minimax_h3._output import (
MiniMaxH3AVOutputCodec,
prepare_h3_target_audio,
@@ -151,9 +155,14 @@ def get_component(self, name: str) -> Any:
return getattr(self, name)
-def _base_adapter(*, latent_storage_dtype: str) -> MiniMaxH3T2VAAdapter:
+def _base_adapter(
+ *,
+ latent_storage_dtype: str,
+ adapter_type: type = MiniMaxH3T2VAAdapter,
+) -> Any:
components = _Adapter()
- adapter = object.__new__(MiniMaxH3T2VAAdapter)
+ components.transformer_ref = components.transformer
+ adapter = object.__new__(adapter_type)
adapter.accelerator = SimpleNamespace(device=torch.device("cpu"))
adapter.training_args = SimpleNamespace(
height=32,
@@ -180,14 +189,20 @@ def _base_adapter(*, latent_storage_dtype: str) -> MiniMaxH3T2VAAdapter:
)
adapter.scheduler.set_timesteps(2, device="cpu")
adapter.audio_scheduler.set_timesteps(2, device="cpu")
+ adapter._effective_pipeline_io_contract = adapter.pipeline_io_contract
+ adapter._condition_state_preparer = adapter.build_condition_state_preparer()
adapter._output_state_codec = adapter.build_output_state_codec()
return adapter
-def _condition() -> dict[str, Any]:
+def _condition(
+ *,
+ video_condition_rows: int = 0,
+ audio_condition_rows: int = 0,
+) -> dict[str, Any]:
text_rows = 2
- audio_rows = 74
- video_rows = 7
+ audio_rows = 74 + audio_condition_rows
+ video_rows = 7 + video_condition_rows
sequence_length = text_rows + audio_rows + video_rows
return {
"height": [32],
@@ -202,8 +217,8 @@ def _condition() -> dict[str, Any]:
"text_indices": torch.arange(text_rows).unsqueeze(0),
"audio_indices": torch.arange(text_rows, text_rows + audio_rows).unsqueeze(0),
"video_indices": torch.arange(text_rows + audio_rows, sequence_length).unsqueeze(0),
- "num_condition_video_rows": [0],
- "num_condition_audio_rows": [0],
+ "num_condition_video_rows": [video_condition_rows],
+ "num_condition_audio_rows": [audio_condition_rows],
}
@@ -243,15 +258,77 @@ def test_t2va_declares_exact_configured_single_sample_av_contract() -> None:
MiniMaxH3T2VAAdapter.validate_offline_output_capability()
-@pytest.mark.parametrize("adapter_type", [MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter])
-def test_conditioned_h3_workflows_require_reproducible_prefix_binder(
- adapter_type: type,
-) -> None:
- with pytest.raises(NotImplementedError, match="conditioned-prefix binder"):
- adapter_type.validate_offline_output_capability()
+def test_fl2va_declares_ordered_first_last_input_and_complete_offline_codec() -> None:
+ contract = MiniMaxH3FL2VAAdapter.pipeline_io_contract
+
+ assert contract is not None
+ assert contract.input_media.binding is InputMediaBinding.GROUPED_BY_TYPE
+ assert contract.input_media.order is InputMediaOrder.WITHIN_TYPE
+ assert len(contract.input_media.rules) == 1
+ rule = contract.input_media.rules[0]
+ assert rule.format.type is MediaType.IMAGE
+ assert (rule.min_count, rule.max_count) == (1, 2)
+ assert rule.slots == ("first_frame", "last_frame")
+ assert rule.required_slots == ()
+ MiniMaxH3FL2VAAdapter.validate_offline_output_capability()
-def test_codec_uses_framework_video_sample_without_condition_fp16_rounding() -> None:
+def test_ref2va_declares_global_ordered_multimodal_input_and_complete_codec() -> None:
+ contract = MiniMaxH3Ref2VAAdapter.pipeline_io_contract
+
+ assert contract is not None
+ assert contract.input_media.binding is InputMediaBinding.ORDERED_REFERENCES
+ assert contract.input_media.order is InputMediaOrder.GLOBAL
+ rules = contract.input_media.rules
+ assert tuple(rule.format.type for rule in rules) == (
+ MediaType.IMAGE,
+ MediaType.VIDEO,
+ MediaType.AUDIO,
+ )
+ assert tuple((rule.min_count, rule.max_count) for rule in rules) == (
+ (0, 9),
+ (0, 3),
+ (0, 3),
+ )
+ assert rules[1].format.fps is RateRequirement.OPTIONAL
+ assert rules[2].format.sample_rate is RateRequirement.OPTIONAL
+ assert contract.input_media.min_total_count == 1
+ assert contract.input_media.max_total_count == 12
+ assert contract.input_media.required_any_types == (MediaType.IMAGE, MediaType.VIDEO)
+ MiniMaxH3Ref2VAAdapter.validate_offline_output_capability()
+
+
+def test_ref2va_contract_fails_before_decode_for_empty_audio_only_and_overall_limit() -> None:
+ contract = MiniMaxH3Ref2VAAdapter.pipeline_io_contract
+
+ def model_input(*media: MediaAsset) -> NormalizedModelInput:
+ return NormalizedModelInput(
+ prompt="describe",
+ negative_prompt=None,
+ media=media,
+ )
+
+ with pytest.raises(ValueError, match="at least 1 input media item"):
+ validate_pipeline_model_input(model_input(), contract)
+ with pytest.raises(ValueError, match="whose type is in.*image.*video"):
+ validate_pipeline_model_input(
+ model_input(MediaAsset(type="audio", path="voice.wav")),
+ contract,
+ )
+ with pytest.raises(ValueError, match="at most 12 input media item"):
+ validate_pipeline_model_input(
+ model_input(
+ *(
+ [MediaAsset(type="image", path=f"image-{index}.png") for index in range(9)]
+ + [MediaAsset(type="video", path=f"video-{index}.mp4") for index in range(3)]
+ + [MediaAsset(type="audio", path="voice.wav")]
+ )
+ ),
+ contract,
+ )
+
+
+def test_codec_uses_deterministic_av_modes_without_condition_fp16_rounding() -> None:
adapter = _Adapter()
generator = torch.Generator().manual_seed(17)
@@ -261,9 +338,9 @@ def test_codec_uses_framework_video_sample_without_condition_fp16_rounding() ->
generator,
)
- assert adapter.video_posterior.sample_calls == 1
- assert adapter.video_posterior.sample_generator is generator
- assert adapter.video_posterior.mode_calls == 0
+ assert adapter.video_posterior.sample_calls == 0
+ assert adapter.video_posterior.sample_generator is None
+ assert adapter.video_posterior.mode_calls == 1
assert adapter.audio_posterior.mode_calls == 1
assert adapter.audio_posterior.sample_calls == 0
assert adapter.vae.encoded_pixels is not None
@@ -273,10 +350,10 @@ def test_codec_uses_framework_video_sample_without_condition_fp16_rounding() ->
assert encoded.clean_state.component_names == ("video", "audio")
assert encoded.clean_state.components["video"].shape == (1, 7, 96)
assert encoded.clean_state.components["audio"].shape == (1, 74, 32)
- sampled_value = adapter.video_posterior.values.flatten()[0]
- condition_rounded_value = sampled_value.to(torch.float16).to(torch.float32)
- assert sampled_value != condition_rounded_value
- assert encoded.clean_state.components["video"].flatten()[0] == sampled_value
+ mode_value = adapter.video_posterior.values.flatten()[0]
+ condition_rounded_value = mode_value.to(torch.float16).to(torch.float32)
+ assert mode_value != condition_rounded_value
+ assert encoded.clean_state.components["video"].flatten()[0] == mode_value
assert encoded.forward_context == {}
assert encoded.decode_context["geometry"] == {
"height": 32,
@@ -327,10 +404,31 @@ def test_geometry_hook_rejects_invalid_input_owned_flat_layout() -> None:
encoded = MiniMaxH3AVOutputCodec(adapter).encode_output_state(_media(), condition)
condition["audio_indices"] = condition["audio_indices"][:, :-1]
- with pytest.raises(ValueError, match="audio layout expected 74 target rows"):
+ with pytest.raises(ValueError, match=r"audio layout expected 0 condition \+ 74 target rows"):
validate_h3_encoded_output_geometry(adapter, _media(), condition, encoded)
+@pytest.mark.parametrize(
+ ("video_condition_rows", "audio_condition_rows"),
+ [(1, 0), (3, 5)],
+)
+def test_codec_accepts_conditioned_layout_while_encoding_target_only_rows(
+ video_condition_rows: int,
+ audio_condition_rows: int,
+) -> None:
+ adapter = _Adapter()
+ condition = _condition(
+ video_condition_rows=video_condition_rows,
+ audio_condition_rows=audio_condition_rows,
+ )
+
+ encoded = MiniMaxH3AVOutputCodec(adapter).encode_output_state(_media(), condition)
+
+ assert encoded.clean_state.components["video"].shape == (1, 7, 96)
+ assert encoded.clean_state.components["audio"].shape == (1, 74, 32)
+ validate_h3_encoded_output_geometry(adapter, _media(), condition, encoded)
+
+
def test_base_lifecycle_casts_clean_state_without_output_context_dtype_drift() -> None:
adapter = _base_adapter(latent_storage_dtype="fp16")
@@ -341,11 +439,30 @@ def test_base_lifecycle_casts_clean_state_without_output_context_dtype_drift() -
assert encoded.forward_context == {}
+def test_offline_flow_objective_sums_modality_means_without_changing_joint_reducer() -> None:
+ adapter = object.__new__(MiniMaxH3T2VAAdapter)
+ values = {
+ "video": torch.full((2, 7, 96), 2.0),
+ "audio": torch.full((2, 74, 32), 5.0),
+ }
+
+ offline = adapter.reduce_flow_matching_objective_values(values)
+ joint = adapter.reduce_latent_values(values)
+
+ assert torch.equal(offline, torch.full((2,), 7.0))
+ assert not torch.equal(joint, offline)
+
+
+@pytest.mark.parametrize(
+ "adapter_type",
+ [MiniMaxH3T2VAAdapter, MiniMaxH3FL2VAAdapter, MiniMaxH3Ref2VAAdapter],
+)
def test_adapter_decode_routes_both_components_with_cached_geometry(
monkeypatch: pytest.MonkeyPatch,
+ adapter_type: type,
) -> None:
encoded = MiniMaxH3AVOutputCodec(_Adapter()).encode_output_state(_media(), _condition())
- adapter = object.__new__(MiniMaxH3T2VAAdapter)
+ adapter = object.__new__(adapter_type)
observed: dict[str, Any] = {}
def decode(self: Any, state: LatentState, **kwargs: Any) -> str:
@@ -353,7 +470,7 @@ def decode(self: Any, state: LatentState, **kwargs: Any) -> str:
observed.update(kwargs)
return "decoded-av"
- monkeypatch.setattr(MiniMaxH3T2VAAdapter, "decode_latents", decode)
+ monkeypatch.setattr(adapter_type, "decode_latents", decode)
assert adapter.decode_output_state(encoded, output_type="np") == "decoded-av"
assert observed["state"] is encoded.clean_state
@@ -498,6 +615,73 @@ def test_h3_offline_dpo_arms_reuse_structured_noise_through_real_forward() -> No
)
+@pytest.mark.parametrize(
+ ("adapter_type", "audio_condition_rows"),
+ [(MiniMaxH3FL2VAAdapter, 0), (MiniMaxH3Ref2VAAdapter, 2)],
+)
+def test_conditioned_h3_prepares_one_prefix_for_codec_and_velocity_forward(
+ monkeypatch: pytest.MonkeyPatch,
+ adapter_type: type,
+ audio_condition_rows: int,
+) -> None:
+ adapter = _base_adapter(latent_storage_dtype="fp32", adapter_type=adapter_type)
+ condition = {
+ **_condition(
+ video_condition_rows=1,
+ audio_condition_rows=audio_condition_rows,
+ ),
+ "prompt_embeds": torch.zeros(1, 2, 4),
+ "condition_latents": [[torch.ones(1, 24, 1, 2, 2)]],
+ "audio_condition_latents": [
+ [torch.ones(audio_condition_rows, 32)] if audio_condition_rows else []
+ ],
+ }
+ prefixes = {
+ "video": torch.full((1, 1, 96), 3.0),
+ "audio": torch.full((1, audio_condition_rows, 32), 4.0),
+ }
+ calls = 0
+
+ def prepare(*args: Any, **kwargs: Any) -> dict[str, torch.Tensor]:
+ nonlocal calls
+ calls += 1
+ return prefixes
+
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3._condition.prepare_h3_condition_prefixes",
+ prepare,
+ )
+
+ prepared = adapter.prepare_condition_state(condition)
+ chosen = adapter.encode_output_state(_media(), prepared)
+ rejected = adapter.encode_output_state(_media(), prepared)
+ model_batch = dict(prepared.model_forward_condition())
+ times, noised = build_noised_output_state(
+ adapter,
+ chosen.clean_state,
+ torch.tensor([500.0]),
+ batch=model_batch,
+ generator=torch.Generator().manual_seed(31),
+ )
+ output = adapter._forward_state(
+ batch=SimpleNamespace(),
+ state=noised.state,
+ times=times,
+ next_state=None,
+ compute_log_prob=False,
+ return_fields=("velocity",),
+ noise_level=0.0,
+ forward_kwargs=model_batch,
+ )
+
+ assert calls == 1
+ validate_preference_output_states(chosen, rejected)
+ assert model_batch["condition_prefixes"] is prefixes
+ assert chosen.clean_state.components["video"].shape == (1, 7, 96)
+ assert output.velocity.components["video"].shape == (1, 7, 96)
+ assert output.velocity.components["audio"].shape == (1, 74, 32)
+
+
def test_conditioned_offline_replay_still_requires_prefix_binder() -> None:
adapter = object.__new__(MiniMaxH3FL2VAAdapter)
state = LatentState(
diff --git a/tests/models/minimax_h3/test_workflow_execution.py b/tests/models/minimax_h3/test_workflow_execution.py
index 46b468aa0..334937956 100644
--- a/tests/models/minimax_h3/test_workflow_execution.py
+++ b/tests/models/minimax_h3/test_workflow_execution.py
@@ -16,7 +16,10 @@
import pytest
import torch
+from PIL import Image
+from flow_factory.data_utils.offline_condition_cache import build_offline_condition_cache
+from flow_factory.data_utils.schema import normalize_v2_record
from flow_factory.models.minimax_h3.adapters import (
MiniMaxH3FL2VAAdapter,
MiniMaxH3Ref2VAAdapter,
@@ -63,6 +66,20 @@ def _adapter(adapter_class: type, transformer: Any = None) -> Any:
"num_frames": 5,
},
),
+ (
+ MiniMaxH3FL2VAAdapter,
+ {
+ "images": [["ending"]],
+ "image_slots": [["last_frame"]],
+ },
+ {
+ "prompt": "describe",
+ "last_image": "ending",
+ "height": 64,
+ "width": 96,
+ "num_frames": 5,
+ },
+ ),
],
)
def test_preprocess_uses_exact_workflow_inputs_and_b1(
@@ -87,14 +104,79 @@ def test_preprocess_uses_exact_workflow_inputs_and_b1(
adapter.preprocess_func(prompt=["one", "two"], height=64, width=96, num_frames=5)
+def test_v2_last_only_condition_reaches_h3_preprocess_through_arrow(tmp_path, monkeypatch) -> None:
+ """The complete public-schema/cache path preserves a sparse last-frame binding."""
+ ending_path = tmp_path / "ending.png"
+ Image.new("RGB", (16, 16), color=(12, 34, 56)).save(ending_path)
+ record = normalize_v2_record(
+ {
+ "schema_version": 2,
+ "input": {
+ "prompt": "Reveal what led to this ending.",
+ "media": [
+ {
+ "type": "image",
+ "path": ending_path.name,
+ "slot": "last_frame",
+ }
+ ],
+ },
+ "supervision": {
+ "type": "demonstration",
+ "target": {
+ "media": [
+ {"type": "video", "path": "target.mp4", "fps": 24.0},
+ {
+ "type": "audio",
+ "path": "target.wav",
+ "sample_rate": 32000,
+ },
+ ]
+ },
+ },
+ "metadata": {},
+ },
+ dataset_dir=tmp_path,
+ )
+ calls: List[Any] = []
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3.workflow.encode_h3_workflow_inputs",
+ lambda pipeline, values, workflow: calls.append((workflow, values))
+ or {"prompt_embeds": torch.zeros(1, 2, 4)},
+ )
+ adapter = _adapter(MiniMaxH3FL2VAAdapter)
+
+ cache = build_offline_condition_cache(
+ [record],
+ source_name="h3-last-only",
+ dataset_dir=tmp_path,
+ cache_dir=tmp_path / "cache",
+ preprocess_func=adapter.preprocess_func,
+ preprocess_kwargs={
+ "height": 64,
+ "width": 96,
+ "num_frames": 124,
+ },
+ pipeline_io_contract=adapter.pipeline_io_contract,
+ preprocessing_batch_size=1,
+ )
+
+ assert len(cache) == 1
+ assert len(calls) == 1
+ workflow, values = calls[0]
+ assert workflow == "fl2va"
+ assert "image" not in values
+ assert isinstance(values["last_image"], Image.Image)
+ assert values["last_image"].size == (16, 16)
+
+
def test_preprocess_adds_outer_batch_to_arrow_cache_fields(monkeypatch) -> None:
monkeypatch.setattr(
"flow_factory.models.minimax_h3.workflow.encode_h3_workflow_inputs",
lambda *args, **kwargs: {
"prompt_embeds": torch.zeros(1, 2, 4),
- "text_token_tags": torch.tensor([1, 1]),
+ "token_tags": torch.tensor([1, 1]),
"height": 64,
- "keyframe_anchors": (),
},
)
adapter = _adapter(MiniMaxH3T2VAAdapter)
@@ -107,10 +189,9 @@ def test_preprocess_adds_outer_batch_to_arrow_cache_fields(monkeypatch) -> None:
)
assert result["prompt_embeds"].shape == (1, 2, 4)
- assert len(result["text_token_tags"]) == 1
- torch.testing.assert_close(result["text_token_tags"][0], torch.tensor([1, 1]))
+ assert len(result["token_tags"]) == 1
+ torch.testing.assert_close(result["token_tags"][0], torch.tensor([1, 1]))
assert result["height"] == [64]
- assert result["keyframe_anchors"] == [[]]
def test_ref_preprocess_builds_ordered_pinned_objects_without_returning_them(monkeypatch) -> None:
diff --git a/tests/models/test_offline_output_capability_matrix.py b/tests/models/test_offline_output_capability_matrix.py
index 6305f3812..ee39b1a58 100644
--- a/tests/models/test_offline_output_capability_matrix.py
+++ b/tests/models/test_offline_output_capability_matrix.py
@@ -12,8 +12,6 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-import pytest
-
from flow_factory.models.ltx2.ltx2_i2av import LTX2_I2AV_Adapter
from flow_factory.models.ltx2.ltx2_t2av import LTX2_T2AV_Adapter
from flow_factory.models.minimax_h3.adapters import (
@@ -22,27 +20,20 @@
MiniMaxH3T2VAAdapter,
)
from flow_factory.models.wan.wan2_i2v import Wan2_I2V_Adapter
+from flow_factory.models.wan.wan2_t2v import Wan2_T2V_Adapter
-@pytest.mark.parametrize(
- ("adapter_type", "reason_fragment"),
- [
- (Wan2_I2V_Adapter, "first-frame VAE condition"),
- (LTX2_T2AV_Adapter, "paired video/audio decoding"),
- (LTX2_I2AV_Adapter, "active mask"),
- (MiniMaxH3FL2VAAdapter, "conditioned-prefix binder"),
- (MiniMaxH3Ref2VAAdapter, "conditioned-prefix binder"),
- ],
-)
-def test_unimplemented_offline_media_semantics_fail_before_model_loading(
- adapter_type: type,
- reason_fragment: str,
-) -> None:
- """Expose actionable blockers instead of silently guessing target encoding."""
- with pytest.raises(NotImplementedError, match=reason_fragment):
- adapter_type.validate_offline_output_capability()
+def test_video_and_av_adapters_declare_complete_offline_capability() -> None:
+ """Every implemented video/AV workflow exposes its complete offline codec."""
+ adapter_types = (
+ Wan2_T2V_Adapter,
+ Wan2_I2V_Adapter,
+ LTX2_T2AV_Adapter,
+ LTX2_I2AV_Adapter,
+ MiniMaxH3T2VAAdapter,
+ MiniMaxH3FL2VAAdapter,
+ MiniMaxH3Ref2VAAdapter,
+ )
-
-def test_minimax_h3_t2va_declares_complete_offline_output_semantics() -> None:
- """T2VA has no conditioned prefix and can encode paired AV targets on demand."""
- MiniMaxH3T2VAAdapter.validate_offline_output_capability()
+ for adapter_type in adapter_types:
+ adapter_type.validate_offline_output_capability()
diff --git a/tests/models/test_output_state_adapter_lifecycle.py b/tests/models/test_output_state_adapter_lifecycle.py
index 1b4fad7cb..0653e8673 100644
--- a/tests/models/test_output_state_adapter_lifecycle.py
+++ b/tests/models/test_output_state_adapter_lifecycle.py
@@ -14,7 +14,7 @@
"""Lightweight tests for the BaseAdapter output-state codec lifecycle seam."""
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
from types import SimpleNamespace
from typing import Any, Mapping, Optional, Tuple
@@ -35,6 +35,7 @@
RateRequirement,
)
from flow_factory.models.abc import BaseAdapter
+from flow_factory.models.condition_state import PreparedConditionState
from flow_factory.models.output_state import (
DecodedMediaBatch,
EncodedOutputState,
@@ -148,6 +149,38 @@ def encode_output_state(
return self.result or _encoded_image_batch(len(media_batch))
+class _ConditionPreparer:
+ required_components = ("vae",)
+
+ def __init__(self) -> None:
+ self.calls = 0
+ self.grad_enabled: Optional[bool] = None
+
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ del generator
+ self.calls += 1
+ self.grad_enabled = torch.is_grad_enabled()
+ return PreparedConditionState(
+ condition=condition,
+ forward_context={"condition_prefix": torch.zeros(1, 2)},
+ output_context={"codec_condition": torch.ones(1, 2)},
+ )
+
+
+class _StochasticConditionPreparer(_ConditionPreparer):
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: Optional[torch.Generator] = None,
+ ) -> PreparedConditionState:
+ torch.rand((), generator=generator)
+ return super().prepare_condition_state(condition, generator)
+
+
class _LifecycleAdapter(BaseAdapter):
pipeline_io_contract = IMAGE_CONTRACT
@@ -234,6 +267,15 @@ class _DefaultGeometryAdapter(_LifecycleAdapter):
_validate_encoded_output_geometry = BaseAdapter._validate_encoded_output_geometry
+class _PreparedLifecycleAdapter(_LifecycleAdapter):
+ def __init__(self, codec: Optional[_Codec], preparer: _ConditionPreparer) -> None:
+ self._preparer_to_build = preparer
+ super().__init__(codec)
+
+ def build_condition_state_preparer(self) -> _ConditionPreparer:
+ return self._preparer_to_build
+
+
def test_codec_build_runs_after_component_and_scheduler_lifecycle() -> None:
codec = _Codec()
@@ -244,6 +286,48 @@ def test_codec_build_runs_after_component_and_scheduler_lifecycle() -> None:
assert adapter.output_state_encoding_modules == ("vae",)
+def test_condition_preparer_is_declaration_only_and_runs_under_no_grad() -> None:
+ preparer = _ConditionPreparer()
+ adapter = _PreparedLifecycleAdapter(_Codec(), preparer)
+
+ prepared = adapter.prepare_condition_state({"prompt_embeds": torch.ones(1, 2)})
+
+ assert adapter.condition_state_preparer is preparer
+ assert adapter.condition_state_encoding_modules == ("vae",)
+ assert preparer.calls == 1
+ assert preparer.grad_enabled is False
+ assert tuple(prepared.forward_context) == ("condition_prefix",)
+ assert tuple(prepared.output_context) == ("codec_condition",)
+
+
+def test_raw_condition_encode_routes_through_declared_preparer() -> None:
+ preparer = _ConditionPreparer()
+ codec = _Codec()
+ adapter = _PreparedLifecycleAdapter(codec, preparer)
+ condition = {"prompt_embeds": torch.ones(1, 2)}
+
+ adapter.encode_output_state(_image_batch(1), condition)
+
+ assert preparer.calls == 1
+ assert codec.received is not None
+ assert tuple(codec.received[1]) == ("prompt_embeds", "codec_condition")
+ assert torch.equal(codec.received[1]["prompt_embeds"], condition["prompt_embeds"])
+ assert torch.equal(codec.received[1]["codec_condition"], torch.ones(1, 2))
+
+
+def test_default_condition_preparation_is_identity() -> None:
+ adapter = _LifecycleAdapter(_Codec())
+ condition = {"prompt_embeds": torch.ones(1, 2)}
+
+ prepared = adapter.prepare_condition_state(condition)
+
+ assert adapter.condition_state_preparer is None
+ assert adapter.condition_state_encoding_modules == ()
+ assert dict(prepared.condition) == condition
+ assert not prepared.forward_context
+ assert not prepared.output_context
+
+
def test_codec_build_must_remain_declaration_only() -> None:
class MaterializingCodecAdapter(_LifecycleAdapter):
def build_output_state_codec(self) -> Optional[_Codec]:
@@ -263,6 +347,22 @@ def test_contract_without_codec_preserves_online_adapter_construction() -> None:
adapter.encode_output_state(_image_batch(1), {})
+def test_effective_pipeline_contract_can_narrow_checkpoint_capability() -> None:
+ class SingleSampleAdapter(_LifecycleAdapter):
+ def _resolve_pipeline_io_contract(self) -> PipelineIOContract:
+ return replace(
+ self.pipeline_io_contract,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
+
+ adapter = SingleSampleAdapter(_Codec())
+
+ assert adapter.pipeline_io_contract.batch_capability is BatchCapability.UNIFORM
+ assert adapter.effective_pipeline_io_contract.batch_capability is BatchCapability.SINGLE_SAMPLE
+ with pytest.raises(ValueError, match=r"single_sample.*batch size 1.*2"):
+ adapter.encode_output_state(_image_batch(2), {})
+
+
def test_known_codec_blocker_is_actionable_at_direct_encode_boundary() -> None:
class KnownUnavailableAdapter(_OnlineOnlyAdapter):
output_state_codec_unavailable_reason = (
@@ -321,6 +421,12 @@ def test_codec_required_components_must_exist_in_runtime() -> None:
def test_public_output_state_wrapper_cannot_be_overridden() -> None:
+ with pytest.raises(TypeError, match=r"must not override BaseAdapter.prepare_condition_state"):
+
+ class InvalidConditionAdapter(_LifecycleAdapter):
+ def prepare_condition_state(self, *args: Any, **kwargs: Any) -> Any:
+ return None
+
with pytest.raises(TypeError, match=r"must not override BaseAdapter.encode_output_state"):
class InvalidAdapter(_LifecycleAdapter):
@@ -363,12 +469,12 @@ def test_encode_output_state_validates_invokes_no_grad_and_applies_storage_dtype
assert codec.grad_enabled is False
assert codec.received is not None
assert codec.received[0] is media_batch
- assert codec.received[1] is condition
+ assert dict(codec.received[1]) == condition
assert codec.received[2] is generator
assert encoded.clean_state.components["latent"].dtype is torch.float16
assert adapter.geometry_validation is not None
assert adapter.geometry_validation[0] is media_batch
- assert adapter.geometry_validation[1] is condition
+ assert dict(adapter.geometry_validation[1]) == condition
assert adapter.geometry_validation[2] is encoded
@@ -384,6 +490,20 @@ def test_encode_output_state_rejects_candidate_before_invoking_codec() -> None:
assert adapter.geometry_validation is None
+def test_encode_output_state_rejects_candidate_before_stochastic_preparation() -> None:
+ preparer = _StochasticConditionPreparer()
+ adapter = _PreparedLifecycleAdapter(_Codec(), preparer)
+ generator = torch.Generator().manual_seed(7)
+ original_state = generator.get_state().clone()
+ wrong_type = ((_DecodedMedia(type="video", payload=torch.zeros(1)),),)
+
+ with pytest.raises(ValueError, match=r"expected.*type 'image'.*'video'"):
+ adapter.encode_output_state(wrong_type, {}, generator)
+
+ assert preparer.calls == 0
+ assert torch.equal(generator.get_state(), original_state)
+
+
@pytest.mark.parametrize(
("condition", "generator", "message"),
[
diff --git a/tests/models/test_wan_output_codec.py b/tests/models/test_wan_output_codec.py
index ce2a0b1a3..e3dab0ae4 100644
--- a/tests/models/test_wan_output_codec.py
+++ b/tests/models/test_wan_output_codec.py
@@ -12,16 +12,38 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+from pathlib import Path
from types import SimpleNamespace
from typing import Any
import numpy as np
import pytest
import torch
-
-from flow_factory.contracts import GeometrySource, MediaType, RateRequirement
-from flow_factory.data_utils.offline_dataset import DecodedMedia
+from PIL import Image
+
+from flow_factory.contracts import (
+ BatchCapability,
+ GeometrySource,
+ InputMediaOrder,
+ MediaType,
+ RateRequirement,
+)
+from flow_factory.data_utils.offline_condition_cache import build_offline_condition_cache
+from flow_factory.data_utils.offline_dataset import (
+ OFFLINE_CONDITION_ID_COLUMN,
+ DecodedMedia,
+ _collate_condition_mappings,
+)
+from flow_factory.data_utils.schema import normalize_v2_record
+from flow_factory.models.abc import BaseAdapter
+from flow_factory.models.wan._conditioning import (
+ WanI2VConditionStatePreparer,
+ normalize_wan_i2v_image_rows,
+ normalize_wan_image_embeds,
+ split_wan_image_embeds,
+)
from flow_factory.models.wan._output import WanVideoOutputCodec
+from flow_factory.models.wan.wan2_i2v import Wan2_I2V_Adapter, WanI2VSample
from flow_factory.models.wan.wan2_t2v import Wan2_T2V_Adapter
@@ -29,11 +51,16 @@ class _Posterior:
def __init__(self, latents: torch.Tensor) -> None:
self.latents = latents
self.generators: list[torch.Generator | None] = []
+ self.mode_calls = 0
def sample(self, generator: torch.Generator | None = None) -> torch.Tensor:
self.generators.append(generator)
return self.latents
+ def mode(self) -> torch.Tensor:
+ self.mode_calls += 1
+ return self.latents
+
class _VAE:
dtype = torch.float32
@@ -56,6 +83,23 @@ def encode(self, pixels: torch.Tensor) -> Any:
class _VideoProcessor:
def __init__(self) -> None:
self.videos: list[list[np.ndarray]] = []
+ self.images: list[list[Image.Image]] = []
+
+ def preprocess(
+ self,
+ images: list[Image.Image],
+ *,
+ height: int,
+ width: int,
+ ) -> torch.Tensor:
+ self.images.append(images)
+ channels = [
+ torch.tensor(np.asarray(image, dtype=np.float32)[0, 0] / 255.0)
+ .view(3, 1, 1)
+ .expand(3, height, width)
+ for image in images
+ ]
+ return torch.stack(channels, dim=0)
def preprocess_video(
self,
@@ -68,6 +112,33 @@ def preprocess_video(
return torch.zeros(len(videos), 3, videos[0].shape[0], height, width)
+class _ProcessorOutput(dict):
+ def to(self, device: torch.device) -> "_ProcessorOutput":
+ del device
+ return self
+
+
+class _ImageProcessor:
+ def __call__(self, *, images: list[Image.Image], return_tensors: str) -> _ProcessorOutput:
+ assert return_tensors == "pt"
+ return _ProcessorOutput(pixel_values=torch.zeros(len(images), 3, 2, 2))
+
+
+class _ImageEncoder:
+ device = torch.device("cpu")
+
+ def __call__(
+ self,
+ pixel_values: torch.Tensor,
+ *,
+ output_hidden_states: bool,
+ ) -> Any:
+ assert output_hidden_states is True
+ count = pixel_values.shape[0]
+ embeds = torch.arange(count * 3 * 5, dtype=torch.float32).view(count, 3, 5)
+ return SimpleNamespace(hidden_states=(embeds, torch.zeros_like(embeds)))
+
+
class _Adapter:
_configured_video_output_geometry = Wan2_T2V_Adapter._configured_video_output_geometry
_resample_output_video = staticmethod(Wan2_T2V_Adapter._resample_output_video)
@@ -88,8 +159,35 @@ def __init__(self) -> None:
transformer=SimpleNamespace(config=SimpleNamespace(patch_size=(1, 2, 2))),
transformer_2=None,
video_processor=_VideoProcessor(),
+ config=SimpleNamespace(expand_timesteps=False),
+ )
+
+
+class _WanPreprocessHarness:
+ """Exercise BaseAdapter preprocessing with Wan's real image encoder hook."""
+
+ preprocess_func = BaseAdapter.preprocess_func
+ encode_image = Wan2_I2V_Adapter.encode_image
+ encode_video = BaseAdapter.encode_video
+ encode_audio = BaseAdapter.encode_audio
+ python_format_columns = frozenset()
+ supports_ordered_references = False
+
+ def __init__(self, *, with_clip: bool = False) -> None:
+ self.device = torch.device("cpu")
+ self.training_args = SimpleNamespace(height=16, width=16)
+ self.image_encoder = _ImageEncoder()
+ self.pipeline = SimpleNamespace(
+ transformer=SimpleNamespace(config=SimpleNamespace(image_dim=8 if with_clip else None)),
+ video_processor=_VideoProcessor(),
+ image_processor=_ImageProcessor(),
+ image_encoder=self.image_encoder,
)
+ def encode_prompt(self, prompt: list[str], **kwargs: Any) -> dict[str, torch.Tensor]:
+ del kwargs
+ return {"prompt_embeds": torch.ones(len(prompt), 3, 4)}
+
def _media(video: np.ndarray, fps: float = 8.0):
return (
@@ -104,12 +202,20 @@ def _media(video: np.ndarray, fps: float = 8.0):
)
+def _condition_pixels(count: int = 1) -> torch.Tensor:
+ pixels = [torch.zeros(3, 16, 16)]
+ if count == 2:
+ pixels.append(torch.ones(3, 16, 16))
+ return torch.stack(pixels, dim=0).unsqueeze(0)
+
+
def test_wan_t2v_declares_required_video_output_semantics() -> None:
contract = Wan2_T2V_Adapter.pipeline_io_contract
assert contract.geometry_source is GeometrySource.CONFIGURED
assert contract.output_media.items[0].type is MediaType.VIDEO
assert contract.output_media.items[0].fps is RateRequirement.REQUIRED
+ assert contract.batch_capability is BatchCapability.SINGLE_SAMPLE
Wan2_T2V_Adapter.validate_offline_output_capability()
adapter = object.__new__(Wan2_T2V_Adapter)
@@ -187,3 +293,242 @@ def test_wan_geometry_validator_rejects_output_context_drift() -> None:
{},
drifted,
)
+
+
+def test_wan_i2v_declares_ordered_first_optional_last_offline_contract() -> None:
+ contract = Wan2_I2V_Adapter.pipeline_io_contract
+
+ assert contract.geometry_source is GeometrySource.CONFIGURED
+ assert contract.batch_capability is BatchCapability.SINGLE_SAMPLE
+ assert contract.input_media.order is InputMediaOrder.WITHIN_TYPE
+ assert len(contract.input_media.rules) == 1
+ rule = contract.input_media.rules[0]
+ assert (rule.format.type, rule.min_count, rule.max_count) == (MediaType.IMAGE, 1, 2)
+ assert dict(Wan2_I2V_Adapter.offline_training_forward_overrides) == {
+ "guidance_scale": 1.0,
+ "guidance_scale_2": 1.0,
+ }
+ Wan2_I2V_Adapter.validate_offline_output_capability()
+
+ adapter = object.__new__(Wan2_I2V_Adapter)
+ assert isinstance(adapter.build_condition_state_preparer(), WanI2VConditionStatePreparer)
+ codec = adapter.build_output_state_codec()
+ assert isinstance(codec, WanVideoOutputCodec)
+ assert codec.bind_condition_active_mask is True
+
+ runtime = object.__new__(Wan2_I2V_Adapter)
+ runtime.pipeline = _Adapter().pipeline
+ effective = Wan2_I2V_Adapter._resolve_pipeline_io_contract(runtime)
+ assert effective.input_media.rules[0].max_count == 2
+ runtime.pipeline.config.expand_timesteps = True
+ effective = Wan2_I2V_Adapter._resolve_pipeline_io_contract(runtime)
+ assert effective.input_media.rules[0].max_count == 1
+
+
+def test_wan_i2v_condition_preparer_preserves_order_and_uses_posterior_mode() -> None:
+ adapter = _Adapter()
+ first = Image.new("RGB", (4, 4), color="red")
+ last = Image.new("RGB", (4, 4), color="blue")
+ image_embeds = torch.ones(1, 2, 4, 6)
+
+ prepared = WanI2VConditionStatePreparer(adapter).prepare_condition_state(
+ {
+ "images": [[first, last]],
+ "condition_images": _condition_pixels(2),
+ "prompt_embeds": torch.ones(1, 3, 5),
+ "image_embeds": image_embeds,
+ }
+ )
+
+ assert tuple(prepared.condition) == ("prompt_embeds",)
+ assert tuple(prepared.forward_context) == ("latent_condition", "image_embeds")
+ assert prepared.output_context == {}
+ assert prepared.forward_context["image_embeds"].shape == (2, 4, 6)
+ condition = prepared.forward_context["latent_condition"]
+ assert condition.shape == (1, 7, 2, 2, 2)
+ torch.testing.assert_close(condition[:, :4, 0], torch.ones(1, 4, 2, 2))
+ expected_last_mask = torch.tensor([0.0, 0.0, 0.0, 1.0]).view(1, 4, 1, 1)
+ torch.testing.assert_close(
+ condition[:, :4, 1],
+ expected_last_mask.expand(1, 4, 2, 2),
+ )
+ assert adapter.vae.posterior.mode_calls == 1
+ assert adapter.vae.posterior.generators == []
+ encoded_pixels = adapter.vae.encoded_pixels[0]
+ torch.testing.assert_close(encoded_pixels[:, :, 0], torch.zeros(1, 3, 16, 16))
+ torch.testing.assert_close(encoded_pixels[:, :, -1], torch.ones(1, 3, 16, 16))
+
+
+def test_wan_i2v_first_only_non_expanded_condition_has_no_output_mask() -> None:
+ adapter = _Adapter()
+ first = Image.new("RGB", (4, 4))
+
+ prepared = WanI2VConditionStatePreparer(adapter).prepare_condition_state(
+ {"images": [[first]], "condition_images": _condition_pixels()}
+ )
+
+ condition = prepared.forward_context["latent_condition"]
+ torch.testing.assert_close(condition[:, :4, 0], torch.ones(1, 4, 2, 2))
+ torch.testing.assert_close(condition[:, :4, 1], torch.zeros(1, 4, 2, 2))
+ assert "first_frame_mask" not in prepared.forward_context
+ assert prepared.output_context == {}
+
+
+def test_wan_i2v_expand_condition_binds_target_active_mask() -> None:
+ adapter = _Adapter()
+ adapter.pipeline.config.expand_timesteps = True
+ first = Image.new("RGB", (4, 4))
+ prepared = WanI2VConditionStatePreparer(adapter).prepare_condition_state(
+ {"images": [[first]], "condition_images": _condition_pixels()}
+ )
+
+ condition = prepared.forward_context["latent_condition"]
+ mask = prepared.forward_context["first_frame_mask"]
+ assert condition.shape == (1, 3, 2, 2, 2)
+ assert mask.shape == (1, 1, 2, 2, 2)
+ torch.testing.assert_close(mask[:, :, 0], torch.zeros(1, 1, 2, 2))
+ torch.testing.assert_close(mask[:, :, 1], torch.ones(1, 1, 2, 2))
+ assert prepared.output_context["first_frame_mask"] is mask
+
+ source = np.zeros((9, 4, 4, 3), dtype=np.uint8)
+ encoded = WanVideoOutputCodec(adapter, bind_condition_active_mask=True).encode_output_state(
+ _media(source),
+ prepared.output_codec_condition(),
+ )
+
+ active_mask = encoded.clean_state.active_masks["latent"]
+ assert active_mask.dtype is torch.bool
+ torch.testing.assert_close(active_mask, mask.bool())
+ assert adapter.vae.posterior.mode_calls == 1
+ assert adapter.vae.posterior.generators == [None]
+
+
+def test_wan_i2v_expand_rejects_last_frame_instead_of_ignoring_it() -> None:
+ adapter = _Adapter()
+ adapter.pipeline.config.expand_timesteps = True
+ first = Image.new("RGB", (4, 4))
+ last = Image.new("RGB", (4, 4))
+
+ with pytest.raises(ValueError, match="expand_timesteps does not support"):
+ WanI2VConditionStatePreparer(adapter).prepare_condition_state(
+ {"images": [[first, last]], "condition_images": _condition_pixels(2)}
+ )
+
+ assert adapter.vae.encoded_pixels == []
+
+
+def test_wan_i2v_expand_target_requires_prepared_active_mask_before_vae() -> None:
+ adapter = _Adapter()
+ adapter.pipeline.config.expand_timesteps = True
+
+ with pytest.raises(ValueError, match="requires first_frame_mask"):
+ WanVideoOutputCodec(adapter, bind_condition_active_mask=True).encode_output_state(
+ _media(np.zeros((9, 4, 4, 3), dtype=np.uint8)),
+ {},
+ )
+
+ assert adapter.vae.encoded_pixels == []
+ assert adapter.vae.posterior.generators == []
+
+
+def test_wan_i2v_normalization_never_truncates_optional_last_frame() -> None:
+ first = Image.new("RGB", (4, 4))
+ last = Image.new("RGB", (4, 4))
+
+ rows = normalize_wan_i2v_image_rows([[first, last]], expected_batch_size=1)
+
+ assert rows == ((first, last),)
+
+
+def test_wan_i2v_online_prepare_latents_reuses_condition_mode_path() -> None:
+ adapter = _Adapter()
+ first_pixels = torch.zeros(1, 3, 16, 16)
+ last_pixels = torch.ones(1, 3, 16, 16)
+ rollout_noise = torch.zeros(1, 3, 2, 2, 2)
+
+ latents, condition = Wan2_I2V_Adapter.prepare_latents(
+ adapter,
+ first_pixels,
+ batch_size=1,
+ num_channels_latents=3,
+ height=16,
+ width=16,
+ num_frames=5,
+ dtype=torch.float32,
+ device=torch.device("cpu"),
+ latents=rollout_noise,
+ last_image=last_pixels,
+ )
+
+ assert latents is rollout_noise
+ assert condition.shape == (1, 7, 2, 2, 2)
+ assert adapter.vae.posterior.mode_calls == 1
+ assert adapter.vae.posterior.generators == []
+
+
+def test_wan_v2_condition_cache_keeps_ordered_pixels_through_prepare(
+ tmp_path: Path,
+) -> None:
+ Image.new("RGB", (4, 4), color="red").save(tmp_path / "first.png")
+ Image.new("RGB", (4, 4), color="blue").save(tmp_path / "last.png")
+ record = normalize_v2_record(
+ {
+ "schema_version": 2,
+ "input": {
+ "prompt": "first and last",
+ "media": [
+ {"type": "image", "path": "first.png"},
+ {"type": "image", "path": "last.png"},
+ ],
+ },
+ "supervision": {
+ "type": "demonstration",
+ "target": {"media": [{"type": "video", "path": "target.mp4"}]},
+ },
+ },
+ dataset_dir=tmp_path,
+ )
+ cache = build_offline_condition_cache(
+ [record],
+ source_name="wan-i2v",
+ dataset_dir=tmp_path,
+ cache_dir=tmp_path / "cache",
+ preprocess_func=_WanPreprocessHarness(with_clip=True).preprocess_func,
+ preprocessing_batch_size=1,
+ force_reprocess=True,
+ )
+ cache_row = dict(cache[0])
+ cache_row.pop(OFFLINE_CONDITION_ID_COLUMN)
+
+ condition = _collate_condition_mappings([cache_row])
+ assert len(condition["images"]) == 1
+ assert len(condition["images"][0]) == 2
+ assert condition["condition_images"].shape == (1, 2, 3, 16, 16)
+ torch.testing.assert_close(
+ condition["condition_images"][0, :, :, 0, 0],
+ torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
+ )
+ assert condition["image_embeds"].shape == (1, 2, 3, 5)
+ prepared = WanI2VConditionStatePreparer(_Adapter()).prepare_condition_state(condition)
+
+ assert "images" not in prepared.condition
+ assert prepared.forward_context["latent_condition"].shape == (1, 7, 2, 2, 2)
+ assert prepared.forward_context["image_embeds"].shape == (2, 3, 5)
+
+
+def test_wan_i2v_two_sample_first_last_embeds_round_trip_through_replay_stack() -> None:
+ packed = torch.arange(4 * 3 * 5, dtype=torch.float32).view(4, 3, 5)
+ per_sample = split_wan_image_embeds(packed, (2, 2))
+ samples = [
+ WanI2VSample(
+ image_embeds=per_sample[index],
+ latent_condition=torch.zeros(3, 2, 2, 2),
+ )
+ for index in range(2)
+ ]
+
+ replay_batch = WanI2VSample.stack(samples)
+ assert replay_batch["image_embeds"].shape == (2, 2, 3, 5)
+ restored = normalize_wan_image_embeds(replay_batch["image_embeds"], batch_size=2)
+
+ torch.testing.assert_close(restored, packed)
diff --git a/tests/models/trajectory/test_base_adapter_trajectory.py b/tests/models/trajectory/test_base_adapter_trajectory.py
index 72eab6c40..6d4360161 100644
--- a/tests/models/trajectory/test_base_adapter_trajectory.py
+++ b/tests/models/trajectory/test_base_adapter_trajectory.py
@@ -489,6 +489,49 @@ def test_reduce_latent_values_uses_global_element_weighting() -> None:
)
+def test_flow_matching_objective_reduction_defaults_to_global_weighting() -> None:
+ adapter = _structured_adapter()
+ values = {
+ "video": torch.tensor([[1.0, 3.0], [2.0, 4.0]]),
+ "audio": torch.tensor([[10.0], [20.0]]),
+ }
+
+ assert torch.equal(
+ adapter.reduce_flow_matching_objective_values(values),
+ adapter.reduce_latent_values(values),
+ )
+
+
+class FlowObjectiveReducerFake(StructuredAdapterFake):
+ """Adapter summing independently normalized modality objectives."""
+
+ def _reduce_flow_matching_objective_values(
+ self,
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: Optional[LatentState] = None,
+ ) -> torch.Tensor:
+ reduced = self.reduce_component_latent_values(values, state=state)
+ return reduced["video"] + reduced["audio"]
+
+
+def test_flow_matching_objective_reduction_can_preserve_online_reducer() -> None:
+ adapter = object.__new__(FlowObjectiveReducerFake)
+ values = {
+ "video": torch.tensor([[1.0, 3.0], [2.0, 4.0]]),
+ "audio": torch.tensor([[10.0], [20.0]]),
+ }
+
+ assert torch.equal(
+ adapter.reduce_flow_matching_objective_values(values),
+ torch.tensor([12.0, 23.0]),
+ )
+ assert torch.equal(
+ adapter.reduce_latent_values(values),
+ torch.tensor([14.0 / 3.0, 26.0 / 3.0]),
+ )
+
+
def test_reduce_latent_values_preserves_single_component_scalar_scale() -> None:
values = torch.tensor([2.0, 3.0])
diff --git a/tests/trainers/test_offline_batch_primitives.py b/tests/trainers/test_offline_batch_primitives.py
index dbc82d9fb..18c3d594c 100644
--- a/tests/trainers/test_offline_batch_primitives.py
+++ b/tests/trainers/test_offline_batch_primitives.py
@@ -19,8 +19,10 @@
import torch
from flow_factory.contracts import NON_MODEL_CONDITION_KEYS
+from flow_factory.models.condition_state import PreparedConditionState
from flow_factory.trainers.common.offline_batch import (
bind_output_forward_context,
+ bind_prepared_condition_output,
move_condition_to_device,
)
@@ -86,6 +88,30 @@ def test_bind_output_context_preserves_input_ownership_without_mutation() -> Non
assert tuple(context) == ("img_ids",)
+def test_bind_prepared_condition_preserves_input_and_output_ownership() -> None:
+ prepared = PreparedConditionState(
+ condition={"prompt_embeds": torch.ones(1, 2)},
+ forward_context={"condition_prefix": torch.zeros(1, 3)},
+ output_context={"codec_only": torch.ones(1)},
+ )
+
+ bound = bind_prepared_condition_output(prepared, {"output_ids": torch.zeros(4, 3)})
+
+ assert tuple(bound) == ("prompt_embeds", "condition_prefix", "output_ids")
+ assert "codec_only" not in bound
+
+
+def test_bind_prepared_condition_rejects_output_collision() -> None:
+ prepared = PreparedConditionState(
+ condition={"prompt_embeds": torch.ones(1, 2)},
+ forward_context={"condition_prefix": torch.zeros(1, 3)},
+ output_context={},
+ )
+
+ with pytest.raises(ValueError, match=r"collides.*condition_prefix"):
+ bind_prepared_condition_output(prepared, {"condition_prefix": torch.ones(1, 3)})
+
+
def test_bind_output_context_rejects_ambiguous_key_ownership() -> None:
with pytest.raises(ValueError, match=r"collides.*\('geometry',\)"):
bind_output_forward_context({"geometry": 1}, {"geometry": 2})
diff --git a/tests/trainers/test_offline_flow_matching.py b/tests/trainers/test_offline_flow_matching.py
index 8f24a0f36..123564a9e 100644
--- a/tests/trainers/test_offline_flow_matching.py
+++ b/tests/trainers/test_offline_flow_matching.py
@@ -245,7 +245,7 @@ def test_flow_matching_loss_computes_fp32_errors_before_adapter_reduction() -> N
received: dict[str, Any] = {}
class Adapter:
- def reduce_latent_values(self, values: Any, *, state: Any):
+ def reduce_flow_matching_objective_values(self, values: Any, *, state: Any):
received["values"] = values
received["state"] = state
total = torch.cat([value.flatten(1) for value in values.values()], dim=1)
diff --git a/tests/trainers/test_offline_trainers.py b/tests/trainers/test_offline_trainers.py
index 547ec44b8..a91fa99e3 100644
--- a/tests/trainers/test_offline_trainers.py
+++ b/tests/trainers/test_offline_trainers.py
@@ -30,6 +30,7 @@
PreferenceOutputBatch,
)
from flow_factory.data_utils.schema import NormalizedModelInput
+from flow_factory.models.condition_state import PreparedConditionState
from flow_factory.models.output_state import (
EncodedOutputState,
GeometrySignature,
@@ -108,6 +109,8 @@ def __init__(self) -> None:
self.pipeline_io_contract = object()
self.train_calls = 0
self.encode_calls: list[str] = []
+ self.prepare_calls = 0
+ self.prepared_condition_ids: list[int] = []
self.forward_events: list[tuple[float, bool, bool]] = []
self.forward_override_events: list[tuple[float, bool, dict[str, float]]] = []
self.drawn_noise: list[LatentState] = []
@@ -119,13 +122,23 @@ def train(self, mode: bool = True) -> None:
assert mode is True
self.train_calls += 1
+ def prepare_condition_state(
+ self,
+ condition: Mapping[str, Any],
+ generator: torch.Generator | None = None,
+ ) -> PreparedConditionState:
+ del generator
+ self.prepare_calls += 1
+ return PreparedConditionState.identity(condition)
+
def encode_output_state(
self,
media_batch: tuple[tuple[DecodedMedia, ...], ...],
- condition: Mapping[str, Any],
+ condition: Mapping[str, Any] | PreparedConditionState,
generator: torch.Generator | None = None,
) -> EncodedOutputState:
- del condition, generator
+ self.prepared_condition_ids.append(id(condition))
+ del generator
arm = str(media_batch[0][0].payload)
self.encode_calls.append(arm)
arm_value = 1.0 if arm == "rejected" else 0.0
@@ -255,6 +268,14 @@ def reduce_latent_values(
del state
return values["latent"].flatten(1).mean(dim=1)
+ @staticmethod
+ def reduce_flow_matching_objective_values(
+ values: Mapping[str, torch.Tensor],
+ *,
+ state: LatentState,
+ ) -> torch.Tensor:
+ return _Adapter.reduce_latent_values(values, state=state)
+
def _media(arm: str, batch_size: int = 2) -> tuple[tuple[DecodedMedia, ...], ...]:
return tuple(
@@ -355,6 +376,7 @@ def test_offline_trainers_build_unprepared_distributed_loaders(
trainer.adapter = SimpleNamespace(
preprocess_func=object(),
pipeline_io_contract=object(),
+ effective_pipeline_io_contract=object(),
)
sentinel = object()
received: dict[str, Any] = {}
@@ -374,7 +396,7 @@ def fake_builder(**kwargs: Any) -> Any:
"accelerator": trainer.accelerator,
"preprocess_func": trainer.adapter.preprocess_func,
"supervision_type": supervision_type,
- "pipeline_io_contract": trainer.adapter.pipeline_io_contract,
+ "pipeline_io_contract": trainer.adapter.effective_pipeline_io_contract,
}
assert trainer.accelerator.prepare_calls == 0
@@ -395,6 +417,7 @@ def test_sft_reencodes_targets_and_preserves_optimizer_cadence(
assert trainer.step == 1
assert adapter.train_calls == 2
+ assert adapter.prepare_calls == 2
assert adapter.encode_calls == ["target", "target"]
assert len(trainer.accelerator.backward_losses) == 2
assert len(trainer.accelerator.accumulate_roots) == 2
@@ -443,6 +466,8 @@ def test_offline_dpo_shares_schedule_noise_and_reference_scope(
assert trainer.step == 1
assert adapter.encode_calls == ["chosen", "rejected"]
+ assert adapter.prepare_calls == 1
+ assert adapter.prepared_condition_ids[0] == adapter.prepared_condition_ids[1]
assert len(adapter.drawn_noise) == 2
assert len(adapter.reused_noise) == 4
assert adapter.reused_noise[0] is adapter.drawn_noise[0]
diff --git a/tests/trainers/test_runtime_identity.py b/tests/trainers/test_runtime_identity.py
index 3f8cbb635..7350e667a 100644
--- a/tests/trainers/test_runtime_identity.py
+++ b/tests/trainers/test_runtime_identity.py
@@ -24,6 +24,7 @@
from accelerate.utils import DistributedType
from torch.utils.data import ConcatDataset, DataLoader, DistributedSampler
+from flow_factory.contracts import NegativePromptPolicy
from flow_factory.contracts.execution import OFFLINE_EXECUTION_CONTRACT
from flow_factory.data_utils.dataset import GeneralDataset
from flow_factory.data_utils.multi_source import (
@@ -36,6 +37,7 @@
AdamWOptimizerArguments,
MultiOptimizerArguments,
)
+from flow_factory.models.pipeline_contracts import image_output_contract
from flow_factory.trainers.common.runtime_identity import (
build_trainer_runtime_identity,
)
@@ -201,10 +203,16 @@ def __init__(
log_every: int = 10,
max_grad_norm: float = 1.0,
update_frequency: int = 1,
+ accepts_image_input: bool = False,
dataloader: DataLoader | None = None,
) -> None:
parameter = torch.nn.Parameter(torch.zeros(width, width))
self.adapter = _Adapter()
+ self.adapter.effective_pipeline_io_contract = image_output_contract(
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ input_image_min_count=0 if accepts_image_input else None,
+ input_image_max_count=1 if accepts_image_input else None,
+ )
self.adapter.component_variant_registry = _Registry(
{"base": (_Record("transformer", "weight", parameter),)}
)
@@ -320,6 +328,15 @@ def test_parameter_and_optimizer_schema_changes_have_independent_digests() -> No
assert changed_optimizer["optimizer_schema_digest"] != baseline["optimizer_schema_digest"]
+def test_effective_pipeline_contract_changes_execution_identity() -> None:
+ """Checkpoint-realized input semantics are exact-resume boundaries."""
+ baseline = build_trainer_runtime_identity(_Trainer(accepts_image_input=False))
+ changed = build_trainer_runtime_identity(_Trainer(accepts_image_input=True))
+
+ assert changed["execution_contract_digest"] != baseline["execution_contract_digest"]
+ assert changed["data_contract_digest"] == baseline["data_contract_digest"]
+
+
@pytest.mark.parametrize(
("field", "changed_value"),
(("max_grad_norm", 0.5), ("update_frequency", 3)),
From 850f61dd746cfd4926628ee3918f26643b01c932 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 10:07:42 +0800
Subject: [PATCH 32/76] feat(dataset): add public offline smoke builders
---
.agents/knowledge/topics/fix_patterns.md | 8 +
.gitignore | 4 +
README.md | 2 +
dataset/offline_smoke/README.md | 56 ++
dataset/offline_smoke/SOURCES.md | 32 +
dataset/offline_smoke/__init__.py | 45 ++
dataset/offline_smoke/build_mini.py | 666 +++++++++++++++++++
dataset/offline_smoke/datasets.lock.json | 15 +
dataset/offline_smoke/prepare.py | 319 +++++++++
dataset/offline_smoke/profiles.py | 317 +++++++++
dataset/offline_smoke/publish.py | 225 +++++++
dataset/offline_smoke/validate.py | 233 +++++++
guidance/datasets.md | 35 +
guidance/gpu_validation.md | 27 +-
tests/dataset/test_offline_smoke_build.py | 243 +++++++
tests/dataset/test_offline_smoke_prepare.py | 260 ++++++++
tests/dataset/test_offline_smoke_profiles.py | 144 ++++
tests/dataset/test_offline_smoke_publish.py | 98 +++
18 files changed, 2728 insertions(+), 1 deletion(-)
create mode 100644 dataset/offline_smoke/README.md
create mode 100644 dataset/offline_smoke/SOURCES.md
create mode 100644 dataset/offline_smoke/__init__.py
create mode 100644 dataset/offline_smoke/build_mini.py
create mode 100644 dataset/offline_smoke/datasets.lock.json
create mode 100644 dataset/offline_smoke/prepare.py
create mode 100644 dataset/offline_smoke/profiles.py
create mode 100644 dataset/offline_smoke/publish.py
create mode 100644 dataset/offline_smoke/validate.py
create mode 100644 tests/dataset/test_offline_smoke_build.py
create mode 100644 tests/dataset/test_offline_smoke_prepare.py
create mode 100644 tests/dataset/test_offline_smoke_profiles.py
create mode 100644 tests/dataset/test_offline_smoke_publish.py
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 401c79356..2223bed35 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -246,6 +246,14 @@ Based on the fix type, write the fix entry to the appropriate document:
numerically lossy or RNG-consuming transition work that is outside their requested output.
- **Related Constraint**: #7
+### Documented dataset tools must use their package invocation mode
+- **Date**: 2026-08-30
+- **Symptom**: Running the documented `python dataset/offline_smoke/prepare.py ...` command failed before argument parsing with `attempted relative import with no known parent package`.
+- **Root Cause**: The package uses relative imports, while the documentation incorrectly advertised direct file execution instead of Python's module mode.
+- **Fix**: All offline-smoke commands now use `python -m dataset.offline_smoke.` with unconditional package imports, and a subprocess regression executes the documented form.
+- **Lesson**: A checked-in CLI example is part of the public interface; standardize on one package-aware invocation and test that exact process rather than adding conditional import fallbacks.
+- **Related Constraint**: N/A
+
### Exact resume must lock the checkpoint-realized pipeline contract
- **Date**: 2026-08-29
- **Symptom**: Exact resume could accept a checkpoint after an in-place model configuration change
diff --git a/.gitignore b/.gitignore
index 2f84609c7..854e74f24 100644
--- a/.gitignore
+++ b/.gitignore
@@ -109,6 +109,10 @@ latest_checkpointed_iteration.txt
# Dataset assets (keep metadata, ignore media)
# ====================
+# Reconstructed/publication staging for the SFT and offline-DPO smoke datasets
+dataset/_prepared_offline_smoke/
+dataset/offline_smoke/_staging/
+
# gitattributes in dataset
dataset/**/.gitattributes
diff --git a/README.md b/README.md
index 8c0a23ca1..6e6ca2788 100644
--- a/README.md
+++ b/README.md
@@ -244,6 +244,8 @@ encoded on the fly; their VAE latents are never stored in the preprocessing cach
epoch is one complete dataloader traversal sharded by PyTorch's official `DistributedSampler`. See the
[dataset guide](guidance/datasets.md#offline-v2-records) for the full schema and cadence rules, and
the [GPU validation plan](guidance/gpu_validation.md) for the 120-job model/backend/algorithm matrix.
+The [offline smoke builder](dataset/offline_smoke/README.md) reconstructs independent SFT and
+offline-DPO mini datasets for every currently implemented image, video, and audio-video profile.
## Text-to-Image & Text-to-Video
diff --git a/dataset/offline_smoke/README.md b/dataset/offline_smoke/README.md
new file mode 100644
index 000000000..09c8533eb
--- /dev/null
+++ b/dataset/offline_smoke/README.md
@@ -0,0 +1,56 @@
+# Offline smoke datasets
+
+This directory contains the reproducible build, preparation, validation, and
+publication tooling for two independent public fixtures:
+
+- [Jayce-Ping/Flow-Factory-SFT-Smoke](https://huggingface.co/datasets/Jayce-Ping/Flow-Factory-SFT-Smoke)
+- [Jayce-Ping/Flow-Factory-Offline-DPO-Smoke](https://huggingface.co/datasets/Jayce-Ping/Flow-Factory-Offline-DPO-Smoke)
+
+The repositories are intentionally small correctness fixtures. They are not
+quality-training corpora. All media are deterministic procedural assets released
+under CC0-1.0; the builder does not download third-party data or create VAE
+latent caches.
+
+See [SOURCES.md](SOURCES.md) for the DiffSynth, DyRef, image, video, and audio-video
+datasets reviewed for schema compatibility and the reasons they are not silently
+mirrored into these public fixtures.
+
+## Build staging repositories
+
+Install the project dependencies, including Pillow, NumPy, and PyAV, then run:
+
+```bash
+python -m dataset.offline_smoke.build_mini \
+ --staging-root dataset/offline_smoke/_staging
+```
+
+The command creates two self-contained trees. It refuses to overwrite an
+existing staging tree unless `--replace` is supplied. A fixed seed and fixed
+dependency versions produce the same logical records and media.
+
+Each tree contains:
+
+```text
+README.md
+LICENSE
+dataset_manifest.json
+provenance.jsonl
+media/
+profiles//train.jsonl
+```
+
+Every runtime alias has 32 strict V2 records, enough for two local batches on up
+to 16 ranks with per-device batch size one. Media paths in each JSONL are
+relative to its profile directory (`../../media/...`). The aliases cover the ten
+main GPU modes plus the supplemental `image-i2i` contract gate.
+
+Video DPO corruptions preserve the decoded first and last frames. Audio-video
+candidates follow the declared exact output order `[video, audio]`. The generic
+candidate projection follows `profiles.py` rather than assuming that every
+multi-component output is AV, so a future pure-audio or other output contract can
+be added without changing the public V2 record model.
+
+Preparation performs V2, task-profile contract, path, and media validation; see
+`python -m dataset.offline_smoke.prepare --help`. Publication is a separate
+external-state operation and requires the explicit confirmation flag shown by
+`python -m dataset.offline_smoke.publish --help`.
diff --git a/dataset/offline_smoke/SOURCES.md b/dataset/offline_smoke/SOURCES.md
new file mode 100644
index 000000000..1d3acf517
--- /dev/null
+++ b/dataset/offline_smoke/SOURCES.md
@@ -0,0 +1,32 @@
+# Dataset source review
+
+The published smoke repositories use only deterministic procedural media generated by
+`build_mini.py`. This keeps the fixtures self-contained, redistributable, and exact enough to
+exercise endpoint and audio/video clock contracts. They are correctness fixtures, not substitutes
+for semantic-quality training data.
+
+The following public datasets informed the schema and remain candidates for optional, separately
+licensed semantic demos:
+
+| Source | Relevant profiles | License/access decision |
+|---|---|---|
+| [ShareGPT-4o-Image](https://huggingface.co/datasets/FreedomIntelligence/ShareGPT-4o-Image) | T2I and I2I SFT | Apache-2.0. A future mini must preserve both input and output images; the legacy Flow-Factory mini retained only the output. |
+| [Open Image Preferences](https://huggingface.co/datasets/data-is-better-together/open-image-preferences-v1-binarized) | T2I offline DPO | Apache-2.0 human preference pairs. Suitable for a separately attributed semantic subset. |
+| [MACRO](https://huggingface.co/datasets/Azily/Macro-Dataset) | Ordered multi-image-to-image SFT | CC-BY-4.0. Preferred over non-commercial multi-reference alternatives when attribution is retained. |
+| [DyRef](https://github.com/Weistrass/DyRef) and [OmniRef training data](https://huggingface.co/datasets/Eason0438/OmniRef-training) | Ordered multi-image SFT and online RL | The pipeline is SFT followed by online RL, not offline DPO. The Hub dataset is tagged Apache-2.0, but per-asset provenance is incomplete, so the smoke repositories do not mirror it. |
+| [T2V Ranking Human Preferences](https://huggingface.co/datasets/datapointai/text-2-video-ranking-human-preferences) | T2V offline DPO | CC-BY-4.0 and gated. A local opt-in importer may derive best/worst pairs after the user accepts upstream access; the public smoke data must not bypass the gate. |
+| [I2V Human Preferences](https://huggingface.co/datasets/datapointai/image-2-video-human-preferences-large) | I2V offline DPO | CC-BY-4.0 and gated. Its shared reference image matches the I2V preference contract, but it is not mirrored here. |
+| [VA-Judger-Bench](https://huggingface.co/datasets/ShareLab-SII/VA-Judger-Bench) | T2AV preference structure | MIT-tagged paired audio/video benchmark. Kept as a schema reference rather than republishing benchmark media as training data. |
+| [DiffSynth-Studio examples](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples) | Wan, LTX2, and MiniMax-H3 SFT field flow | Apache-2.0 code. Example media lack per-file provenance, so only the field and media-clock conventions were reused. |
+
+The following sources are intentionally excluded from the public smoke repositories:
+
+- SafeSora and ONE-Lab MultiRef: non-commercial dataset terms.
+- Pico-Banana preference data: non-commercial, no-derivatives terms.
+- VideoDPO release media: no sufficiently clear data redistribution license was found.
+- MiniMax-H3-generated media: the model output license is not suitable for a model-neutral,
+ cross-family training fixture.
+
+Any future semantic subset must pin an immutable upstream revision and retain source row IDs,
+license, attribution, transforms, and per-asset SHA-256 values. Mixed-license assets must not be
+relicensed as CC0 or Apache-2.0.
diff --git a/dataset/offline_smoke/__init__.py b/dataset/offline_smoke/__init__.py
new file mode 100644
index 000000000..c606120cc
--- /dev/null
+++ b/dataset/offline_smoke/__init__.py
@@ -0,0 +1,45 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Build declarations for Flow-Factory's public offline smoke datasets."""
+
+from .profiles import (
+ CANONICAL_PROFILES,
+ DATASET_REPO_IDS,
+ GPU_ALIAS_TO_PROFILE,
+ MAIN_GPU_ALIASES,
+ OFFLINE_DPO_REPO_ID,
+ SFT_REPO_ID,
+ SUPPLEMENTAL_GPU_ALIASES,
+ GPUSmokeCase,
+ OfflineSmokeProfile,
+ SmokeGeometry,
+ get_profile,
+ output_media_types,
+)
+
+__all__ = [
+ "CANONICAL_PROFILES",
+ "DATASET_REPO_IDS",
+ "GPUSmokeCase",
+ "GPU_ALIAS_TO_PROFILE",
+ "MAIN_GPU_ALIASES",
+ "OFFLINE_DPO_REPO_ID",
+ "OfflineSmokeProfile",
+ "SFT_REPO_ID",
+ "SUPPLEMENTAL_GPU_ALIASES",
+ "SmokeGeometry",
+ "get_profile",
+ "output_media_types",
+]
diff --git a/dataset/offline_smoke/build_mini.py b/dataset/offline_smoke/build_mini.py
new file mode 100644
index 000000000..6c5378f39
--- /dev/null
+++ b/dataset/offline_smoke/build_mini.py
@@ -0,0 +1,666 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Build the two deterministic, procedural offline-smoke HF staging trees."""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import math
+import os
+import shutil
+import tempfile
+import wave
+from dataclasses import dataclass
+from fractions import Fraction
+from pathlib import Path
+from typing import Any, Iterable, Literal, Mapping, Sequence
+
+import av
+import numpy as np
+from PIL import Image, ImageDraw
+
+from .profiles import (
+ GPU_ALIAS_TO_PROFILE,
+ MAIN_GPU_ALIASES,
+ OFFLINE_DPO_REPO_ID,
+ SFT_REPO_ID,
+)
+
+RUNTIME_ALIASES = MAIN_GPU_ALIASES + ("image-i2i",)
+RECORDS_PER_ALIAS = 32
+DEFAULT_SEED = 20260830
+FPS = 24
+CC0 = "CC0-1.0"
+COLORS = (
+ ("coral", (224, 83, 74)),
+ ("amber", (234, 166, 52)),
+ ("teal", (42, 157, 143)),
+ ("blue", (65, 105, 225)),
+ ("violet", (139, 92, 246)),
+)
+SHAPES = ("circle", "square", "triangle", "diamond", "star")
+
+
+@dataclass(frozen=True, slots=True)
+class _Asset:
+ """One repo-root-relative generated media asset."""
+
+ type: Literal["image", "video", "audio"]
+ path: str
+ fps: int | None = None
+ sample_rate: int | None = None
+
+
+class _Writer:
+ """Write one self-contained repository and its provenance index."""
+
+ def __init__(
+ self,
+ root: Path,
+ repo_id: str,
+ supervision: Literal["demonstration", "preference"],
+ seed: int,
+ records_per_alias: int,
+ ) -> None:
+ self.root = root
+ self.repo_id = repo_id
+ self.supervision = supervision
+ self.seed = seed
+ self.records_per_alias = records_per_alias
+ self.provenance: list[dict[str, Any]] = []
+ self.script_sha = _sha256(Path(__file__))
+
+ def image(self, relative: str, value: Image.Image, seed: int) -> _Asset:
+ path = self._new_path(relative)
+ value.convert("RGB").save(path, format="PNG", compress_level=9)
+ self._record(relative, "image", seed, width=value.width, height=value.height)
+ return _Asset("image", relative)
+
+ def video(self, relative: str, frames: Sequence[np.ndarray], seed: int) -> _Asset:
+ path = self._new_path(relative)
+ height, width, channels = frames[0].shape
+ if channels != 3 or any(frame.shape != frames[0].shape for frame in frames):
+ raise ValueError("video frames must share one RGB geometry")
+ with av.open(str(path), "w", format="avi") as container:
+ stream = container.add_stream("ffv1", rate=Fraction(FPS, 1))
+ stream.width, stream.height, stream.pix_fmt = width, height, "bgr0"
+ for index, array in enumerate(frames):
+ frame = av.VideoFrame.from_ndarray(array, format="rgb24")
+ frame.pts = index
+ for packet in stream.encode(frame):
+ container.mux(packet)
+ for packet in stream.encode():
+ container.mux(packet)
+ decoded = _decode(path)
+ if len(decoded) != len(frames) or decoded[0].shape != frames[0].shape:
+ raise RuntimeError(f"generated video failed decode probe: {relative}")
+ self._record(
+ relative,
+ "video",
+ seed,
+ width=width,
+ height=height,
+ num_frames=len(frames),
+ fps=FPS,
+ codec="ffv1",
+ )
+ return _Asset("video", relative, fps=FPS)
+
+ def audio(self, relative: str, samples: np.ndarray, rate: int, seed: int) -> _Asset:
+ path = self._new_path(relative)
+ with wave.open(str(path), "wb") as output:
+ output.setparams((2, 2, rate, samples.shape[0], "NONE", "not compressed"))
+ output.writeframes(samples.astype(" None:
+ for alias in RUNTIME_ALIASES:
+ alias_dir = self.root / "profiles" / alias
+ alias_dir.mkdir(parents=True)
+ _write_jsonl(alias_dir / "train.jsonl", records[alias])
+ _write_jsonl(
+ self.root / "provenance.jsonl", sorted(self.provenance, key=lambda x: x["path"])
+ )
+ manifest = {
+ "schema_version": 1,
+ "repository_id": self.repo_id,
+ "flow_factory_schema_version": 2,
+ "supervision_type": self.supervision,
+ "records_per_alias": self.records_per_alias,
+ "runtime_aliases": list(RUNTIME_ALIASES),
+ "license": CC0,
+ "condition_endpoint_check": {
+ "metric": "decoded_rgb_max_absolute_difference",
+ "tolerance": 0,
+ },
+ "generator": {
+ "script": "dataset/offline_smoke/build_mini.py",
+ "script_sha256": self.script_sha,
+ "seed": self.seed,
+ },
+ }
+ _write_json(self.root / "dataset_manifest.json", manifest)
+ (self.root / "README.md").write_text(_card(self), encoding="utf-8")
+ (self.root / "LICENSE").write_text(_cc0_notice(), encoding="utf-8")
+
+ def metadata(self, alias: str, index: int, seed: int) -> dict[str, Any]:
+ value: dict[str, Any] = {
+ "sample_id": f"{alias}-{index:04d}",
+ "profile": GPU_ALIAS_TO_PROFILE[alias].name,
+ "gpu_alias": alias,
+ "usage_tier": "smoke_only",
+ "source": {
+ "origin": "flow_factory_procedural",
+ "license": CC0,
+ "generator_seed": seed,
+ "generator_script_sha256": self.script_sha,
+ },
+ }
+ if self.supervision == "preference":
+ value.update(
+ preference_origin="deterministic_corruption",
+ semantic_preference_claim=False,
+ )
+ return value
+
+ def _new_path(self, relative: str) -> Path:
+ path = self.root / relative
+ if not path.resolve().is_relative_to(self.root.resolve()) or path.exists():
+ raise ValueError(f"invalid or duplicate asset path: {relative!r}")
+ path.parent.mkdir(parents=True, exist_ok=True)
+ return path
+
+ def _record(self, relative: str, media_type: str, seed: int, **media: Any) -> None:
+ path = self.root / relative
+ self.provenance.append(
+ {
+ "path": relative,
+ "type": media_type,
+ "sha256": _sha256(path),
+ "size_bytes": path.stat().st_size,
+ "origin": "flow_factory_procedural",
+ "license": CC0,
+ "generator_seed": seed,
+ "media": media,
+ }
+ )
+
+
+def build(
+ staging_root: Path,
+ seed: int,
+ replace: bool,
+ records_per_alias: int = RECORDS_PER_ALIAS,
+) -> tuple[Path, Path]:
+ """Build independent SFT and offline-DPO staging trees atomically.
+
+ Args:
+ staging_root: Directory that receives the two repository trees.
+ seed: Base seed for deterministic procedural records.
+ replace: Whether to replace the two exact destination directories.
+ records_per_alias: Number of records generated for every runtime alias.
+
+ Returns:
+ Paths to the SFT and offline-DPO staging trees, in that order.
+
+ Raises:
+ ValueError: If ``records_per_alias`` is not positive.
+ FileExistsError: If a destination exists and ``replace`` is false.
+ """
+ if records_per_alias < 1:
+ raise ValueError("records_per_alias must be positive")
+ staging_root.mkdir(parents=True, exist_ok=True)
+ destinations = tuple(
+ staging_root / repo_id.rsplit("/", 1)[1] for repo_id in (SFT_REPO_ID, OFFLINE_DPO_REPO_ID)
+ )
+ if not replace and any(path.exists() for path in destinations):
+ raise FileExistsError("staging repo exists; pass --replace for the two exact targets")
+ temporary = Path(tempfile.mkdtemp(prefix=".build-mini-", dir=staging_root))
+ try:
+ for repo_id, supervision, destination in zip(
+ (SFT_REPO_ID, OFFLINE_DPO_REPO_ID),
+ ("demonstration", "preference"),
+ destinations,
+ ):
+ writer = _Writer(
+ temporary / destination.name,
+ repo_id,
+ supervision,
+ seed,
+ records_per_alias,
+ )
+ writer.root.mkdir()
+ records = {alias: [] for alias in RUNTIME_ALIASES}
+ for index in range(records_per_alias):
+ row_seed = _seed(seed, supervision, index)
+ pools = _assets(writer, index, row_seed)
+ for alias in RUNTIME_ALIASES:
+ records[alias].append(_record(writer, alias, index, row_seed, pools))
+ writer.finish(records)
+ if replace:
+ for destination in destinations:
+ if destination.exists():
+ shutil.rmtree(destination)
+ for destination in destinations:
+ os.replace(temporary / destination.name, destination)
+ finally:
+ shutil.rmtree(temporary, ignore_errors=True)
+ return destinations
+
+
+def _assets(writer: _Writer, index: int, seed: int) -> dict[str, dict[str, _Asset]]:
+ """Generate shared image, Wan, LTX, and H3 pools for one logical row."""
+ first = (SHAPES[index % len(SHAPES)], COLORS[index % len(COLORS)])
+ second = (SHAPES[(index + 2) % len(SHAPES)], COLORS[(index + 2) % len(COLORS)])
+ prefix = f"media/{index:04d}"
+ image_pool = {
+ "input": writer.image(f"{prefix}/image/input.png", _scene(first, "center"), seed),
+ "ref1": writer.image(f"{prefix}/image/ref1.png", _scene(first, "center"), seed),
+ "ref2": writer.image(f"{prefix}/image/ref2.png", _scene(second, "center"), seed),
+ "t2i_chosen": writer.image(f"{prefix}/image/t2i_chosen.png", _scene(first, "center"), seed),
+ "i2i_chosen": writer.image(
+ f"{prefix}/image/i2i_chosen.png", _scene(second, "upper_left"), seed
+ ),
+ "multi_chosen": writer.image(
+ f"{prefix}/image/multi_chosen.png", _composition(first, second, False), seed
+ ),
+ }
+ if writer.supervision == "preference":
+ image_pool.update(
+ t2i_rejected=writer.image(
+ f"{prefix}/image/t2i_rejected.png", _scene(second, "lower_right"), seed
+ ),
+ i2i_rejected=writer.image(
+ f"{prefix}/image/i2i_rejected.png", _scene(first, "lower_right"), seed
+ ),
+ multi_rejected=writer.image(
+ f"{prefix}/image/multi_rejected.png",
+ _composition(first, second, True),
+ seed,
+ ),
+ )
+ pools = {"image": image_pool}
+ for family, alias in (("wan", "wan-t2v"), ("ltx", "ltx2-t2av"), ("h3", "h3-t2va")):
+ geometry = _geometry(alias)
+ frames = _motion(geometry.width, geometry.height, geometry.num_frames, first, second)
+ pool: dict[str, _Asset] = {
+ "chosen_video": writer.video(f"{prefix}/{family}/chosen.avi", frames, seed)
+ }
+ decoded = _decode(writer.root / pool["chosen_video"].path)
+ pool["first"] = writer.image(
+ f"{prefix}/{family}/first.png", Image.fromarray(decoded[0]), seed
+ )
+ pool["last"] = writer.image(
+ f"{prefix}/{family}/last.png", Image.fromarray(decoded[-1]), seed
+ )
+ if geometry.sample_rate is not None:
+ samples = _audio(geometry.sample_rate, geometry.num_frames / geometry.frame_rate, index)
+ pool["chosen_audio"] = writer.audio(
+ f"{prefix}/{family}/chosen.wav", samples, geometry.sample_rate, seed
+ )
+ if writer.supervision == "preference":
+ rejected_frames = _corrupt_frames(frames, second[1][1])
+ pool["rejected_video"] = writer.video(
+ f"{prefix}/{family}/rejected.avi", rejected_frames, seed
+ )
+ rejected_decoded = _decode(writer.root / pool["rejected_video"].path)
+ if not (
+ np.array_equal(decoded[0], rejected_decoded[0])
+ and np.array_equal(decoded[-1], rejected_decoded[-1])
+ ):
+ raise RuntimeError(f"{family} DPO corruption changed a decoded endpoint")
+ if geometry.sample_rate is not None:
+ pool["rejected_audio"] = writer.audio(
+ f"{prefix}/{family}/rejected.wav",
+ _corrupt_audio(samples),
+ geometry.sample_rate,
+ seed,
+ )
+ pools[family] = pool
+ return pools
+
+
+def _record(
+ writer: _Writer,
+ alias: str,
+ index: int,
+ seed: int,
+ pools: Mapping[str, Mapping[str, _Asset]],
+) -> dict[str, Any]:
+ profile = GPU_ALIAS_TO_PROFILE[alias]
+ first = (SHAPES[index % len(SHAPES)], COLORS[index % len(COLORS)][0])
+ second = (SHAPES[(index + 2) % len(SHAPES)], COLORS[(index + 2) % len(COLORS)][0])
+ prompt = f"A {first[1]} {first[0]} moves left to right beside a {second[1]} {second[0]}."
+ family = _family(alias)
+ input_media: list[dict[str, Any]] = []
+ if alias == "image-i2i":
+ prompt = f"Move the {first[1]} {first[0]} to the upper left and recolor it {second[1]}."
+ input_media = [_media(pools["image"]["input"])]
+ elif alias == "bagel-mri2i":
+ prompt = "Place reference 1 on the left and reference 2 on the right."
+ input_media = [_media(pools["image"]["ref1"]), _media(pools["image"]["ref2"])]
+ elif alias in ("wan-i2v-first", "ltx2-i2av"):
+ input_media = [_media(pools[family]["first"], slot="first_frame")]
+ elif alias == "wan-flf2v":
+ input_media = [
+ _media(pools[family]["first"], slot="first_frame"),
+ _media(pools[family]["last"], slot="last_frame"),
+ ]
+ elif alias == "h3-fl2va":
+ cases = (
+ [_media(pools[family]["first"], slot="first_frame")],
+ [_media(pools[family]["last"], slot="last_frame")],
+ [
+ _media(pools[family]["first"], slot="first_frame"),
+ _media(pools[family]["last"], slot="last_frame"),
+ ],
+ )
+ input_media = cases[index % len(cases)]
+ elif alias == "h3-ref2va":
+ refs = [
+ _media(pools["wan"]["first"]),
+ _media(pools["wan"]["chosen_video"]),
+ _media(pools["h3"]["chosen_audio"]),
+ ]
+ offset = index % len(refs)
+ input_media = refs[offset:] + refs[:offset]
+
+ output_types = tuple(item.type.value for item in profile.contract.output_media.items)
+ candidates = _candidate_assets(alias, family, pools)
+ chosen = _candidate(output_types, candidates, "chosen")
+ if writer.supervision == "demonstration":
+ supervision: dict[str, Any] = {"type": "demonstration", "target": chosen}
+ else:
+ supervision = {
+ "type": "preference",
+ "chosen": chosen,
+ "rejected": _candidate(output_types, candidates, "rejected"),
+ }
+ return {
+ "schema_version": 2,
+ "input": {"prompt": prompt, "media": input_media},
+ "supervision": supervision,
+ "metadata": writer.metadata(alias, index, seed),
+ }
+
+
+def _candidate_assets(
+ alias: str,
+ family: str,
+ pools: Mapping[str, Mapping[str, _Asset]],
+) -> Mapping[str, Mapping[str, _Asset]]:
+ if alias == "sd35-t2i":
+ names = ("t2i_chosen", "t2i_rejected")
+ elif alias == "image-i2i":
+ names = ("i2i_chosen", "i2i_rejected")
+ elif alias == "bagel-mri2i":
+ names = ("multi_chosen", "multi_rejected")
+ else:
+ pool = pools[family]
+ return {
+ "chosen": {
+ kind: pool[f"chosen_{kind}"]
+ for kind in ("video", "audio")
+ if f"chosen_{kind}" in pool
+ },
+ "rejected": {
+ kind: pool[f"rejected_{kind}"]
+ for kind in ("video", "audio")
+ if f"rejected_{kind}" in pool
+ },
+ }
+ pool = pools["image"]
+ return {
+ "chosen": {"image": pool[names[0]]},
+ "rejected": {"image": pool[names[1]]} if names[1] in pool else {},
+ }
+
+
+def _candidate(
+ output_types: Sequence[str],
+ candidates: Mapping[str, Mapping[str, _Asset]],
+ side: str,
+) -> dict[str, Any]:
+ """Project any declared output sequence without assuming AV-only candidates."""
+ available = candidates[side]
+ media = [_media(available[media_type]) for media_type in output_types]
+ if tuple(item["type"] for item in media) != tuple(output_types):
+ raise RuntimeError("candidate media order diverged from the profile contract")
+ return {"media": media}
+
+
+def _geometry(alias: str) -> Any:
+ profile = GPU_ALIAS_TO_PROFILE[alias]
+ return next(case.geometry for case in profile.gpu_cases if case.alias == alias)
+
+
+def _family(alias: str) -> str:
+ if alias.startswith("wan-"):
+ return "wan"
+ if alias.startswith("ltx2-"):
+ return "ltx"
+ if alias.startswith("h3-"):
+ return "h3"
+ return "image"
+
+
+def _media(asset: _Asset, slot: str | None = None) -> dict[str, Any]:
+ value: dict[str, Any] = {"type": asset.type, "path": f"../../{asset.path}"}
+ if asset.fps is not None:
+ value["fps"] = asset.fps
+ if asset.sample_rate is not None:
+ value["sample_rate"] = asset.sample_rate
+ if slot is not None:
+ value["slot"] = slot
+ return value
+
+
+def _scene(item: tuple[str, tuple[str, tuple[int, int, int]]], position: str) -> Image.Image:
+ shape, (_, color) = item
+ image = Image.new("RGB", (256, 256), (235, 239, 242))
+ centers = {"center": (128, 128), "upper_left": (80, 80), "lower_right": (176, 176)}
+ _draw_shape(ImageDraw.Draw(image), shape, centers[position], 42, color)
+ return image
+
+
+def _composition(
+ first: tuple[str, tuple[str, tuple[int, int, int]]],
+ second: tuple[str, tuple[str, tuple[int, int, int]]],
+ swap: bool,
+) -> Image.Image:
+ image = Image.new("RGB", (256, 256), (235, 239, 242))
+ positions = ((76, 128), (180, 128)) if not swap else ((180, 128), (76, 128))
+ draw = ImageDraw.Draw(image)
+ _draw_shape(draw, first[0], positions[0], 36, first[1][1])
+ _draw_shape(draw, second[0], positions[1], 36, second[1][1])
+ return image
+
+
+def _motion(
+ width: int,
+ height: int,
+ count: int,
+ first: tuple[str, tuple[str, tuple[int, int, int]]],
+ second: tuple[str, tuple[str, tuple[int, int, int]]],
+) -> list[np.ndarray]:
+ frames = []
+ radius = max(min(width, height) // 12, 4)
+ for index in range(count):
+ phase = index / max(count - 1, 1)
+ image = Image.new("RGB", (width, height), (232, 238, 242))
+ draw = ImageDraw.Draw(image)
+ x = round(radius * 2 + phase * (width - radius * 4))
+ y = round(height * (0.55 + 0.1 * math.sin(phase * math.tau)))
+ _draw_shape(draw, first[0], (x, y), radius, first[1][1])
+ _draw_shape(draw, second[0], (width * 3 // 4, height // 3), radius, second[1][1])
+ frames.append(np.asarray(image, dtype=np.uint8).copy())
+ return frames
+
+
+def _draw_shape(
+ draw: ImageDraw.ImageDraw,
+ shape: str,
+ center: tuple[int, int],
+ radius: int,
+ color: tuple[int, int, int],
+) -> None:
+ x, y = center
+ if shape == "circle":
+ draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=color)
+ elif shape == "square":
+ draw.rectangle((x - radius, y - radius, x + radius, y + radius), fill=color)
+ elif shape == "triangle":
+ draw.polygon(
+ ((x, y - radius), (x - radius, y + radius), (x + radius, y + radius)), fill=color
+ )
+ elif shape == "diamond":
+ draw.polygon(
+ ((x, y - radius), (x - radius, y), (x, y + radius), (x + radius, y)), fill=color
+ )
+ else:
+ points = []
+ for point in range(10):
+ angle = -math.pi / 2 + point * math.pi / 5
+ distance = radius if point % 2 == 0 else radius * 0.45
+ points.append(
+ (round(x + math.cos(angle) * distance), round(y + math.sin(angle) * distance))
+ )
+ draw.polygon(points, fill=color)
+
+
+def _audio(rate: int, duration: float, index: int) -> np.ndarray:
+ timeline = np.arange(round(rate * duration), dtype=np.float64) / rate
+ frequency = 180 + 17 * (index % 7)
+ phase = math.tau * (frequency * timeline + 20 * timeline**2)
+ envelope = np.sin(np.pi * timeline / duration) ** 0.5
+ stereo = np.stack((np.sin(phase), np.sin(phase + math.pi / 3)), axis=-1)
+ return np.rint(stereo * envelope[:, None] * 9000).astype(np.int16)
+
+
+def _corrupt_frames(frames: Sequence[np.ndarray], accent: tuple[int, int, int]) -> list[np.ndarray]:
+ output = [frames[0].copy()]
+ for index, frame in enumerate(reversed(frames[1:-1]), start=1):
+ changed = frame.copy()
+ width = changed.shape[1]
+ x = (index * 7) % max(width - 2, 1)
+ changed[:, x : x + 2] = accent
+ output.append(changed)
+ output.append(frames[-1].copy())
+ return output
+
+
+def _corrupt_audio(samples: np.ndarray) -> np.ndarray:
+ output = np.roll(samples, max(samples.shape[0] // 9, 1), axis=0).copy()
+ output[:, 1] *= -1
+ return output
+
+
+def _decode(path: Path) -> list[np.ndarray]:
+ with av.open(str(path)) as container:
+ frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)]
+ if not frames:
+ raise RuntimeError(f"generated video decoded no frames: {path}")
+ return frames
+
+
+def _seed(base: int, supervision: str, index: int) -> int:
+ digest = hashlib.sha256(f"{base}:{supervision}:{index}".encode()).digest()
+ return int.from_bytes(digest[:8], "big")
+
+
+def _sha256(path: Path) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as source:
+ for chunk in iter(lambda: source.read(1024 * 1024), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _write_json(path: Path, value: Mapping[str, Any]) -> None:
+ path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+
+
+def _write_jsonl(path: Path, values: Iterable[Mapping[str, Any]]) -> None:
+ with path.open("w", encoding="utf-8") as output:
+ for value in values:
+ output.write(json.dumps(value, separators=(",", ":"), sort_keys=True) + "\n")
+
+
+def _card(writer: _Writer) -> str:
+ return f"""---
+license: cc0-1.0
+pretty_name: Flow-Factory {writer.supervision} Smoke Fixtures
+---
+
+# {writer.repo_id}
+
+Deterministic procedural correctness fixtures for Flow-Factory's strict V2
+offline schema. Each runtime alias under `profiles/` has {writer.records_per_alias} records and resolves
+its `../../media/...` paths inside this repository. Joint outputs are ordered
+`[video, audio]`. Preference pairs are synthetic smoke-only corruptions, not
+human labels. See `dataset_manifest.json` and `provenance.jsonl` for identities.
+"""
+
+
+def _cc0_notice() -> str:
+ return """CC0 1.0 Universal
+
+To the extent possible under law, the authors waive copyright and related
+rights in these generated dataset assets.
+https://creativecommons.org/publicdomain/zero/1.0/legalcode
+"""
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+ """Build both public dataset staging trees from command-line arguments.
+
+ Args:
+ argv: Optional argument sequence. Uses ``sys.argv`` when omitted.
+
+ Returns:
+ Process exit code zero after a successful build.
+ """
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument(
+ "--staging-root",
+ type=Path,
+ default=Path(__file__).resolve().parent / "_staging",
+ )
+ parser.add_argument("--seed", type=int, default=DEFAULT_SEED)
+ parser.add_argument("--records-per-alias", type=int, default=RECORDS_PER_ALIAS)
+ parser.add_argument("--replace", action="store_true")
+ args = parser.parse_args(argv)
+ paths = build(
+ args.staging_root.expanduser().resolve(),
+ args.seed,
+ args.replace,
+ args.records_per_alias,
+ )
+ print(json.dumps({"sft": str(paths[0]), "offline_dpo": str(paths[1])}, sort_keys=True))
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/dataset/offline_smoke/datasets.lock.json b/dataset/offline_smoke/datasets.lock.json
new file mode 100644
index 000000000..a1f3d4575
--- /dev/null
+++ b/dataset/offline_smoke/datasets.lock.json
@@ -0,0 +1,15 @@
+{
+ "schema_version": 1,
+ "datasets": {
+ "offline-dpo": {
+ "repo_id": "Jayce-Ping/Flow-Factory-Offline-DPO-Smoke",
+ "revision": "b6ec289238eefdd6c855b7fb8f87ccc80ee3040f",
+ "supervision_type": "preference"
+ },
+ "sft": {
+ "repo_id": "Jayce-Ping/Flow-Factory-SFT-Smoke",
+ "revision": "50e7f0e897ed2ce4160b14340fc8a982fe6c6919",
+ "supervision_type": "demonstration"
+ }
+ }
+}
diff --git a/dataset/offline_smoke/prepare.py b/dataset/offline_smoke/prepare.py
new file mode 100644
index 000000000..95dbfe060
--- /dev/null
+++ b/dataset/offline_smoke/prepare.py
@@ -0,0 +1,319 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Download and atomically materialize one pinned offline-smoke profile."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import re
+import shutil
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Dict, Literal, Mapping, Sequence
+
+from huggingface_hub import snapshot_download
+
+from flow_factory.data_utils.schema import DatasetRecordV2
+
+from .profiles import DATASET_REPO_IDS, GPU_ALIAS_TO_PROFILE
+from .validate import validate_dataset
+
+Algorithm = Literal["sft", "offline-dpo"]
+SnapshotDownload = Callable[..., str]
+DEFAULT_LOCK_PATH = Path(__file__).with_name("datasets.lock.json")
+DEFAULT_OUTPUT_ROOT = Path("dataset/_prepared_offline_smoke")
+PENDING_REVISION = "PENDING_INITIAL_UPLOAD"
+_COMMIT_SHA = re.compile(r"^[0-9a-f]{40}$")
+
+
+@dataclass(frozen=True, slots=True)
+class DatasetLock:
+ """One immutable Hub dataset selection."""
+
+ repo_id: str
+ revision: str
+ supervision_type: str
+
+
+def load_dataset_lock(
+ algorithm: Algorithm, lock_path: str | Path = DEFAULT_LOCK_PATH
+) -> DatasetLock:
+ """Load and verify one independent repository lock.
+
+ Args:
+ algorithm: Offline supervision family to resolve.
+ lock_path: JSON lock file containing immutable Hub revisions.
+
+ Returns:
+ Validated repository selection for the requested algorithm.
+
+ Raises:
+ ValueError: If the algorithm, lock schema, repository, supervision, or revision is invalid.
+ """
+ if algorithm not in DATASET_REPO_IDS:
+ raise ValueError(f"unsupported offline smoke algorithm: {algorithm!r}")
+ path = Path(lock_path).expanduser()
+ with path.open(encoding="utf-8") as handle:
+ payload = json.load(handle)
+ if payload.get("schema_version") != 1 or not isinstance(payload.get("datasets"), dict):
+ raise ValueError(f"invalid offline smoke lock schema: {path}")
+ entry = payload["datasets"].get(algorithm)
+ if not isinstance(entry, dict):
+ raise ValueError(f"offline smoke lock has no {algorithm!r} entry: {path}")
+ lock = DatasetLock(
+ repo_id=entry.get("repo_id"),
+ revision=entry.get("revision"),
+ supervision_type=entry.get("supervision_type"),
+ )
+ expected_supervision = "demonstration" if algorithm == "sft" else "preference"
+ if lock.repo_id != DATASET_REPO_IDS[algorithm]:
+ raise ValueError(
+ f"locked repo_id {lock.repo_id!r} disagrees with profiles.py "
+ f"{DATASET_REPO_IDS[algorithm]!r}"
+ )
+ if lock.supervision_type != expected_supervision:
+ raise ValueError(
+ f"locked supervision {lock.supervision_type!r} must be {expected_supervision!r}"
+ )
+ if not isinstance(lock.revision, str) or not _COMMIT_SHA.fullmatch(lock.revision):
+ detail = (
+ "initial publication is still pending"
+ if lock.revision == PENDING_REVISION
+ else "not a commit SHA"
+ )
+ raise ValueError(
+ f"offline smoke revision for {algorithm!r} is {detail}: {lock.revision!r}; "
+ "publish the dataset and pin its 40-character commit SHA"
+ )
+ return lock
+
+
+def prepare_dataset(
+ *,
+ algorithm: Algorithm,
+ profile_name: str,
+ world_size: int,
+ output_root: str | Path = DEFAULT_OUTPUT_ROOT,
+ per_device_batch_size: int = 1,
+ batches_per_rank: int = 2,
+ allow_repeat: bool = False,
+ offline: bool = False,
+ lock_path: str | Path = DEFAULT_LOCK_PATH,
+ download_fn: SnapshotDownload | None = None,
+) -> Path:
+ """Materialize exactly the records required by one distributed smoke run.
+
+ Args:
+ algorithm: ``sft`` or ``offline-dpo`` supervision family.
+ profile_name: Model-specific runtime alias published by the dataset repository.
+ world_size: Number of distributed ranks in the target run.
+ output_root: Root directory for ready-to-use datasets.
+ per_device_batch_size: Number of records consumed by each rank per batch.
+ batches_per_rank: Exact number of rank-local batches to materialize.
+ allow_repeat: Whether to cycle the base fixture when more rows are requested.
+ offline: Whether Hub access must use already downloaded local files only.
+ lock_path: JSON lock file containing immutable Hub revisions.
+ download_fn: Optional snapshot downloader used by tests.
+
+ Returns:
+ Path to the atomically published, self-contained profile directory.
+
+ Raises:
+ FileExistsError: If the requested output profile already exists.
+ ValueError: If sizing, lock data, source records, or media violate their contracts.
+ """
+ for name, value in (
+ ("world_size", world_size),
+ ("per_device_batch_size", per_device_batch_size),
+ ("batches_per_rank", batches_per_rank),
+ ):
+ _require_positive_int(value, name)
+ if profile_name not in GPU_ALIAS_TO_PROFILE:
+ aliases = ", ".join(GPU_ALIAS_TO_PROFILE)
+ raise ValueError(
+ f"unknown published runtime alias {profile_name!r}; choose one of: {aliases}. "
+ "Canonical task profile names describe contracts but are not Hub directories."
+ )
+ source_profile = profile_name
+ requested_rows = world_size * per_device_batch_size * batches_per_rank
+ lock = load_dataset_lock(algorithm, lock_path)
+ root = Path(output_root).expanduser().resolve()
+ target = root / algorithm / source_profile
+ if target.exists():
+ raise FileExistsError(f"offline smoke output already exists: {target}")
+ target.parent.mkdir(parents=True, exist_ok=True)
+
+ download = snapshot_download if download_fn is None else download_fn
+ with tempfile.TemporaryDirectory(prefix=".offline-smoke-download-", dir=root) as download_dir:
+ snapshot = Path(
+ download(
+ repo_id=lock.repo_id,
+ repo_type="dataset",
+ revision=lock.revision,
+ local_dir=download_dir,
+ local_files_only=offline,
+ )
+ ).resolve()
+ source_manifest = snapshot / "profiles" / source_profile / "train.jsonl"
+ rows = _load_rows(source_manifest, supervision_type=lock.supervision_type)
+ if requested_rows > len(rows) and not allow_repeat:
+ raise ValueError(
+ f"profile {source_profile!r} has {len(rows)} base rows but {requested_rows} "
+ "are required; pass --allow-repeat to cycle smoke-only records explicitly"
+ )
+ selected = [rows[index % len(rows)] for index in range(requested_rows)]
+ staging = Path(tempfile.mkdtemp(prefix=f".{source_profile}-", dir=target.parent))
+ try:
+ with (staging / "train.jsonl").open("w", encoding="utf-8") as handle:
+ for row in selected:
+ rewritten = _materialize_row(row, source_manifest.parent, snapshot, staging)
+ handle.write(
+ json.dumps(
+ rewritten, ensure_ascii=False, separators=(",", ":"), sort_keys=True
+ )
+ + "\n"
+ )
+ summary = validate_dataset(
+ staging,
+ algorithm=algorithm,
+ profile_name=profile_name,
+ expected_rows=requested_rows,
+ )
+ materialization = {
+ "algorithm_profile": profile_name,
+ "allow_repeat": allow_repeat,
+ "repo_id": lock.repo_id,
+ "revision": lock.revision,
+ "row_count": requested_rows,
+ "source_profile": source_profile,
+ "validation": summary,
+ }
+ (staging / "materialization.json").write_text(
+ json.dumps(materialization, indent=2, ensure_ascii=False, sort_keys=True) + "\n",
+ encoding="utf-8",
+ )
+ os.replace(staging, target)
+ finally:
+ if staging.exists():
+ shutil.rmtree(staging)
+ return target
+
+
+def _load_rows(path: Path, *, supervision_type: str) -> list[DatasetRecordV2]:
+ if not path.is_file():
+ raise FileNotFoundError(f"pinned snapshot has no ready-to-use profile manifest: {path}")
+ rows = []
+ with path.open(encoding="utf-8") as handle:
+ for line_number, line in enumerate(handle, start=1):
+ if not line.strip():
+ raise ValueError(f"blank JSONL row at {path}:{line_number}")
+ parsed = DatasetRecordV2.model_validate_json(line)
+ if parsed.supervision.type != supervision_type:
+ raise ValueError(
+ f"{path}:{line_number} has {parsed.supervision.type!r} supervision, "
+ f"expected {supervision_type!r}"
+ )
+ rows.append(parsed)
+ if not rows:
+ raise ValueError(f"offline smoke profile manifest is empty: {path}")
+ return rows
+
+
+def _materialize_row(
+ record: DatasetRecordV2,
+ manifest_dir: Path,
+ snapshot_root: Path,
+ staging: Path,
+) -> Dict[str, Any]:
+ payload = record.model_dump(mode="json", exclude_none=True)
+ for media in _media_dicts(payload):
+ raw_path = Path(media["path"])
+ if raw_path.is_absolute():
+ raise ValueError(f"Hub smoke manifests cannot use absolute media paths: {raw_path}")
+ source = (manifest_dir / raw_path).resolve(strict=True)
+ if not source.is_relative_to(snapshot_root) or not source.is_file():
+ raise ValueError(f"Hub smoke media path escapes its pinned snapshot: {raw_path}")
+ relative = source.relative_to(snapshot_root)
+ if not relative.parts or relative.parts[0] != "media":
+ raise ValueError(f"Hub smoke media must live under repo-root media/: {raw_path}")
+ destination = staging / relative
+ if not destination.exists():
+ destination.parent.mkdir(parents=True, exist_ok=True)
+ try:
+ os.link(source, destination)
+ except OSError:
+ shutil.copy2(source, destination)
+ media["path"] = relative.as_posix()
+ return payload
+
+
+def _media_dicts(payload: Mapping[str, Any]):
+ yield from payload["input"].get("media", ())
+ supervision = payload["supervision"]
+ for candidate_name in ("target", "chosen", "rejected"):
+ candidate = supervision.get(candidate_name)
+ if candidate is not None:
+ yield from candidate["media"]
+
+
+def _require_positive_int(value: object, name: str) -> None:
+ if isinstance(value, bool) or not isinstance(value, int) or value < 1:
+ raise ValueError(f"{name} must be an integer >= 1, got {value!r}")
+
+
+def _build_parser() -> argparse.ArgumentParser:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--algorithm", choices=tuple(DATASET_REPO_IDS), required=True)
+ parser.add_argument("--profile", choices=tuple(GPU_ALIAS_TO_PROFILE), required=True)
+ parser.add_argument("--world-size", type=int, required=True)
+ parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT)
+ parser.add_argument("--per-device-batch-size", type=int, default=1)
+ parser.add_argument("--batches-per-rank", type=int, default=2)
+ parser.add_argument("--allow-repeat", action="store_true")
+ parser.add_argument("--offline", action="store_true")
+ parser.add_argument("--lock-path", type=Path, default=DEFAULT_LOCK_PATH)
+ return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+ """Prepare one pinned dataset profile from command-line arguments.
+
+ Args:
+ argv: Optional argument sequence. Uses ``sys.argv`` when omitted.
+
+ Returns:
+ Process exit code zero after successful materialization.
+ """
+ args = _build_parser().parse_args(argv)
+ path = prepare_dataset(
+ algorithm=args.algorithm,
+ profile_name=args.profile,
+ world_size=args.world_size,
+ output_root=args.output_root,
+ per_device_batch_size=args.per_device_batch_size,
+ batches_per_rank=args.batches_per_rank,
+ allow_repeat=args.allow_repeat,
+ offline=args.offline,
+ lock_path=args.lock_path,
+ )
+ print(path)
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/dataset/offline_smoke/profiles.py b/dataset/offline_smoke/profiles.py
new file mode 100644
index 000000000..d273719ee
--- /dev/null
+++ b/dataset/offline_smoke/profiles.py
@@ -0,0 +1,317 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Canonical task catalog for the public SFT and offline-DPO smoke datasets.
+
+Profiles reuse Flow-Factory's official pipeline contract rather than defining a
+second dataset-only type system. Model identifiers remain strings, so this
+module never imports model adapters such as Bagel or their optional runtimes.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import MappingProxyType
+from typing import Mapping
+
+from flow_factory.contracts import (
+ InputMediaBinding,
+ InputMediaOrder,
+ InputMediaRule,
+ MediaFormat,
+ MediaType,
+ NegativePromptPolicy,
+ PipelineIOContract,
+ RateRequirement,
+)
+from flow_factory.models.pipeline_contracts import (
+ IMAGE_FORMAT,
+ VIDEO_FORMAT_OPTIONAL_FPS,
+ audio_video_output_contract,
+ image_output_contract,
+ video_output_contract,
+)
+
+SFT_REPO_ID = "Jayce-Ping/Flow-Factory-SFT-Smoke"
+OFFLINE_DPO_REPO_ID = "Jayce-Ping/Flow-Factory-Offline-DPO-Smoke"
+DATASET_REPO_IDS: Mapping[str, str] = MappingProxyType(
+ {"sft": SFT_REPO_ID, "offline-dpo": OFFLINE_DPO_REPO_ID}
+)
+
+
+@dataclass(frozen=True, slots=True)
+class SmokeGeometry:
+ """Small real-weight geometry used by a dataset or GPU variant."""
+
+ height: int
+ width: int
+ num_frames: int | None = None
+ frame_rate: float | None = None
+ sample_rate: int | None = None
+ num_inference_steps: int = 2
+
+ @property
+ def duration_seconds(self) -> float | None:
+ """Return the declared video duration when a frame clock is available.
+
+ Returns:
+ Duration in seconds, or ``None`` for image-only geometry.
+ """
+ if self.num_frames is None or self.frame_rate is None:
+ return None
+ return self.num_frames / self.frame_rate
+
+
+@dataclass(frozen=True, slots=True)
+class GPUSmokeCase:
+ """One model-specific alias and geometry in the GPU handoff catalog."""
+
+ alias: str
+ model_type: str
+ checkpoint: str
+ geometry: SmokeGeometry
+ main_matrix: bool = True
+
+
+@dataclass(frozen=True, slots=True)
+class OfflineSmokeProfile:
+ """One canonical dataset profile shared by SFT and offline DPO."""
+
+ name: str
+ compatible_model_types: tuple[str, ...]
+ contract: PipelineIOContract
+ default_geometry: SmokeGeometry
+ gpu_cases: tuple[GPUSmokeCase, ...] = ()
+
+ @property
+ def profile_id(self) -> str:
+ """Return the stable canonical profile identifier.
+
+ Returns:
+ Canonical profile name.
+ """
+ return self.name
+
+ @property
+ def gpu_aliases(self) -> tuple[str, ...]:
+ """Return model-specific aliases that use this task profile.
+
+ Returns:
+ Ordered runtime alias tuple.
+ """
+ return tuple(case.alias for case in self.gpu_cases)
+
+
+def output_media_types(contract: PipelineIOContract) -> tuple[str, ...]:
+ """Return the exact output sequence without assuming a modality family.
+
+ Args:
+ contract: Pipeline contract whose output sequence is projected.
+
+ Returns:
+ Ordered public media type names.
+ """
+ return tuple(item.type.value for item in contract.output_media.items)
+
+
+_NO_NEGATIVE = NegativePromptPolicy.UNSUPPORTED
+_IMAGE = SmokeGeometry(256, 256)
+_WAN = SmokeGeometry(240, 240, 5, 24.0)
+_LTX = SmokeGeometry(128, 192, 9, 24.0, 16000)
+_H3 = SmokeGeometry(64, 96, 124, 24.0, 32000)
+_AUDIO_REFERENCE_FORMAT = MediaFormat(
+ type=MediaType.AUDIO,
+ fps=RateRequirement.NOT_APPLICABLE,
+ sample_rate=RateRequirement.OPTIONAL,
+)
+
+
+def _image_rule(
+ min_count: int,
+ max_count: int,
+ slots: tuple[str, ...] = (),
+ required_slots: tuple[str, ...] = (),
+) -> InputMediaRule:
+ return InputMediaRule(IMAGE_FORMAT, min_count, max_count, slots, required_slots)
+
+
+_T2I = image_output_contract(negative_prompt=_NO_NEGATIVE)
+_I2I = image_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_image_min_count=1,
+ input_image_max_count=1,
+)
+_MRI2I = image_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_image_min_count=2,
+ input_image_max_count=2,
+ input_order=InputMediaOrder.WITHIN_TYPE,
+)
+_T2V = video_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ output_fps=RateRequirement.REQUIRED,
+)
+_I2V = video_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_image_min_count=1,
+ input_image_max_count=1,
+ input_image_slots=("first_frame",),
+ required_input_image_slots=("first_frame",),
+ output_fps=RateRequirement.REQUIRED,
+)
+_FL2V = video_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_image_min_count=1,
+ input_image_max_count=2,
+ input_image_slots=("first_frame", "last_frame"),
+ required_input_image_slots=("first_frame",),
+ output_fps=RateRequirement.REQUIRED,
+)
+_T2AV = audio_video_output_contract(negative_prompt=_NO_NEGATIVE)
+_I2AV = audio_video_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_rules=(_image_rule(1, 1, ("first_frame",), ("first_frame",)),),
+)
+_FL2AV = audio_video_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_rules=(_image_rule(1, 2, ("first_frame", "last_frame")),),
+ input_order=InputMediaOrder.WITHIN_TYPE,
+)
+_REF2AV = audio_video_output_contract(
+ negative_prompt=_NO_NEGATIVE,
+ input_rules=(
+ InputMediaRule(format=IMAGE_FORMAT, min_count=0, max_count=9),
+ InputMediaRule(format=VIDEO_FORMAT_OPTIONAL_FPS, min_count=0, max_count=3),
+ InputMediaRule(format=_AUDIO_REFERENCE_FORMAT, min_count=0, max_count=3),
+ ),
+ input_binding=InputMediaBinding.ORDERED_REFERENCES,
+ input_order=InputMediaOrder.GLOBAL,
+ min_input_media_count=1,
+ max_input_media_count=12,
+ required_any_input_types=(MediaType.IMAGE, MediaType.VIDEO),
+)
+
+
+_CHECKPOINTS = {
+ "sd35-t2i": "stabilityai/stable-diffusion-3.5-medium",
+ "image-i2i": "black-forest-labs/FLUX.1-Kontext-dev",
+ "bagel-mri2i": "ByteDance-Seed/BAGEL-7B-MoT",
+ "wan-t2v": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+ "wan-i2v-first": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+ "wan-flf2v": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+ "ltx2-t2av": "Lightricks/LTX-2",
+ "ltx2-i2av": "Lightricks/LTX-2",
+ "h3-t2va": "MiniMaxAI/MiniMax-H3",
+ "h3-fl2va": "MiniMaxAI/MiniMax-H3",
+ "h3-ref2va": "MiniMaxAI/MiniMax-H3",
+}
+
+
+def _case(
+ alias: str,
+ model: str,
+ geometry: SmokeGeometry,
+ main: bool = True,
+) -> GPUSmokeCase:
+ return GPUSmokeCase(alias, model, _CHECKPOINTS[alias], geometry, main)
+
+
+_MODEL_TYPES = {
+ "text_to_image": "sd3-5 flux1 flux2 flux2-klein qwen-image z-image bagel sensenova",
+ "image_to_image": "flux1-kontext flux2 flux2-klein qwen-image-edit-plus bagel sensenova",
+ "multi_image_to_image": "flux2 flux2-klein qwen-image-edit-plus bagel sensenova",
+ "text_to_video": "wan2_t2v",
+ "first_frame_to_video": "wan2_i2v",
+ "first_last_frame_to_video": "wan2_i2v",
+ "text_to_audio_video": "ltx2_t2av minimax-h3-t2va",
+ "first_frame_to_audio_video": "ltx2_i2av minimax-h3-fl2va",
+ "first_last_frame_to_audio_video": "minimax-h3-fl2va",
+ "ordered_references_to_audio_video": "minimax-h3-ref2va",
+}
+_PROFILE_SPECS = (
+ ("text_to_image", _T2I, _IMAGE, (_case("sd35-t2i", "sd3-5", _IMAGE),)),
+ ("image_to_image", _I2I, _IMAGE, (_case("image-i2i", "flux1-kontext", _IMAGE, False),)),
+ ("multi_image_to_image", _MRI2I, _IMAGE, (_case("bagel-mri2i", "bagel", _IMAGE),)),
+ ("text_to_video", _T2V, _WAN, (_case("wan-t2v", "wan2_t2v", _WAN),)),
+ ("first_frame_to_video", _I2V, _WAN, (_case("wan-i2v-first", "wan2_i2v", _WAN),)),
+ ("first_last_frame_to_video", _FL2V, _WAN, (_case("wan-flf2v", "wan2_i2v", _WAN),)),
+ (
+ "text_to_audio_video",
+ _T2AV,
+ _LTX,
+ (_case("ltx2-t2av", "ltx2_t2av", _LTX), _case("h3-t2va", "minimax-h3-t2va", _H3)),
+ ),
+ (
+ "first_frame_to_audio_video",
+ _I2AV,
+ _LTX,
+ (_case("ltx2-i2av", "ltx2_i2av", _LTX),),
+ ),
+ (
+ "first_last_frame_to_audio_video",
+ _FL2AV,
+ _H3,
+ (_case("h3-fl2va", "minimax-h3-fl2va", _H3),),
+ ),
+ (
+ "ordered_references_to_audio_video",
+ _REF2AV,
+ _H3,
+ (_case("h3-ref2va", "minimax-h3-ref2va", _H3),),
+ ),
+)
+_PROFILES = tuple(
+ OfflineSmokeProfile(name, tuple(_MODEL_TYPES[name].split()), contract, geometry, cases)
+ for name, contract, geometry, cases in _PROFILE_SPECS
+)
+
+CANONICAL_PROFILES: Mapping[str, OfflineSmokeProfile] = MappingProxyType(
+ {profile.name: profile for profile in _PROFILES}
+)
+GPU_ALIAS_TO_PROFILE: Mapping[str, OfflineSmokeProfile] = MappingProxyType(
+ {case.alias: profile for profile in _PROFILES for case in profile.gpu_cases}
+)
+MAIN_GPU_ALIASES = tuple(
+ "sd35-t2i bagel-mri2i wan-t2v wan-i2v-first wan-flf2v "
+ "ltx2-t2av ltx2-i2av h3-t2va h3-fl2va h3-ref2va".split()
+)
+SUPPLEMENTAL_GPU_ALIASES = ("image-i2i",)
+
+
+def get_profile(name_or_gpu_alias: str) -> OfflineSmokeProfile:
+ """Resolve a canonical profile name or a model-specific GPU alias.
+
+ Args:
+ name_or_gpu_alias: Canonical profile name or runtime alias.
+
+ Returns:
+ Matching immutable smoke profile.
+
+ Raises:
+ KeyError: If no canonical profile or runtime alias matches the value.
+ """
+ try:
+ return CANONICAL_PROFILES[name_or_gpu_alias]
+ except KeyError:
+ try:
+ return GPU_ALIAS_TO_PROFILE[name_or_gpu_alias]
+ except KeyError as error:
+ raise KeyError(f"unknown offline smoke profile {name_or_gpu_alias!r}") from error
+
+
+__all__ = tuple(
+ "CANONICAL_PROFILES DATASET_REPO_IDS GPUSmokeCase GPU_ALIAS_TO_PROFILE MAIN_GPU_ALIASES "
+ "OFFLINE_DPO_REPO_ID OfflineSmokeProfile SFT_REPO_ID SUPPLEMENTAL_GPU_ALIASES "
+ "SmokeGeometry get_profile output_media_types".split()
+)
diff --git a/dataset/offline_smoke/publish.py b/dataset/offline_smoke/publish.py
new file mode 100644
index 000000000..f8304b360
--- /dev/null
+++ b/dataset/offline_smoke/publish.py
@@ -0,0 +1,225 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Publish one validated offline-smoke staging tree to the Hugging Face Hub.
+
+This command is intentionally separate from dataset construction. Building and
+validating fixtures are local, reversible operations; publication creates public
+external state and therefore requires the explicit ``--confirm-public`` flag.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Literal, Sequence
+
+from huggingface_hub import HfApi
+
+from .profiles import OFFLINE_DPO_REPO_ID, SFT_REPO_ID
+
+Algorithm = Literal["sft", "offline-dpo"]
+_SUPERVISION_BY_ALGORITHM = {"sft": "demonstration", "offline-dpo": "preference"}
+
+
+def default_repo_id(algorithm: Algorithm) -> str:
+ """Return the public repository owned by one supervision family.
+
+ Args:
+ algorithm: Offline supervision family.
+
+ Returns:
+ Canonical public Hugging Face dataset repository ID.
+
+ Raises:
+ ValueError: If the algorithm is unsupported.
+ """
+ if algorithm == "sft":
+ return SFT_REPO_ID
+ if algorithm == "offline-dpo":
+ return OFFLINE_DPO_REPO_ID
+ raise ValueError(f"unsupported offline smoke algorithm: {algorithm!r}")
+
+
+def validate_staging_tree(
+ staging_dir: Path,
+ *,
+ algorithm: Algorithm,
+ destination_repo_id: str,
+) -> None:
+ """Reject mismatched or unsafe publication trees before any Hub mutation.
+
+ Args:
+ staging_dir: Local self-contained repository tree to publish.
+ algorithm: Supervision family selected by the publication command.
+ destination_repo_id: Hub dataset repository that would receive the tree.
+
+ Returns:
+ None.
+
+ Raises:
+ ValueError: If identity, supervision, structure, or paths are invalid.
+ """
+ staging_dir = staging_dir.resolve()
+ if not staging_dir.is_dir():
+ raise ValueError(f"staging directory does not exist: {staging_dir}")
+ for relative_path in (
+ "README.md",
+ "LICENSE",
+ "dataset_manifest.json",
+ "media",
+ "profiles",
+ "provenance.jsonl",
+ ):
+ if not (staging_dir / relative_path).exists():
+ raise ValueError(
+ f"staging directory is incomplete: missing {relative_path!r} under {staging_dir}"
+ )
+
+ with (staging_dir / "dataset_manifest.json").open(encoding="utf-8") as handle:
+ manifest = json.load(handle)
+ canonical_repo_id = default_repo_id(algorithm)
+ expected_supervision = _SUPERVISION_BY_ALGORITHM[algorithm]
+ if manifest.get("schema_version") != 1 or manifest.get("flow_factory_schema_version") != 2:
+ raise ValueError("staging dataset manifest has an unsupported schema version")
+ if manifest.get("supervision_type") != expected_supervision:
+ raise ValueError(
+ f"{algorithm!r} publication requires {expected_supervision!r} supervision, "
+ f"found {manifest.get('supervision_type')!r}"
+ )
+ if (
+ destination_repo_id == canonical_repo_id
+ and manifest.get("repository_id") != canonical_repo_id
+ ):
+ raise ValueError(
+ f"canonical destination {canonical_repo_id!r} requires a matching staging repository_id, "
+ f"found {manifest.get('repository_id')!r}"
+ )
+ runtime_aliases = manifest.get("runtime_aliases")
+ if (
+ not isinstance(runtime_aliases, list)
+ or not runtime_aliases
+ or any(not isinstance(alias, str) or not alias for alias in runtime_aliases)
+ or len(runtime_aliases) != len(set(runtime_aliases))
+ ):
+ raise ValueError("staging dataset manifest must declare unique non-empty runtime aliases")
+ profile_aliases = {path.name for path in (staging_dir / "profiles").iterdir() if path.is_dir()}
+ if profile_aliases != set(runtime_aliases):
+ raise ValueError("staging profile directories disagree with manifest runtime_aliases")
+
+ for path in staging_dir.rglob("*"):
+ if path.is_symlink():
+ raise ValueError(f"public staging trees cannot contain symlinks: {path}")
+ if path.is_file() and not path.resolve().is_relative_to(staging_dir):
+ raise ValueError(f"staging file escapes publication root: {path}")
+
+
+def publish_staging_tree(
+ *,
+ algorithm: Algorithm,
+ staging_dir: Path,
+ repo_id: str | None = None,
+ commit_message: str,
+ confirm_public: bool,
+) -> str:
+ """Create or update one public dataset repository.
+
+ Args:
+ algorithm: Offline supervision family being published.
+ staging_dir: Validated local repository tree.
+ repo_id: Optional noncanonical destination repository override.
+ commit_message: Hub commit message.
+ confirm_public: Explicit acknowledgement of public external state.
+
+ Returns:
+ Immutable Hub commit SHA returned by the upload.
+
+ Raises:
+ ValueError: If public confirmation or staging identity is invalid.
+ """
+ if not confirm_public:
+ raise ValueError("public dataset publication requires --confirm-public")
+ resolved_repo_id = default_repo_id(algorithm) if repo_id is None else repo_id
+ validate_staging_tree(
+ staging_dir,
+ algorithm=algorithm,
+ destination_repo_id=resolved_repo_id,
+ )
+ api = HfApi()
+ api.create_repo(
+ repo_id=resolved_repo_id,
+ repo_type="dataset",
+ private=False,
+ exist_ok=True,
+ )
+ commit = api.upload_folder(
+ repo_id=resolved_repo_id,
+ repo_type="dataset",
+ folder_path=str(staging_dir.resolve()),
+ commit_message=commit_message,
+ )
+ return commit.oid
+
+
+def _build_parser() -> argparse.ArgumentParser:
+ """Build the publication-only command-line interface."""
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--algorithm", choices=("sft", "offline-dpo"), required=True)
+ parser.add_argument("--staging-dir", type=Path, required=True)
+ parser.add_argument("--repo-id", default=None)
+ parser.add_argument(
+ "--commit-message",
+ default="Publish deterministic Flow-Factory offline smoke fixtures",
+ )
+ parser.add_argument(
+ "--confirm-public",
+ action="store_true",
+ help="Acknowledge that this command creates or updates a public Hub dataset.",
+ )
+ return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+ """Publish one staging tree and print a machine-readable result.
+
+ Args:
+ argv: Optional argument sequence. Uses ``sys.argv`` when omitted.
+
+ Returns:
+ Process exit code zero after a successful Hub commit.
+ """
+ args = _build_parser().parse_args(argv)
+ commit_sha = publish_staging_tree(
+ algorithm=args.algorithm,
+ staging_dir=args.staging_dir,
+ repo_id=args.repo_id,
+ commit_message=args.commit_message,
+ confirm_public=args.confirm_public,
+ )
+ print(
+ json.dumps(
+ {
+ "algorithm": args.algorithm,
+ "repo_id": args.repo_id or default_repo_id(args.algorithm),
+ "revision": commit_sha,
+ },
+ sort_keys=True,
+ )
+ )
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/dataset/offline_smoke/validate.py b/dataset/offline_smoke/validate.py
new file mode 100644
index 000000000..1842c8e63
--- /dev/null
+++ b/dataset/offline_smoke/validate.py
@@ -0,0 +1,233 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Thin validation boundary for materialized offline-smoke datasets.
+
+The model adapter remains authoritative for encoded geometry and condition/output
+semantics. This module only applies the existing public V2 and pipeline contracts,
+checks that media stay inside the dataset root, uniquely decodes referenced files,
+and rejects byte-identical offline-DPO candidates.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from pathlib import Path
+from typing import Any, Dict, Literal, Mapping, Tuple
+
+from flow_factory.contracts import (
+ validate_pipeline_model_input,
+ validate_pipeline_output_candidate,
+)
+from flow_factory.data_utils.offline_dataset import (
+ DEFAULT_MEDIA_DECODERS,
+ load_offline_manifest,
+)
+from flow_factory.data_utils.schema import (
+ DemonstrationSupervision,
+ MediaAsset,
+ NormalizedDatasetRecord,
+ NormalizedOutputCandidate,
+ PreferenceSupervision,
+)
+
+from .profiles import get_profile
+
+Algorithm = Literal["sft", "offline-dpo"]
+_SUPERVISION_BY_ALGORITHM: Mapping[Algorithm, str] = {
+ "sft": "demonstration",
+ "offline-dpo": "preference",
+}
+_HASH_CHUNK_SIZE = 1024 * 1024
+
+
+def validate_dataset(
+ dataset_dir: str | Path,
+ *,
+ algorithm: Algorithm,
+ profile_name: str,
+ expected_rows: int | None = None,
+) -> Dict[str, Any]:
+ """Validate one self-contained materialized dataset without loading a model.
+
+ Args:
+ dataset_dir: Directory containing ``train.jsonl`` and local media.
+ algorithm: Supervision family required by the trainer.
+ profile_name: Canonical task profile or model-specific runtime alias.
+ expected_rows: Optional exact row-count requirement.
+
+ Returns:
+ Validation summary with profile identity and unique media hashes.
+
+ Raises:
+ FileNotFoundError: If the dataset or a referenced media file is missing.
+ TypeError: If a record violates a typed schema or contract field.
+ ValueError: If records, paths, media, or preference arms are invalid.
+ """
+ supervision_type = _require_algorithm(algorithm)
+ if expected_rows is not None:
+ _require_positive_int(expected_rows, "expected_rows")
+ root = Path(dataset_dir).expanduser().resolve()
+ if not root.is_dir():
+ raise FileNotFoundError(f"offline smoke dataset directory does not exist: {root}")
+
+ profile = get_profile(profile_name)
+ records = load_offline_manifest(
+ root / "train.jsonl",
+ supervision_type=supervision_type,
+ dataset_dir=root,
+ )
+ if expected_rows is not None and len(records) != expected_rows:
+ raise ValueError(
+ f"offline smoke dataset must contain exactly {expected_rows} rows, "
+ f"found {len(records)}"
+ )
+
+ digest_cache: Dict[str, Tuple[Path, str]] = {}
+ decode_cache: Dict[Tuple[str, Path], object] = {}
+ for row_index, record in enumerate(records):
+ _validate_contract(record, profile.contract, row_index=row_index)
+ for asset in _record_media(record):
+ cached = digest_cache.get(asset.path)
+ if cached is None:
+ path = _require_contained_file(asset.path, root=root)
+ cached = (path, _sha256_file(path))
+ digest_cache[asset.path] = cached
+ else:
+ path = cached[0]
+ cache_key = (asset.type, path)
+ if cache_key not in decode_cache:
+ decode_cache[cache_key] = _decode_media(asset, context=f"row {row_index}")
+ if isinstance(record.supervision, PreferenceSupervision):
+ _require_distinct_candidates(
+ record.supervision.chosen,
+ record.supervision.rejected,
+ digest_cache=digest_cache,
+ context=f"row {row_index}",
+ )
+
+ return {
+ "algorithm": algorithm,
+ "requested_profile": profile_name,
+ "canonical_profile": profile.profile_id,
+ "row_count": len(records),
+ "media_file_count": len(digest_cache),
+ "media_sha256": {
+ path.relative_to(root).as_posix(): digest for path, digest in digest_cache.values()
+ },
+ }
+
+
+def _require_algorithm(algorithm: str) -> str:
+ if algorithm not in _SUPERVISION_BY_ALGORITHM:
+ raise ValueError(
+ f"unsupported offline smoke algorithm {algorithm!r}; "
+ f"expected one of {tuple(_SUPERVISION_BY_ALGORITHM)!r}"
+ )
+ return _SUPERVISION_BY_ALGORITHM[algorithm]
+
+
+def _require_positive_int(value: object, name: str) -> None:
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise TypeError(f"{name} must be an integer >= 1, got {type(value).__name__}: {value!r}")
+ if value < 1:
+ raise ValueError(f"{name} must be >= 1, got {value}")
+
+
+def _validate_contract(record, contract, *, row_index: int) -> None:
+ try:
+ validate_pipeline_model_input(record.model_input, contract)
+ for _, candidate in _record_candidates(record):
+ validate_pipeline_output_candidate(candidate.media, contract)
+ except (TypeError, ValueError) as exc:
+ raise type(exc)(
+ f"offline smoke row {row_index} violates its profile contract: {exc}"
+ ) from exc
+
+
+def _record_candidates(
+ record: NormalizedDatasetRecord,
+) -> Tuple[Tuple[str, NormalizedOutputCandidate], ...]:
+ supervision = record.supervision
+ if isinstance(supervision, DemonstrationSupervision):
+ return (("target", supervision.target),)
+ if isinstance(supervision, PreferenceSupervision):
+ return (("chosen", supervision.chosen), ("rejected", supervision.rejected))
+ raise TypeError(f"unsupported normalized supervision: {type(supervision).__name__}")
+
+
+def _record_media(record: NormalizedDatasetRecord):
+ yield from record.model_input.media
+ for _, candidate in _record_candidates(record):
+ yield from candidate.media
+
+
+def _require_contained_file(path_value: str, *, root: Path) -> Path:
+ path = Path(path_value)
+ try:
+ lexical_relative = path.relative_to(root)
+ except ValueError as exc:
+ raise ValueError(f"offline smoke media path escapes dataset root: {path}") from exc
+ cursor = root
+ for part in lexical_relative.parts:
+ cursor /= part
+ if cursor.is_symlink():
+ raise ValueError(f"offline smoke media paths cannot traverse symlinks: {cursor}")
+ try:
+ resolved = path.resolve(strict=True)
+ except FileNotFoundError as exc:
+ raise FileNotFoundError(f"offline smoke media file does not exist: {path}") from exc
+ if not resolved.is_relative_to(root):
+ raise ValueError(f"offline smoke media path escapes dataset root: {path}")
+ if not resolved.is_file():
+ raise ValueError(f"offline smoke media path is not a regular file: {path}")
+ return resolved
+
+
+def _decode_media(asset: MediaAsset, *, context: str) -> object:
+ decoder = DEFAULT_MEDIA_DECODERS.get(asset.type)
+ if decoder is None:
+ raise ValueError(f"{context} has no decoder for media type {asset.type!r}")
+ try:
+ return decoder(asset)
+ except (ImportError, OSError, TypeError, ValueError) as exc:
+ raise type(exc)(f"{context} failed to decode {asset.type} {asset.path!r}: {exc}") from exc
+
+
+def _sha256_file(path: Path) -> str:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ for chunk in iter(lambda: handle.read(_HASH_CHUNK_SIZE), b""):
+ digest.update(chunk)
+ return digest.hexdigest()
+
+
+def _require_distinct_candidates(
+ chosen: NormalizedOutputCandidate,
+ rejected: NormalizedOutputCandidate,
+ *,
+ digest_cache: Mapping[str, Tuple[Path, str]],
+ context: str,
+) -> None:
+ def signature(candidate: NormalizedOutputCandidate) -> Tuple[Tuple[str, str], ...]:
+ return tuple(
+ (
+ asset.type,
+ digest_cache[asset.path][1],
+ )
+ for asset in candidate.media
+ )
+
+ if signature(chosen) == signature(rejected):
+ raise ValueError(f"{context} has byte-identical chosen and rejected candidates")
diff --git a/guidance/datasets.md b/guidance/datasets.md
index 7a87075e0..81da9aaac 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -125,6 +125,41 @@ Tiny schema-complete fixtures and configs are available for
[SFT](../examples/sft/lora/sd3_5/default.yaml) and
[offline DPO](../examples/offline_dpo/lora/sd3_5/default.yaml).
+### Public offline smoke datasets
+
+The repository includes one builder for two independent, self-contained public mini datasets:
+
+- [Jayce-Ping/Flow-Factory-SFT-Smoke](https://huggingface.co/datasets/Jayce-Ping/Flow-Factory-SFT-Smoke)
+- [Jayce-Ping/Flow-Factory-Offline-DPO-Smoke](https://huggingface.co/datasets/Jayce-Ping/Flow-Factory-Offline-DPO-Smoke)
+
+Their runtime aliases cover the currently implemented image, video, and ordered `(video, audio)`
+output adapters. The catalog is expressed with the general `PipelineIOContract` output sequence;
+it does not encode audio-video as a special media type. A future audio-only or other ordered output
+can therefore use the same V2 schema and preparation path, although no fixture is published for a
+model family that does not exist in the framework today.
+
+Before a GPU smoke run, materialize exactly two rank-local batches from the pinned Hub revision:
+
+```bash
+python -m dataset.offline_smoke.prepare \
+ --algorithm sft \
+ --profile ltx2-t2av \
+ --world-size 8
+
+python -m dataset.offline_smoke.prepare \
+ --algorithm offline-dpo \
+ --profile h3-ref2va \
+ --world-size 8
+```
+
+The default output is `dataset/_prepared_offline_smoke///train.jsonl` with
+profile-local media. The preparation tool uses the locked immutable dataset revision, selects
+`world_size * per_device_batch_size * batches_per_rank` rows, and validates through the existing
+V2 and canonical task-profile contract boundaries. It does not preprocess target media or create
+latent caches.
+See [`dataset/offline_smoke`](../dataset/offline_smoke/README.md) for construction, provenance, and
+publication details.
+
Evaluation still uses generation acquisition, including when the trainer is SFT or offline DPO.
Consequently, a split enabled through `data.datasets[*].eval` must use one of the legacy
prompt/condition formats documented under [Common task formats](#common-task-formats), not a
diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md
index ede5e295f..92935b7a3 100644
--- a/guidance/gpu_validation.md
+++ b/guidance/gpu_validation.md
@@ -101,6 +101,31 @@ recorded in the resolved YAML; do not silently return to a large quality recipe.
| `h3-fl2va` | `examples/grpo/lora/minimax_h3_fl2va/default.yaml` plus H3 debug geometry | `resolution: [64, 96]`, `num_frames: 124`, `frame_rate: 24.0` | `num_inference_steps: 2` | The two offline records are explicit `first_frame`-only and `last_frame`-only cases. Cover both slots together in the additional variant gate. |
| `h3-ref2va` | `examples/grpo/lora/minimax_h3_ref2va/default.yaml` plus H3 debug geometry | `resolution: [64, 96]`, `num_frames: 124`, `frame_rate: 24.0` | `num_inference_steps: 2` | Preserve heterogeneous global reference order; include image, video, and audio references. |
+For the two offline algorithms, construct each ready-to-use alias from its independent pinned
+public dataset before launching distributed workers:
+
+```bash
+python -m dataset.offline_smoke.prepare \
+ --algorithm sft \
+ --profile "${MODE}" \
+ --world-size "${WORLD_SIZE}" \
+ --per-device-batch-size 1 \
+ --batches-per-rank 2
+
+python -m dataset.offline_smoke.prepare \
+ --algorithm offline-dpo \
+ --profile "${MODE}" \
+ --world-size "${WORLD_SIZE}" \
+ --per-device-batch-size 1 \
+ --batches-per-rank 2
+```
+
+Run the command once per shared filesystem rather than once per distributed rank. The SFT and DPO
+repositories have the same alias/input distribution but different supervision and self-contained
+media. `image-i2i` is an additional contract gate outside the 120-job main matrix. The catalog and
+materializer operate on arbitrary ordered output media sequences; currently published fixtures are
+limited to the image, video, and `(video, audio)` outputs implemented by registered adapters.
+
For GRPO set `group_size: 2`, `unique_sample_num_per_epoch: 1`, and
`gradient_accumulation_steps: auto`; this avoids a degenerate one-candidate
advantage while retaining one optimizer step per epoch. For TDM set
@@ -139,7 +164,7 @@ backends and all four algorithms.
| Wan T2V | `Wan2.1-T2V-14B-Diffusers`, `Wan2.2-TI2V-5B-Diffusers`, `Wan2.2-T2V-A14B-Diffusers` |
| Wan I2V first-only | `Wan2.1-I2V-14B-480P-Diffusers`, `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
| Wan first/last | `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
-| LTX2 T2AV and I2AV | `Lightricks/LTX-2.3` or its canonical official-Diffusers repository revision |
+| LTX2 T2AV and I2AV | `dg845/LTX-2.3-Diffusers` |
| MiniMax H3 | The same checkpoint is covered separately by T2VA, FL2VA, and Ref2VA inputs. Add an FL2VA first-plus-last fixture to complement the first-only/last-only main jobs. |
Wan2.2 TI2V-5B uses expanded timesteps. Official Diffusers ignores a supplied
diff --git a/tests/dataset/test_offline_smoke_build.py b/tests/dataset/test_offline_smoke_build.py
new file mode 100644
index 000000000..fe01f2263
--- /dev/null
+++ b/tests/dataset/test_offline_smoke_build.py
@@ -0,0 +1,243 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import hashlib
+import json
+import os
+from pathlib import Path
+from typing import Any, Iterator, Mapping
+
+import av
+import numpy as np
+import pytest
+from PIL import Image
+
+from dataset.offline_smoke.build_mini import RUNTIME_ALIASES, build
+from dataset.offline_smoke.profiles import GPU_ALIAS_TO_PROFILE
+from dataset.offline_smoke.validate import validate_dataset
+from flow_factory.contracts import (
+ validate_pipeline_model_input,
+ validate_pipeline_output_candidate,
+)
+from flow_factory.data_utils.offline_dataset import load_offline_manifest
+from flow_factory.models.registry import get_model_adapter_class
+
+ROWS = 3
+AV_ALIASES = (
+ "ltx2-t2av",
+ "ltx2-i2av",
+ "h3-t2va",
+ "h3-fl2va",
+ "h3-ref2va",
+)
+
+
+@pytest.fixture(scope="module")
+def built_repos(
+ tmp_path_factory: pytest.TempPathFactory,
+) -> tuple[Path, Path, Path]:
+ root = tmp_path_factory.mktemp("offline-smoke-build")
+ sft_root, dpo_root = build(
+ root / "staging",
+ seed=20260830,
+ replace=False,
+ records_per_alias=ROWS,
+ )
+ return root, sft_root, dpo_root
+
+
+def test_all_runtime_aliases_pass_public_validation(
+ built_repos: tuple[Path, Path, Path],
+) -> None:
+ root, sft_root, dpo_root = built_repos
+ for algorithm, repo, supervision_type in (
+ ("sft", sft_root, "demonstration"),
+ ("offline-dpo", dpo_root, "preference"),
+ ):
+ manifest = json.loads((repo / "dataset_manifest.json").read_text(encoding="utf-8"))
+ assert manifest["condition_endpoint_check"] == {
+ "metric": "decoded_rgb_max_absolute_difference",
+ "tolerance": 0,
+ }
+ profile_names = {path.name for path in (repo / "profiles").iterdir()}
+ assert profile_names == set(RUNTIME_ALIASES)
+ for alias in RUNTIME_ALIASES:
+ rows = _rows(repo, alias)
+ assert len(rows) == ROWS
+ assert {row["supervision"]["type"] for row in rows} == {supervision_type}
+ materialized = _materialize(repo, alias, root / "materialized" / algorithm / alias)
+ summary = validate_dataset(
+ materialized,
+ algorithm=algorithm,
+ profile_name=alias,
+ expected_rows=ROWS,
+ )
+ assert summary["row_count"] == ROWS
+
+
+def test_runtime_geometries_and_av_order(
+ built_repos: tuple[Path, Path, Path],
+) -> None:
+ _, sft_root, dpo_root = built_repos
+ expected_frames = {"wan-t2v": 5, "ltx2-t2av": 9, "h3-t2va": 124}
+ for alias, count in expected_frames.items():
+ row = _rows(sft_root, alias)[0]
+ video = _resolve(sft_root, alias, row["supervision"]["target"]["media"][0])
+ assert len(_decode_video(video)) == count
+
+ for repo, candidate_names in ((sft_root, ("target",)), (dpo_root, ("chosen", "rejected"))):
+ for alias in AV_ALIASES:
+ for row in _rows(repo, alias):
+ for name in candidate_names:
+ media = row["supervision"][name]["media"]
+ assert [item["type"] for item in media] == ["video", "audio"]
+
+
+def test_generated_rows_remain_legal_subsets_of_real_adapter_contracts(
+ built_repos: tuple[Path, Path, Path],
+) -> None:
+ _, sft_root, _ = built_repos
+ for alias in RUNTIME_ALIASES:
+ if alias == "bagel-mri2i":
+ # The Bagel contract is covered behind its optional-kernel seam in
+ # tests/models/test_bagel_output_codec.py.
+ continue
+ profile = GPU_ALIAS_TO_PROFILE[alias]
+ case = next(case for case in profile.gpu_cases if case.alias == alias)
+ adapter_type = get_model_adapter_class(case.model_type)
+ contract = adapter_type.pipeline_io_contract
+ manifest = sft_root / "profiles" / alias / "train.jsonl"
+ records = load_offline_manifest(
+ manifest,
+ supervision_type="demonstration",
+ dataset_dir=manifest.parent,
+ )
+ for record in records:
+ validate_pipeline_model_input(record.model_input, contract)
+ validate_pipeline_output_candidate(record.supervision.target.media, contract)
+
+
+def test_h3_sparse_frame_slots_and_dpo_endpoints(
+ built_repos: tuple[Path, Path, Path],
+) -> None:
+ _, _, dpo_root = built_repos
+ rows = _rows(dpo_root, "h3-fl2va")
+ assert [[item["slot"] for item in row["input"]["media"]] for row in rows] == [
+ ["first_frame"],
+ ["last_frame"],
+ ["first_frame", "last_frame"],
+ ]
+
+ for alias in ("wan-i2v-first", "wan-flf2v", "ltx2-i2av", "h3-fl2va"):
+ for row in _rows(dpo_root, alias):
+ chosen = _candidate_video(dpo_root, alias, row, "chosen")
+ rejected = _candidate_video(dpo_root, alias, row, "rejected")
+ for condition in row["input"]["media"]:
+ slot = condition.get("slot")
+ if slot not in {"first_frame", "last_frame"}:
+ continue
+ endpoint = 0 if slot == "first_frame" else -1
+ expected = np.asarray(
+ Image.open(_resolve(dpo_root, alias, condition)).convert("RGB")
+ )
+ assert np.array_equal(chosen[endpoint], expected)
+ assert np.array_equal(rejected[endpoint], expected)
+
+
+def test_dpo_arms_are_content_distinct(
+ built_repos: tuple[Path, Path, Path],
+) -> None:
+ _, _, dpo_root = built_repos
+ digest_cache: dict[Path, str] = {}
+ for alias in RUNTIME_ALIASES:
+ for row in _rows(dpo_root, alias):
+ chosen = _signature(dpo_root, alias, row["supervision"]["chosen"], digest_cache)
+ rejected = _signature(
+ dpo_root,
+ alias,
+ row["supervision"]["rejected"],
+ digest_cache,
+ )
+ assert chosen != rejected
+
+
+def _rows(repo: Path, alias: str) -> list[dict[str, Any]]:
+ path = repo / "profiles" / alias / "train.jsonl"
+ return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines()]
+
+
+def _all_media(row: Mapping[str, Any]) -> Iterator[dict[str, Any]]:
+ yield from row["input"].get("media", ())
+ for name in ("target", "chosen", "rejected"):
+ candidate = row["supervision"].get(name)
+ if candidate is not None:
+ yield from candidate["media"]
+
+
+def _resolve(repo: Path, alias: str, media: Mapping[str, Any]) -> Path:
+ return (repo / "profiles" / alias / media["path"]).resolve()
+
+
+def _materialize(repo: Path, alias: str, destination: Path) -> Path:
+ rows = _rows(repo, alias)
+ destination.mkdir(parents=True)
+ repo = repo.resolve()
+ for row in rows:
+ for media in _all_media(row):
+ source = _resolve(repo, alias, media)
+ relative = source.relative_to(repo)
+ target = destination / relative
+ target.parent.mkdir(parents=True, exist_ok=True)
+ if not target.exists():
+ os.link(source, target)
+ media["path"] = relative.as_posix()
+ (destination / "train.jsonl").write_text(
+ "".join(json.dumps(row, sort_keys=True) + "\n" for row in rows),
+ encoding="utf-8",
+ )
+ return destination
+
+
+def _decode_video(path: Path) -> list[np.ndarray]:
+ with av.open(str(path)) as container:
+ return [frame.to_ndarray(format="rgb24") for frame in container.decode(video=0)]
+
+
+def _candidate_video(
+ repo: Path,
+ alias: str,
+ row: Mapping[str, Any],
+ name: str,
+) -> list[np.ndarray]:
+ media = next(item for item in row["supervision"][name]["media"] if item["type"] == "video")
+ return _decode_video(_resolve(repo, alias, media))
+
+
+def _signature(
+ repo: Path,
+ alias: str,
+ candidate: Mapping[str, Any],
+ cache: dict[Path, str],
+) -> tuple[tuple[str, str], ...]:
+ values = []
+ for media in candidate["media"]:
+ path = _resolve(repo, alias, media)
+ if path not in cache:
+ digest = hashlib.sha256()
+ with path.open("rb") as handle:
+ for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+ digest.update(chunk)
+ cache[path] = digest.hexdigest()
+ values.append((media["type"], cache[path]))
+ return tuple(values)
diff --git a/tests/dataset/test_offline_smoke_prepare.py b/tests/dataset/test_offline_smoke_prepare.py
new file mode 100644
index 000000000..eac5e5011
--- /dev/null
+++ b/tests/dataset/test_offline_smoke_prepare.py
@@ -0,0 +1,260 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, Mapping
+
+import pytest
+from PIL import Image
+
+from dataset.offline_smoke.prepare import PENDING_REVISION, prepare_dataset
+from dataset.offline_smoke.profiles import DATASET_REPO_IDS
+from dataset.offline_smoke.validate import validate_dataset
+
+PINNED_SHA = "a" * 40
+
+
+def test_prepare_cli_supports_documented_module_execution() -> None:
+ repo_root = Path(__file__).resolve().parents[2]
+ result = subprocess.run(
+ [sys.executable, "-m", "dataset.offline_smoke.prepare", "--help"],
+ cwd=repo_root,
+ capture_output=True,
+ text=True,
+ )
+ assert result.returncode == 0, result.stderr
+ assert "--algorithm" in result.stdout
+
+
+def _write_lock(path: Path, *, pending: bool = False) -> None:
+ revision = PENDING_REVISION if pending else PINNED_SHA
+ path.write_text(
+ json.dumps(
+ {
+ "schema_version": 1,
+ "datasets": {
+ "sft": {
+ "repo_id": DATASET_REPO_IDS["sft"],
+ "revision": revision,
+ "supervision_type": "demonstration",
+ },
+ "offline-dpo": {
+ "repo_id": DATASET_REPO_IDS["offline-dpo"],
+ "revision": revision,
+ "supervision_type": "preference",
+ },
+ },
+ }
+ ),
+ encoding="utf-8",
+ )
+
+
+def _record(algorithm: str, *, target: str, rejected: str) -> Dict[str, Any]:
+ supervision: Dict[str, Any]
+ if algorithm == "sft":
+ supervision = {
+ "type": "demonstration",
+ "target": {"media": [{"type": "image", "path": target}]},
+ }
+ else:
+ supervision = {
+ "type": "preference",
+ "chosen": {"media": [{"type": "image", "path": target}]},
+ "rejected": {"media": [{"type": "image", "path": rejected}]},
+ }
+ return {
+ "schema_version": 2,
+ "input": {"prompt": "A deterministic smoke fixture.", "media": []},
+ "supervision": supervision,
+ "metadata": {"usage_tier": "smoke_only"},
+ }
+
+
+def _populate_snapshot(
+ local_dir: Path,
+ *,
+ algorithm: str,
+ profile_dir: str = "sd35-t2i",
+ row_count: int = 4,
+ target_path: str = "../../media/target.png",
+ identical_pair: bool = False,
+) -> None:
+ media_dir = local_dir / "media"
+ manifest_dir = local_dir / "profiles" / profile_dir
+ media_dir.mkdir(parents=True)
+ manifest_dir.mkdir(parents=True)
+ Image.new("RGB", (8, 8), color=(10, 20, 30)).save(media_dir / "target.png")
+ rejected_color = (10, 20, 30) if identical_pair else (30, 20, 10)
+ Image.new("RGB", (8, 8), color=rejected_color).save(media_dir / "rejected.png")
+ rows = [
+ _record(
+ algorithm,
+ target=target_path,
+ rejected="../../media/rejected.png",
+ )
+ for _ in range(row_count)
+ ]
+ (manifest_dir / "train.jsonl").write_text(
+ "".join(json.dumps(row) + "\n" for row in rows),
+ encoding="utf-8",
+ )
+
+
+@pytest.mark.parametrize("algorithm", ["sft", "offline-dpo"])
+def test_prepare_uses_independent_pinned_repo_and_exact_rank_geometry(
+ tmp_path: Path,
+ algorithm: str,
+) -> None:
+ lock_path = tmp_path / "lock.json"
+ _write_lock(lock_path)
+ calls = []
+
+ def download(**kwargs: Any) -> str:
+ calls.append(kwargs)
+ local_dir = Path(kwargs["local_dir"])
+ _populate_snapshot(local_dir, algorithm=algorithm)
+ return str(local_dir)
+
+ prepared = prepare_dataset(
+ algorithm=algorithm,
+ profile_name="sd35-t2i",
+ world_size=2,
+ output_root=tmp_path / "prepared",
+ offline=True,
+ lock_path=lock_path,
+ download_fn=download,
+ )
+
+ assert calls == [
+ {
+ "repo_id": DATASET_REPO_IDS[algorithm],
+ "repo_type": "dataset",
+ "revision": PINNED_SHA,
+ "local_dir": calls[0]["local_dir"],
+ "local_files_only": True,
+ }
+ ]
+ rows = [json.loads(line) for line in (prepared / "train.jsonl").read_text().splitlines()]
+ assert len(rows) == 4
+ assert all(".." not in media["path"] for row in rows for media in _all_media(row))
+ assert (prepared / "media" / "target.png").is_file()
+ assert json.loads((prepared / "materialization.json").read_text())["revision"] == PINNED_SHA
+ summary = validate_dataset(
+ prepared,
+ algorithm=algorithm,
+ profile_name="sd35-t2i",
+ expected_rows=4,
+ )
+ assert summary["row_count"] == 4
+
+
+def test_pending_revision_fails_before_download(tmp_path: Path) -> None:
+ lock_path = tmp_path / "lock.json"
+ _write_lock(lock_path, pending=True)
+
+ def unexpected_download(**kwargs: Any) -> str:
+ raise AssertionError(f"download must not run: {kwargs}")
+
+ with pytest.raises(ValueError, match="publication is still pending"):
+ prepare_dataset(
+ algorithm="sft",
+ profile_name="sd35-t2i",
+ world_size=1,
+ output_root=tmp_path / "prepared",
+ lock_path=lock_path,
+ download_fn=unexpected_download,
+ )
+
+
+def test_canonical_profile_name_fails_before_download(tmp_path: Path) -> None:
+ lock_path = tmp_path / "lock.json"
+ _write_lock(lock_path)
+
+ def unexpected_download(**kwargs: Any) -> str:
+ raise AssertionError(f"download must not run: {kwargs}")
+
+ with pytest.raises(ValueError, match="published runtime alias"):
+ prepare_dataset(
+ algorithm="sft",
+ profile_name="text_to_image",
+ world_size=1,
+ output_root=tmp_path / "prepared",
+ lock_path=lock_path,
+ download_fn=unexpected_download,
+ )
+
+
+def test_repeat_requires_explicit_opt_in(tmp_path: Path) -> None:
+ lock_path = tmp_path / "lock.json"
+ _write_lock(lock_path)
+
+ def download(**kwargs: Any) -> str:
+ local_dir = Path(kwargs["local_dir"])
+ _populate_snapshot(local_dir, algorithm="sft", row_count=2)
+ return str(local_dir)
+
+ kwargs = {
+ "algorithm": "sft",
+ "profile_name": "sd35-t2i",
+ "world_size": 2,
+ "output_root": tmp_path / "prepared",
+ "lock_path": lock_path,
+ "download_fn": download,
+ }
+ with pytest.raises(ValueError, match="--allow-repeat"):
+ prepare_dataset(**kwargs)
+ prepared = prepare_dataset(**kwargs, allow_repeat=True)
+ assert len((prepared / "train.jsonl").read_text().splitlines()) == 4
+
+
+@pytest.mark.parametrize("failure", ["escape", "identical_pair"])
+def test_invalid_snapshot_never_publishes_partial_output(tmp_path: Path, failure: str) -> None:
+ lock_path = tmp_path / "lock.json"
+ _write_lock(lock_path)
+
+ def download(**kwargs: Any) -> str:
+ local_dir = Path(kwargs["local_dir"])
+ target_path = "../../../escape.png" if failure == "escape" else "../../media/target.png"
+ _populate_snapshot(
+ local_dir,
+ algorithm="offline-dpo",
+ target_path=target_path,
+ identical_pair=failure == "identical_pair",
+ )
+ if failure == "escape":
+ Image.new("RGB", (8, 8)).save(local_dir.parent / "escape.png")
+ return str(local_dir)
+
+ with pytest.raises(ValueError, match="escapes|byte-identical"):
+ prepare_dataset(
+ algorithm="offline-dpo",
+ profile_name="sd35-t2i",
+ world_size=2,
+ output_root=tmp_path / "prepared",
+ lock_path=lock_path,
+ download_fn=download,
+ )
+ assert not (tmp_path / "prepared" / "offline-dpo" / "sd35-t2i").exists()
+
+
+def _all_media(row: Mapping[str, Any]):
+ yield from row["input"].get("media", ())
+ for name in ("target", "chosen", "rejected"):
+ candidate = row["supervision"].get(name)
+ if candidate is not None:
+ yield from candidate["media"]
diff --git a/tests/dataset/test_offline_smoke_profiles.py b/tests/dataset/test_offline_smoke_profiles.py
new file mode 100644
index 000000000..d615cd1b0
--- /dev/null
+++ b/tests/dataset/test_offline_smoke_profiles.py
@@ -0,0 +1,144 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Lock the public offline-smoke task catalog and official I/O contracts."""
+
+import subprocess
+import sys
+from dataclasses import FrozenInstanceError
+from types import SimpleNamespace
+
+import pytest
+
+from dataset.offline_smoke import profiles as p
+from flow_factory import contracts as c
+from flow_factory.models.registry import list_registered_models
+
+PROFILE_IDS = tuple(
+ "text_to_image image_to_image multi_image_to_image text_to_video first_frame_to_video "
+ "first_last_frame_to_video text_to_audio_video first_frame_to_audio_video "
+ "first_last_frame_to_audio_video ordered_references_to_audio_video".split()
+)
+
+
+def _input(profile: str, media_type: c.MediaType) -> c.InputMediaRule:
+ return next(
+ rule
+ for rule in p.CANONICAL_PROFILES[profile].contract.input_media.rules
+ if rule.format.type is media_type
+ )
+
+
+def test_catalog_covers_all_actual_profiles_models_repositories_and_aliases() -> None:
+ assert tuple(p.CANONICAL_PROFILES) == PROFILE_IDS
+ assert {
+ model
+ for profile in p.CANONICAL_PROFILES.values()
+ for model in profile.compatible_model_types
+ } == set(list_registered_models())
+ assert p.SFT_REPO_ID == "Jayce-Ping/Flow-Factory-SFT-Smoke"
+ assert p.OFFLINE_DPO_REPO_ID == "Jayce-Ping/Flow-Factory-Offline-DPO-Smoke"
+ assert set(p.MAIN_GPU_ALIASES) == set(p.GPU_ALIAS_TO_PROFILE) - {"image-i2i"}
+ assert p.SUPPLEMENTAL_GPU_ALIASES == ("image-i2i",)
+
+
+def test_actual_profiles_preserve_exact_output_sequences() -> None:
+ outputs = tuple(p.output_media_types(item.contract) for item in p.CANONICAL_PROFILES.values())
+ assert outputs == (("image",),) * 3 + (("video",),) * 3 + (("video", "audio"),) * 4
+ for profile in tuple(p.CANONICAL_PROFILES.values())[6:]:
+ contract = profile.contract
+ assert contract.output_media.items[0].fps is c.RateRequirement.REQUIRED
+ assert contract.output_media.items[1].sample_rate is c.RateRequirement.REQUIRED
+
+
+def test_image_and_endpoint_profiles_keep_cardinality_order_and_slots() -> None:
+ i2i = _input("image_to_image", c.MediaType.IMAGE)
+ multi = _input("multi_image_to_image", c.MediaType.IMAGE)
+ first_video = _input("first_frame_to_video", c.MediaType.IMAGE)
+ fl_video = _input("first_last_frame_to_video", c.MediaType.IMAGE)
+ first_av = _input("first_frame_to_audio_video", c.MediaType.IMAGE)
+ fl_av = _input("first_last_frame_to_audio_video", c.MediaType.IMAGE)
+
+ assert (i2i.min_count, i2i.max_count) == (1, 1)
+ assert (multi.min_count, multi.max_count) == (2, 2)
+ assert (
+ p.CANONICAL_PROFILES["multi_image_to_image"].contract.input_media.order
+ is c.InputMediaOrder.WITHIN_TYPE
+ )
+ assert (first_video.slots, first_video.required_slots) == (("first_frame",), ("first_frame",))
+ assert (fl_video.slots, fl_video.required_slots) == (
+ ("first_frame", "last_frame"),
+ ("first_frame",),
+ )
+ assert (first_av.slots, first_av.required_slots) == (("first_frame",), ("first_frame",))
+ assert (fl_av.slots, fl_av.required_slots) == (("first_frame", "last_frame"), ())
+
+
+def test_h3_reference_contract_is_global_heterogeneous_and_bounded() -> None:
+ contract = p.CANONICAL_PROFILES["ordered_references_to_audio_video"].contract
+ assert contract.input_media.binding is c.InputMediaBinding.ORDERED_REFERENCES
+ assert contract.input_media.order is c.InputMediaOrder.GLOBAL
+ assert (contract.input_media.min_total_count, contract.input_media.max_total_count) == (1, 12)
+ assert contract.input_media.required_any_types == (c.MediaType.IMAGE, c.MediaType.VIDEO)
+ counts = tuple(
+ (rule.format.type.value, rule.min_count, rule.max_count)
+ for rule in contract.input_media.rules
+ )
+ assert counts == (("image", 0, 9), ("video", 0, 3), ("audio", 0, 3))
+
+
+def test_gpu_variants_keep_model_specific_clocks() -> None:
+ cases = {
+ case.alias: case for profile in p.CANONICAL_PROFILES.values() for case in profile.gpu_cases
+ }
+ wan, ltx, h3 = (cases[name].geometry for name in ("wan-t2v", "ltx2-t2av", "h3-t2va"))
+ assert (wan.height, wan.num_frames) == (240, 5)
+ assert (ltx.height, ltx.width, ltx.num_frames, ltx.sample_rate) == (128, 192, 9, 16000)
+ assert (h3.height, h3.width, h3.num_frames, h3.sample_rate) == (64, 96, 124, 32000)
+ assert p.get_profile("image-i2i") is p.CANONICAL_PROFILES["image_to_image"]
+
+
+def test_official_contract_supports_ordered_heterogeneous_input_and_audio_only_output() -> None:
+ rate = c.RateRequirement
+ video = c.MediaFormat(c.MediaType.VIDEO, rate.OPTIONAL, rate.NOT_APPLICABLE)
+ audio_input = c.MediaFormat(c.MediaType.AUDIO, rate.NOT_APPLICABLE, rate.OPTIONAL)
+ audio_output = c.MediaFormat(c.MediaType.AUDIO, rate.NOT_APPLICABLE, rate.REQUIRED)
+ rules = (c.InputMediaRule(video, 1, 2), c.InputMediaRule(audio_input, 1, 2))
+ inputs = c.InputMediaSpec(
+ rules, c.InputMediaBinding.ORDERED_REFERENCES, c.InputMediaOrder.GLOBAL, 2, 4
+ )
+ contract = c.PipelineIOContract(
+ inputs,
+ c.NegativePromptPolicy.UNSUPPORTED,
+ c.OutputMediaSequence((audio_output,)),
+ c.GeometrySource.OUTPUT_MEDIA,
+ c.BatchCapability.UNIFORM,
+ )
+ media = (
+ SimpleNamespace(type="video", fps=24.0, sample_rate=None),
+ SimpleNamespace(type="audio", fps=None, sample_rate=16000),
+ )
+ model_input = SimpleNamespace(prompt="Compose a sound.", negative_prompt=None, media=media)
+ c.validate_pipeline_model_input(model_input, contract)
+ assert p.output_media_types(contract) == ("audio",)
+ output = SimpleNamespace(type="audio", fps=None, sample_rate=16000)
+ c.validate_pipeline_output_candidate((output,), contract)
+ assert "audio_only" not in p.CANONICAL_PROFILES
+
+
+def test_catalog_is_frozen_and_import_does_not_load_bagel_adapter() -> None:
+ with pytest.raises(FrozenInstanceError):
+ p.CANONICAL_PROFILES["text_to_image"].name = "changed"
+ script = "import sys; import dataset.offline_smoke.profiles; assert 'flow_factory.models.bagel.bagel' not in sys.modules"
+ subprocess.run([sys.executable, "-c", script], check=True)
diff --git a/tests/dataset/test_offline_smoke_publish.py b/tests/dataset/test_offline_smoke_publish.py
new file mode 100644
index 000000000..ac2237374
--- /dev/null
+++ b/tests/dataset/test_offline_smoke_publish.py
@@ -0,0 +1,98 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from dataset.offline_smoke import publish
+from dataset.offline_smoke.profiles import SFT_REPO_ID
+
+
+def _staging_tree(
+ root: Path,
+ *,
+ repository_id: str = SFT_REPO_ID,
+ supervision_type: str = "demonstration",
+) -> Path:
+ for directory in ("media", "profiles/sd35-t2i"):
+ (root / directory).mkdir(parents=True, exist_ok=True)
+ for filename in ("README.md", "LICENSE", "provenance.jsonl"):
+ (root / filename).write_text("fixture\n", encoding="utf-8")
+ (root / "dataset_manifest.json").write_text(
+ json.dumps(
+ {
+ "schema_version": 1,
+ "flow_factory_schema_version": 2,
+ "repository_id": repository_id,
+ "supervision_type": supervision_type,
+ "runtime_aliases": ["sd35-t2i"],
+ }
+ ),
+ encoding="utf-8",
+ )
+ return root
+
+
+def test_wrong_supervision_fails_before_any_hub_mutation(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ staging = _staging_tree(
+ tmp_path / "dpo-staging",
+ repository_id="Jayce-Ping/Flow-Factory-Offline-DPO-Smoke",
+ supervision_type="preference",
+ )
+
+ def unexpected_api() -> None:
+ raise AssertionError("HfApi must not be constructed for mismatched staging")
+
+ monkeypatch.setattr(publish, "HfApi", unexpected_api)
+ with pytest.raises(ValueError, match="requires 'demonstration' supervision"):
+ publish.publish_staging_tree(
+ algorithm="sft",
+ staging_dir=staging,
+ commit_message="must not publish",
+ confirm_public=True,
+ )
+
+
+def test_valid_staging_binds_default_repo_and_returns_commit(
+ tmp_path: Path,
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ staging = _staging_tree(tmp_path / "sft-staging")
+ calls = []
+
+ class FakeApi:
+ def create_repo(self, **kwargs) -> None:
+ calls.append(("create_repo", kwargs))
+
+ def upload_folder(self, **kwargs):
+ calls.append(("upload_folder", kwargs))
+ return SimpleNamespace(oid="b" * 40)
+
+ monkeypatch.setattr(publish, "HfApi", FakeApi)
+ revision = publish.publish_staging_tree(
+ algorithm="sft",
+ staging_dir=staging,
+ commit_message="publish test fixture",
+ confirm_public=True,
+ )
+
+ assert revision == "b" * 40
+ assert [name for name, _ in calls] == ["create_repo", "upload_folder"]
+ assert all(kwargs["repo_id"] == SFT_REPO_ID for _, kwargs in calls)
From 0e098eb9d81b957a9ddea46b06ca56b6d8278389 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 18:47:35 +0800
Subject: [PATCH 33/76] [samples,reward] fix: preserve reconstruction fields in
distributed rewards
---
.agents/knowledge/topics/fix_patterns.md | 8 +++
src/flow_factory/rewards/reward_processor.py | 5 ++
src/flow_factory/samples/samples.py | 9 ++-
.../test_reward_processor_reconstruction.py | 57 +++++++++++++++++++
4 files changed, 77 insertions(+), 2 deletions(-)
create mode 100644 tests/rewards/test_reward_processor_reconstruction.py
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 2223bed35..a3a7387c1 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -54,6 +54,14 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Lesson**: When extending a base contract for a partial-coverage feature (where only some subclasses will participate), no-op default + opt-in override beats forcing every subclass to acknowledge it. Reserve `@abstractmethod` for invariants that ALL subclasses must implement (e.g. `load_pipeline`, `decode_latents`, `forward`, `inference`).
- **Related Constraint**: #12 (post-update text codifies "Optional encoder overrides (no-op default)").
+### Distributed reward gathering must preserve reconstruction invariants
+- **Date**: 2026-08-30
+- **Symptom**: Two-rank H3 Ref2VA GRPO failed before reward execution because `MiniMaxH3Ref2VASample` was reconstructed with `reference_manifest=None`.
+- **Root Cause**: The distributed group-reward path gathered only reward-consumed fields, although `gather_samples` reconstructs the concrete sample class and that class can require additional state.
+- **Fix**: `BaseSample` now declares an empty `reconstruction_required_fields` contract, `OrderedReferenceConditionSample` adds `reference_manifest`, and `RewardProcessor` unions that contract into its distributed gather fields without forwarding it to the reward call.
+- **Lesson**: Communication payload requirements and reward-call requirements are distinct contracts; partial gathers must preserve constructor invariants even for fields that downstream computation does not consume.
+- **Related Constraint**: N/A
+
### Preference-arm conditioning ownership
- **Date**: 2026-08-11
- **Symptom**: DPO evaluated the rejected H3 state with the chosen sample's prompt/reference conditioning.
diff --git a/src/flow_factory/rewards/reward_processor.py b/src/flow_factory/rewards/reward_processor.py
index ccd000867..8632ff06c 100644
--- a/src/flow_factory/rewards/reward_processor.py
+++ b/src/flow_factory/rewards/reward_processor.py
@@ -474,6 +474,11 @@ def _compute_groupwise_distributed(
for model in models.values():
required_fields.update(model.required_fields)
+ # ``gather_samples`` reconstructs the concrete sample class from the
+ # transported fields. Preserve fields required by that class's constructor
+ # invariants even when the reward itself does not consume them.
+ required_fields.update(type(samples[0]).reconstruction_required_fields)
+
# Always include the typed source bookkeeping — the gate needs
# `source` (and ideally `source_id`) on the gathered side. Now
# that they're real dataclass fields on `BaseSample`, gathering
diff --git a/src/flow_factory/samples/samples.py b/src/flow_factory/samples/samples.py
index 71c914fa7..86d2e7dd0 100644
--- a/src/flow_factory/samples/samples.py
+++ b/src/flow_factory/samples/samples.py
@@ -97,6 +97,10 @@ class BaseSample:
{"height", "width", "latent_index_map", "log_prob_index_map"}
)
+ # Fields that must be transported whenever a concrete sample is reconstructed
+ # from a partial cross-rank gather, even when no downstream consumer reads them.
+ reconstruction_required_fields: ClassVar[frozenset[str]] = frozenset()
+
# Denoiseing trajectory
timesteps: Optional[torch.Tensor] = None # (T+1,)
all_latents: Optional[torch.Tensor] = None # (num_steps, Seq_len, C)
@@ -633,8 +637,9 @@ class T2AVSample(BaseSample):
class OrderedReferenceConditionSample(BaseSample):
"""Sample conditioned by an ordered heterogeneous reference manifest."""
- _id_fields: ClassVar[frozenset[str]] = BaseSample._id_fields | frozenset(
- {"reference_manifest"}
+ _id_fields: ClassVar[frozenset[str]] = BaseSample._id_fields | frozenset({"reference_manifest"})
+ reconstruction_required_fields: ClassVar[frozenset[str]] = (
+ BaseSample.reconstruction_required_fields | frozenset({"reference_manifest"})
)
reference_manifest: Optional[str] = None
diff --git a/tests/rewards/test_reward_processor_reconstruction.py b/tests/rewards/test_reward_processor_reconstruction.py
new file mode 100644
index 000000000..f9fd45574
--- /dev/null
+++ b/tests/rewards/test_reward_processor_reconstruction.py
@@ -0,0 +1,57 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import torch
+
+from flow_factory.rewards import GroupwiseRewardModel, RewardModelOutput
+from flow_factory.rewards.reward_processor import RewardProcessor
+from flow_factory.samples import MiniMaxH3Ref2VASample
+
+
+class PromptOnlyGroupReward(GroupwiseRewardModel):
+ required_fields = ("prompt",)
+
+ def __init__(self) -> None:
+ pass
+
+ def __call__(self, prompt: list[str]) -> RewardModelOutput:
+ return RewardModelOutput(rewards=torch.zeros(len(prompt)))
+
+
+def test_distributed_group_reward_preserves_sample_reconstruction_fields() -> None:
+ accelerator = SimpleNamespace(
+ device=torch.device("cpu"),
+ process_index=0,
+ num_processes=1,
+ is_local_main_process=True,
+ wait_for_everyone=lambda: None,
+ reduce=lambda tensor, reduction: tensor,
+ )
+ model = PromptOnlyGroupReward()
+ processor = RewardProcessor(
+ accelerator=accelerator,
+ reward_models={"prompt_only": model},
+ group_on_same_rank=False,
+ verbose=False,
+ )
+ sample = MiniMaxH3Ref2VASample(
+ prompt="A reference-conditioned prompt",
+ reference_manifest='[{"kind":"image","path":"condition.png"}]',
+ )
+
+ rewards = processor.compute_rewards([sample], store_to_samples=False)
+
+ torch.testing.assert_close(rewards["prompt_only"], torch.zeros(1))
From 527092d78d42b64b07d62f18f2e1242e0231cfc8 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 19:12:18 +0800
Subject: [PATCH 34/76] [trainer] fix: preserve ZeRO logical optimizer identity
---
.agents/knowledge/topics/fix_patterns.md | 14 +++
.../trainers/common/runtime_identity.py | 48 ++++++++--
tests/trainers/test_runtime_identity.py | 92 +++++++++++++++++++
3 files changed, 146 insertions(+), 8 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index a3a7387c1..e6747a3e7 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -333,6 +333,20 @@ Based on the fix type, write the fix entry to the appropriate document:
and set-like identity fields require one canonical representation.
- **Related Constraint**: #5
+### ZeRO optimizer identity must use logical model groups
+- **Date**: 2026-08-30
+- **Symptom**: Every ZeRO-2 trainer failed during initialization because its optimizer schema
+ contained a parameter not owned by the rebound component-variant registry.
+- **Root Cause**: DeepSpeed ZeRO-1/2 replaces each public optimizer group with a rank-local flat
+ FP32 master partition, while runtime identity incorrectly treated those partitions as the live
+ model parameters owned by the registry.
+- **Fix**: Runtime identity now maps stable parameter ownership through DeepSpeed's retained
+ `bit16_groups` and continues to serialize settings from the public optimizer groups. It fails
+ closed if logical groups are absent or do not match the partitioned group count.
+- **Lesson**: A distributed optimizer's public parameter groups may be physical state partitions;
+ exact-resume identity must separate logical model ownership from physical group settings.
+- **Related Constraint**: #18a
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/trainers/common/runtime_identity.py b/src/flow_factory/trainers/common/runtime_identity.py
index 356df6f9d..5fc6148e2 100644
--- a/src/flow_factory/trainers/common/runtime_identity.py
+++ b/src/flow_factory/trainers/common/runtime_identity.py
@@ -25,6 +25,7 @@
from typing import Any
import torch
+from accelerate.utils import DistributedType
from torch.utils.data import ConcatDataset, DataLoader, Subset
_EXECUTION_IDENTITY_HOOK = "runtime_execution_identity_payload"
@@ -930,10 +931,13 @@ def _parameter_schema(trainer: Any) -> tuple[list[dict[str, Any]], dict[int, str
def _optimizer_schema(trainer: Any, parameter_keys: Mapping[int, str]) -> dict[str, Any]:
"""Return ordered optimizer groups linked to the stable parameter schema."""
optimizer = trainer.optimizer
+ optimizer_groups = optimizer.param_groups
+ logical_parameter_groups = _logical_optimizer_parameter_groups(trainer, optimizer_groups)
groups = []
- consumed_parameters: set[int] = set()
- for group_index, group in enumerate(optimizer.param_groups):
- raw_parameters = group.get("params")
+ consumed_parameters: set[str] = set()
+ for group_index, (group, raw_parameters) in enumerate(
+ zip(optimizer_groups, logical_parameter_groups)
+ ):
if not isinstance(raw_parameters, Sequence):
raise TypeError(
f"optimizer group {group_index} params must be a sequence, "
@@ -947,9 +951,9 @@ def _optimizer_schema(trainer: Any, parameter_keys: Mapping[int, str]) -> dict[s
"optimizer schema contains a parameter not owned by the rebound "
f"variant registry at group {group_index}, index {parameter_index}"
)
- if id(parameter) in consumed_parameters:
+ if key in consumed_parameters:
raise ValueError(f"optimizer schema references parameter {key!r} more than once")
- consumed_parameters.add(id(parameter))
+ consumed_parameters.add(key)
group_parameters.append(key)
settings = {}
for key, value in group.items():
@@ -970,11 +974,11 @@ def _optimizer_schema(trainer: Any, parameter_keys: Mapping[int, str]) -> dict[s
"settings": settings,
}
)
- missing_parameters = frozenset(parameter_keys).difference(consumed_parameters)
+ missing_parameters = frozenset(parameter_keys.values()).difference(consumed_parameters)
if missing_parameters:
- missing_keys = tuple(parameter_keys[identity] for identity in missing_parameters)
raise ValueError(
- "optimizer schema does not exhaust rebound variant parameters: " f"{missing_keys!r}"
+ "optimizer schema does not exhaust rebound variant parameters: "
+ f"{tuple(sorted(missing_parameters))!r}"
)
return {
"type_chain": _optimizer_type_chain(optimizer),
@@ -982,6 +986,34 @@ def _optimizer_schema(trainer: Any, parameter_keys: Mapping[int, str]) -> dict[s
}
+def _logical_optimizer_parameter_groups(
+ trainer: Any,
+ optimizer_groups: Sequence[Mapping[str, Any]],
+) -> Sequence[Sequence[torch.Tensor]]:
+ """Return model-owned parameters before ZeRO replaces groups with flat partitions."""
+ accelerator = trainer.accelerator
+ deepspeed_plugin = getattr(getattr(accelerator, "state", None), "deepspeed_plugin", None)
+ if getattr(accelerator, "distributed_type", None) != DistributedType.DEEPSPEED or getattr(
+ deepspeed_plugin, "zero_stage", None
+ ) not in (1, 2):
+ return tuple(group.get("params") for group in optimizer_groups)
+
+ deepspeed_optimizer = getattr(trainer.optimizer, "optimizer", None)
+ logical_groups = getattr(deepspeed_optimizer, "bit16_groups", None)
+ if not isinstance(logical_groups, Sequence) or isinstance(logical_groups, (str, bytes)):
+ raise TypeError(
+ "DeepSpeed ZeRO-1/2 optimizer schema requires the logical model parameter "
+ f"groups from optimizer.bit16_groups, received {type(logical_groups).__name__}: "
+ f"{logical_groups!r}"
+ )
+ if len(logical_groups) != len(optimizer_groups):
+ raise ValueError(
+ "DeepSpeed ZeRO-1/2 optimizer schema expected logical and partitioned group "
+ f"counts to match, received {len(logical_groups)} and {len(optimizer_groups)}"
+ )
+ return logical_groups
+
+
def _optimizer_type_chain(optimizer: Any) -> list[str]:
"""Describe transparent optimizer wrappers without following cycles."""
names = []
diff --git a/tests/trainers/test_runtime_identity.py b/tests/trainers/test_runtime_identity.py
index 7350e667a..f1b368b0f 100644
--- a/tests/trainers/test_runtime_identity.py
+++ b/tests/trainers/test_runtime_identity.py
@@ -362,6 +362,95 @@ def test_optimizer_schema_rejects_parameters_outside_rebound_registry() -> None:
build_trainer_runtime_identity(trainer)
+@pytest.mark.parametrize("zero_stage", (1, 2))
+def test_deepspeed_optimizer_schema_maps_flat_partitions_through_logical_groups(
+ zero_stage: int,
+) -> None:
+ """ZeRO flat optimizer partitions retain their logical variant ownership."""
+
+ def identity(partition_widths: tuple[int, int]) -> dict[str, Any]:
+ trainer = _Trainer()
+ fake_parameter = torch.nn.Parameter(torch.ones(3, 3))
+ trainer.adapter.component_variant_registry.records["fake"] = (
+ _Record("transformer", "fake_weight", fake_parameter),
+ )
+ trainer.optimizer.add_param_group(
+ {"params": [fake_parameter], "role_name": "fake", "lr": 2e-3}
+ )
+ trainer.config.optimizer_args = MultiOptimizerArguments(
+ optimizer_configs=[
+ AdamWOptimizerArguments(name="base", learning_rate=1e-3),
+ AdamWOptimizerArguments(name="fake", learning_rate=2e-3),
+ ]
+ )
+ trainer._required_trainable_roles = lambda: ("base", "fake")
+
+ basic_optimizer = trainer.optimizer
+ logical_groups = [list(group["params"]) for group in basic_optimizer.param_groups]
+ for group, width in zip(basic_optimizer.param_groups, partition_widths):
+ group["params"] = [torch.nn.Parameter(torch.zeros(width), requires_grad=True)]
+ zero_optimizer = SimpleNamespace(
+ bit16_groups=logical_groups,
+ optimizer=basic_optimizer,
+ )
+ trainer.optimizer = SimpleNamespace(
+ param_groups=basic_optimizer.param_groups,
+ optimizer=zero_optimizer,
+ )
+ trainer.accelerator.distributed_type = DistributedType.DEEPSPEED
+ trainer.accelerator.state.deepspeed_plugin = SimpleNamespace(
+ zero_stage=zero_stage,
+ deepspeed_config={"zero_optimization": {"stage": zero_stage}},
+ gradient_accumulation_steps=1,
+ gradient_clipping=1.0,
+ is_train_batch_min=True,
+ )
+ return build_trainer_runtime_identity(trainer)
+
+ rank_zero = identity((5, 7))
+ rank_one = identity((6, 8))
+
+ assert rank_zero["parameter_schema_digest"] == rank_one["parameter_schema_digest"]
+ assert rank_zero["optimizer_schema_digest"] == rank_one["optimizer_schema_digest"]
+
+
+def test_deepspeed_optimizer_schema_rejects_foreign_logical_parameters() -> None:
+ """ZeRO logical groups must still exhaust only registry-owned parameters."""
+ trainer = _Trainer()
+ trainer.accelerator.distributed_type = DistributedType.DEEPSPEED
+ trainer.accelerator.state.deepspeed_plugin = SimpleNamespace(zero_stage=2)
+ trainer.optimizer.optimizer = SimpleNamespace(
+ bit16_groups=[[torch.nn.Parameter(torch.zeros(1))]]
+ )
+
+ with pytest.raises(ValueError, match="not owned by the rebound variant registry"):
+ build_trainer_runtime_identity(trainer)
+
+
+@pytest.mark.parametrize(
+ ("logical_groups", "error_type", "message"),
+ (
+ (None, TypeError, "requires the logical model parameter groups"),
+ ([], ValueError, "logical and partitioned group counts to match"),
+ ),
+)
+def test_deepspeed_optimizer_schema_rejects_invalid_logical_groups(
+ logical_groups: list[list[torch.nn.Parameter]] | None,
+ error_type: type[Exception],
+ message: str,
+) -> None:
+ """ZeRO identity fails closed when its logical parameter seam is unavailable."""
+ trainer = _Trainer()
+ trainer.accelerator.distributed_type = DistributedType.DEEPSPEED
+ trainer.accelerator.state.deepspeed_plugin = SimpleNamespace(zero_stage=2)
+ trainer.optimizer.optimizer = SimpleNamespace()
+ if logical_groups is not None:
+ trainer.optimizer.optimizer.bit16_groups = logical_groups
+
+ with pytest.raises(error_type, match=message):
+ build_trainer_runtime_identity(trainer)
+
+
def test_backend_and_precision_drift_change_the_resume_identity() -> None:
"""Backend checkpoint layouts are rejected before prepared-state mutation."""
baseline_trainer = _Trainer()
@@ -457,6 +546,9 @@ def identity(*, micro_batch: int, accumulation_steps: int) -> dict[str, Any]:
gradient_clipping="auto",
is_train_batch_min=True,
)
+ trainer.optimizer.optimizer = SimpleNamespace(
+ bit16_groups=[list(group["params"]) for group in trainer.optimizer.param_groups]
+ )
return build_trainer_runtime_identity(trainer)
baseline = identity(micro_batch=1, accumulation_steps=2)
From e570d1430d32daf87b43f5d8600a63b2c9eea3b4 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 19:26:39 +0800
Subject: [PATCH 35/76] [models] fix: enable Bagel gradient checkpointing
---
.agents/knowledge/topics/fix_patterns.md | 16 +++
.../bagel/modeling/bagel/qwen2_navit.py | 34 ++++--
tests/models/test_bagel_tdm_contracts.py | 112 ++++++++++++++++++
3 files changed, 152 insertions(+), 10 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index e6747a3e7..ca4c4b76f 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -347,6 +347,22 @@ Based on the fix type, write the fix entry to the appropriate document:
exact-resume identity must separate logical model ownership from physical group settings.
- **Related Constraint**: #18a
+### Custom Transformers models must implement their declared checkpointing seam
+- **Date**: 2026-08-30
+- **Symptom**: Every Bagel trainer failed during initialization when full gradient checkpointing
+ called `Qwen2ForCausalLM.gradient_checkpointing_enable()` and Transformers reported that the
+ architecture was incompatible.
+- **Root Cause**: Bagel's custom Qwen2-NaViT model inherited
+ `supports_gradient_checkpointing = True` but did not expose a `gradient_checkpointing` state or
+ invoke the installed checkpoint function in its active decoder loop.
+- **Fix**: The custom Qwen2 model now owns the standard checkpointing flag and routes each pure
+ decoder layer through Transformers' installed non-reentrant checkpoint function while training.
+ Cache-updating and TaylorSeer paths remain direct to avoid replaying mutations. A backward
+ regression proves the layer is recomputed rather than merely accepting the API call.
+- **Lesson**: A custom `PreTrainedModel` must pair its capability declaration with both the state
+ seam expected by Transformers and an execution-time checkpoint boundary in the forward path.
+- **Related Constraint**: #7
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
index aa4f172fb..43e13ee6c 100644
--- a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
+++ b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
@@ -1083,6 +1083,7 @@ def __init__(self, config):
if self.use_moe:
self.norm_moe_gen = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.rotary_emb = Qwen2RotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
# Initialize weights and apply final processing
self.post_init()
@@ -1194,18 +1195,31 @@ def forward_inference(
decoder_layer.cache_dic = self.cache_dic
decoder_layer.enable_taylorseer = True
self.current["layer"] = layer_idx
- packed_query_sequence, past_key_values = decoder_layer(
- packed_query_sequence=packed_query_sequence,
- query_lens=query_lens,
- packed_query_position_embeddings=packed_query_position_embeddings,
- packed_query_indexes=packed_query_indexes,
- past_key_values=past_key_values,
- key_values_lens=key_values_lens,
- packed_key_value_indexes=packed_key_value_indexes,
- update_past_key_values=update_past_key_values,
- is_causal=is_causal,
+ layer_kwargs = {
+ "packed_query_sequence": packed_query_sequence,
+ "query_lens": query_lens,
+ "packed_query_position_embeddings": packed_query_position_embeddings,
+ "packed_query_indexes": packed_query_indexes,
+ "past_key_values": past_key_values,
+ "key_values_lens": key_values_lens,
+ "packed_key_value_indexes": packed_key_value_indexes,
+ "update_past_key_values": update_past_key_values,
+ "is_causal": is_causal,
**extra_inputs,
+ }
+ checkpoint_layer = (
+ self.gradient_checkpointing
+ and self.training
+ and not update_past_key_values
+ and not enable_taylorseer
)
+ if checkpoint_layer:
+ packed_query_sequence, past_key_values = self._gradient_checkpointing_func(
+ decoder_layer.__call__,
+ **layer_kwargs,
+ )
+ else:
+ packed_query_sequence, past_key_values = decoder_layer(**layer_kwargs)
if self.use_moe:
if mode == "und":
diff --git a/tests/models/test_bagel_tdm_contracts.py b/tests/models/test_bagel_tdm_contracts.py
index 949359ea5..bb5c01a3a 100644
--- a/tests/models/test_bagel_tdm_contracts.py
+++ b/tests/models/test_bagel_tdm_contracts.py
@@ -48,6 +48,70 @@ def _load_bagel_types(monkeypatch: pytest.MonkeyPatch) -> tuple[type, type]:
return module.BagelAdapter, module.BagelSample
+def _load_bagel_qwen_types(monkeypatch: pytest.MonkeyPatch) -> tuple[type, type]:
+ """Load Bagel's Qwen2-NaViT model behind the optional-kernel seam."""
+ _load_bagel_types(monkeypatch)
+ module = importlib.import_module("flow_factory.models.bagel.modeling.bagel.qwen2_navit")
+ return module.Qwen2Config, module.Qwen2ForCausalLM
+
+
+class _FakeRotaryEmbedding(torch.nn.Module):
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_ids: torch.Tensor,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ del position_ids
+ values = torch.ones_like(hidden_states).unsqueeze(0)
+ return values, values
+
+
+class _CheckpointedDecoderLayer(torch.nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.calls = 0
+
+ def forward(
+ self,
+ *,
+ packed_query_sequence: torch.Tensor,
+ past_key_values: Any,
+ **kwargs: Any,
+ ) -> tuple[torch.Tensor, Any]:
+ del kwargs
+ self.calls += 1
+ return packed_query_sequence.square(), past_key_values
+
+
+def _checkpointing_qwen_model(
+ monkeypatch: pytest.MonkeyPatch,
+) -> tuple[torch.nn.Module, _CheckpointedDecoderLayer]:
+ qwen_config_type, qwen_causal_lm_type = _load_bagel_qwen_types(monkeypatch)
+ config = qwen_config_type(
+ vocab_size=8,
+ hidden_size=8,
+ intermediate_size=16,
+ num_hidden_layers=1,
+ num_attention_heads=2,
+ num_key_value_heads=2,
+ max_position_embeddings=16,
+ layer_module="Qwen2DecoderLayer",
+ qk_norm=False,
+ _attn_implementation="eager",
+ pad_token_id=0,
+ )
+ causal_lm = qwen_causal_lm_type(config)
+ model = causal_lm.model
+ decoder = _CheckpointedDecoderLayer()
+ model.layers = torch.nn.ModuleList([decoder])
+ model.rotary_emb = _FakeRotaryEmbedding()
+ model.norm = torch.nn.Identity()
+ model.use_moe = False
+ causal_lm.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+ model.train()
+ return model, decoder
+
+
def _adapter(adapter_type: type, decoder: Any) -> Any:
adapter = object.__new__(adapter_type)
adapter.decode_latents = MethodType(decoder, adapter)
@@ -69,6 +133,54 @@ def _result(batch_size: int = 2) -> dict[str, Any]:
}
+def test_bagel_qwen_gradient_checkpointing_recomputes_decoder(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ model, decoder = _checkpointing_qwen_model(monkeypatch)
+
+ packed = torch.tensor([[2.0, 3.0]], requires_grad=True)
+ cache = object()
+ output = model.forward_inference(
+ packed_query_sequence=packed,
+ query_lens=torch.tensor([1]),
+ packed_query_position_ids=torch.tensor([0]),
+ packed_query_indexes=torch.tensor([0]),
+ past_key_values=cache,
+ key_values_lens=torch.tensor([0]),
+ packed_key_value_indexes=torch.tensor([], dtype=torch.long),
+ update_past_key_values=False,
+ is_causal=False,
+ )
+ output.packed_query_sequence.sum().backward()
+
+ assert model.gradient_checkpointing is True
+ assert decoder.calls == 2
+ assert output.past_key_values is cache
+ torch.testing.assert_close(packed.grad, torch.tensor([[4.0, 6.0]]))
+
+
+def test_bagel_qwen_checkpointing_does_not_replay_cache_updates(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ model, decoder = _checkpointing_qwen_model(monkeypatch)
+
+ packed = torch.tensor([[2.0, 3.0]], requires_grad=True)
+ output = model.forward_inference(
+ packed_query_sequence=packed,
+ query_lens=torch.tensor([1]),
+ packed_query_position_ids=torch.tensor([0]),
+ packed_query_indexes=torch.tensor([0]),
+ past_key_values=object(),
+ key_values_lens=torch.tensor([0]),
+ packed_key_value_indexes=torch.tensor([], dtype=torch.long),
+ update_past_key_values=True,
+ is_causal=False,
+ )
+ output.packed_query_sequence.sum().backward()
+
+ assert decoder.calls == 1
+
+
def test_bagel_assembles_samples_from_one_batched_decode(
monkeypatch: pytest.MonkeyPatch,
) -> None:
From 05c1d592fadd1ea1088ee601e1eee79fc16caf73 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 19:39:10 +0800
Subject: [PATCH 36/76] [models] fix: align Bagel FSDP wrap classes
---
.agents/knowledge/topics/fix_patterns.md | 13 ++++++
.../bagel/modeling/bagel/qwen2_navit.py | 2 +
tests/models/test_bagel_tdm_contracts.py | 45 +++++++++++++++++++
3 files changed, 60 insertions(+)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index ca4c4b76f..a21928655 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -363,6 +363,19 @@ Based on the fix type, write the fix entry to the appropriate document:
seam expected by Transformers and an execution-time checkpoint boundary in the forward path.
- **Related Constraint**: #7
+### FSDP wrap metadata must follow the instantiated architecture variant
+- **Date**: 2026-08-30
+- **Symptom**: Every Bagel FSDP2 trainer failed during `accelerator.prepare()` because Accelerate
+ could not find the declared `Qwen2DecoderLayer` in the loaded model.
+- **Root Cause**: Bagel's custom Qwen2 classes inherited fixed `_no_split_modules` metadata from the
+ standard decoder even though `config.layer_module` instantiated a MoE or MoT decoder variant.
+- **Fix**: Both the inner Qwen2 model and outer causal LM now derive their no-split class from the
+ realized `layer_module`. CPU FSDP2 auto-wrap regressions verify all decoder variants resolve
+ through the same PEFT wrapper used by LoRA training.
+- **Lesson**: Distributed wrap metadata is realized model state. Config-selectable architectures
+ must not inherit a fixed block class that may be absent from the instantiated module tree.
+- **Related Constraint**: #9
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
index 43e13ee6c..d81e96f3c 100644
--- a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
+++ b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
@@ -1069,6 +1069,7 @@ def forward_inference(
class Qwen2Model(Qwen2PreTrainedModel):
def __init__(self, config):
super().__init__(config)
+ self._no_split_modules = [config.layer_module]
self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size
self.use_moe = "Mo" in config.layer_module
@@ -1250,6 +1251,7 @@ class Qwen2ForCausalLM(Qwen2PreTrainedModel):
def __init__(self, config):
super().__init__(config)
+ self._no_split_modules = [config.layer_module]
self.model = Qwen2Model(config)
self.vocab_size = config.vocab_size
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
diff --git a/tests/models/test_bagel_tdm_contracts.py b/tests/models/test_bagel_tdm_contracts.py
index bb5c01a3a..9d6bf1010 100644
--- a/tests/models/test_bagel_tdm_contracts.py
+++ b/tests/models/test_bagel_tdm_contracts.py
@@ -23,8 +23,11 @@
import pytest
import torch
+from accelerate import FullyShardedDataParallelPlugin
+from peft import LoraConfig, get_peft_model
import flow_factory.utils.imports as import_utils
+from flow_factory.models.model_bundle import ModelBundle
from flow_factory.samples import BaseSample
from flow_factory.trainers.distillation.distillation_runtime import (
validate_media_free_rollout,
@@ -181,6 +184,48 @@ def test_bagel_qwen_checkpointing_does_not_replay_cache_updates(
assert decoder.calls == 1
+@pytest.mark.parametrize(
+ "layer_module",
+ ["Qwen2DecoderLayer", "Qwen2MoEDecoderLayer", "Qwen2MoTDecoderLayer"],
+)
+def test_bagel_qwen_fsdp_wrap_policy_matches_decoder_variant(
+ monkeypatch: pytest.MonkeyPatch,
+ layer_module: str,
+) -> None:
+ qwen_config_type, qwen_causal_lm_type = _load_bagel_qwen_types(monkeypatch)
+ config = qwen_config_type(
+ vocab_size=8,
+ hidden_size=8,
+ intermediate_size=16,
+ num_hidden_layers=1,
+ num_attention_heads=2,
+ num_key_value_heads=2,
+ max_position_embeddings=16,
+ layer_module=layer_module,
+ qk_norm=False,
+ _attn_implementation="eager",
+ pad_token_id=0,
+ )
+ base_model = qwen_causal_lm_type(config)
+ decoder_type = type(base_model.model.layers[0])
+ model = get_peft_model(
+ base_model,
+ LoraConfig(r=2, lora_alpha=2, target_modules=["q_proj"]),
+ )
+ bundle = ModelBundle({"transformer": model})
+ plugin = FullyShardedDataParallelPlugin(
+ fsdp_version=2,
+ auto_wrap_policy="transformer_based_wrap",
+ )
+
+ plugin.set_auto_wrap_policy(bundle)
+
+ assert set(model._no_split_modules) == {decoder_type.__name__}
+ assert set(base_model.model._no_split_modules) == {decoder_type.__name__}
+ assert set(bundle._no_split_modules) == {decoder_type.__name__}
+ assert plugin.auto_wrap_policy.keywords["transformer_layer_cls"] == {decoder_type}
+
+
def test_bagel_assembles_samples_from_one_batched_decode(
monkeypatch: pytest.MonkeyPatch,
) -> None:
From 95d05b468032b1bd2ccaa557bda86a981c5b468c Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:04:08 +0800
Subject: [PATCH 37/76] [models] fix: route Bagel FSDP language model forwards
---
.agents/knowledge/topics/fix_patterns.md | 19 ++
src/flow_factory/models/bagel/bagel.py | 33 ++-
.../models/bagel/modeling/bagel/bagel.py | 40 ++--
.../bagel/modeling/bagel/qwen2_navit.py | 33 ++-
tests/models/test_bagel_tdm_contracts.py | 212 +++++++++++++++++-
5 files changed, 309 insertions(+), 28 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index a21928655..1f2d9efb1 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -376,6 +376,25 @@ Based on the fix type, write the fix entry to the appropriate document:
must not inherit a fixed block class that may be absent from the instantiated module tree.
- **Related Constraint**: #9
+### Parameter-sharded submodule work must remain inside the prepared root forward
+- **Date**: 2026-08-30
+- **Symptom**: Every Bagel FSDP2 trainer reached sampling or offline replay but failed at token
+ embedding with a mixed `torch.Tensor` and `DTensor` operator error.
+- **Root Cause**: Bagel's cache helpers and denoising path reached through the physical pipeline to
+ `language_model.model.embed_tokens` before calling the routed transformer. Decoder layers owned
+ nested FSDP groups, but embedding and final normalization belonged to the prepared `ModelBundle`
+ root, whose unshard hook was bypassed by those direct calls.
+- **Fix**: The outer Qwen forward now accepts raw packed token IDs and inserts their embeddings into
+ an optional query-local auxiliary sequence. Bagel cache helpers accept an injected language-model
+ forward, and the adapter supplies its routed transformer for text, VAE, ViT, denoising, and CFG
+ passes. Each logical language-model pass therefore performs embedding, decoder execution, and
+ final normalization inside one prepared-root call.
+- **Lesson**: Under compositional FSDP, wrapping transformer blocks does not make arbitrary child
+ access safe. Any computation using parameters owned by the prepared root must execute beneath
+ that root's forward hooks; converting ordinary inputs to `DTensor` or manually unsharding only
+ masks the first failure and breaks distributed lifecycle semantics.
+- **Related Constraint**: #9
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/bagel/bagel.py b/src/flow_factory/models/bagel/bagel.py
index 678ec42bd..bef659b2b 100644
--- a/src/flow_factory/models/bagel/bagel.py
+++ b/src/flow_factory/models/bagel/bagel.py
@@ -814,7 +814,9 @@ def _update_context_text(self, text: Union[str, List[str]], gen_context: Dict) -
)
generation_input = move_tensors_to_device(generation_input, device, max_depth=1)
past_key_values = bagel.forward_cache_update_text(
- gen_context["past_key_values"], **generation_input
+ gen_context["past_key_values"],
+ language_model_forward=self.transformer,
+ **generation_input,
)
return {"kv_lens": kv_lens, "ropes": ropes, "past_key_values": past_key_values}
@@ -861,7 +863,10 @@ def _update_context_image(
)
gen_input = move_tensors_to_device(gen_input, device, max_depth=1)
past_key_values = bagel.forward_cache_update_vae(
- vae_model, past_key_values, **gen_input
+ vae_model,
+ past_key_values,
+ language_model_forward=self.transformer,
+ **gen_input,
)
if vit:
@@ -873,7 +878,11 @@ def _update_context_image(
new_token_ids=self.new_token_ids,
)
gen_input = move_tensors_to_device(gen_input, device, max_depth=1)
- past_key_values = bagel.forward_cache_update_vit(past_key_values, **gen_input)
+ past_key_values = bagel.forward_cache_update_vit(
+ past_key_values,
+ language_model_forward=self.transformer,
+ **gen_input,
+ )
return {"kv_lens": kv_lens, "ropes": ropes, "past_key_values": past_key_values}
@@ -951,14 +960,6 @@ def _forward_flow(
cfg_img_packed_key_value_indexes: Optional[torch.LongTensor] = None,
cfg_type: str = "parallel",
):
- packed_text_embedding = self.pipeline.transformer.model.embed_tokens(
- packed_text_ids
- ).float()
- packed_sequence = packed_text_embedding.new_zeros(
- (sum(packed_seqlens), self.pipeline.bagel.hidden_size), dtype=torch.float32
- )
- packed_sequence[packed_text_indexes] = packed_text_embedding
-
# ``x_t`` is the packed VAE-token tensor (sum_vae_tokens, patch_dim). A stray
# leading batch dim of 1 (callers passing (1, tokens, dim)) is squeezed off.
# ``timestep`` is one sigma per VAE token (expanded per sample upstream), so it
@@ -973,6 +974,9 @@ def _forward_flow(
packed_pos_embed = self.pipeline.bagel.latent_pos_embed(packed_vae_position_ids)
packed_timestep_embeds = self.pipeline.bagel.time_embedder(timestep)
x_t = self.pipeline.bagel.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed
+ packed_sequence = x_t.new_zeros(
+ (sum(packed_seqlens), self.pipeline.bagel.hidden_size), dtype=torch.float32
+ )
if x_t.dtype != packed_sequence.dtype:
x_t = x_t.to(packed_sequence.dtype)
packed_sequence[packed_vae_token_indexes] = x_t
@@ -982,7 +986,6 @@ def _forward_flow(
extra_inputs = {
"mode": "gen",
"packed_vae_token_indexes": packed_vae_token_indexes,
- "packed_text_indexes": packed_text_indexes,
}
output = self.transformer(
packed_query_sequence=packed_sequence,
@@ -994,6 +997,8 @@ def _forward_flow(
packed_key_value_indexes=packed_key_value_indexes,
update_past_key_values=False,
is_causal=False,
+ packed_text_ids=packed_text_ids,
+ packed_text_indexes=packed_text_indexes,
**extra_inputs,
)
v_t = self.pipeline.bagel.llm2vae(output.packed_query_sequence)
@@ -1009,6 +1014,8 @@ def _forward_flow(
packed_key_value_indexes=cfg_text_packed_key_value_indexes,
update_past_key_values=False,
is_causal=False,
+ packed_text_ids=packed_text_ids,
+ packed_text_indexes=packed_text_indexes,
**extra_inputs,
)
cfg_text_v_t = self.pipeline.bagel.llm2vae(cfg_text_output.packed_query_sequence)
@@ -1024,6 +1031,8 @@ def _forward_flow(
packed_key_value_indexes=cfg_img_packed_key_value_indexes,
update_past_key_values=False,
is_causal=False,
+ packed_text_ids=packed_text_ids,
+ packed_text_indexes=packed_text_indexes,
**extra_inputs,
)
cfg_img_v_t = self.pipeline.bagel.llm2vae(cfg_img_output.packed_query_sequence)
diff --git a/src/flow_factory/models/bagel/modeling/bagel/bagel.py b/src/flow_factory/models/bagel/modeling/bagel/bagel.py
index e7eedf3a1..06286fa99 100644
--- a/src/flow_factory/models/bagel/modeling/bagel/bagel.py
+++ b/src/flow_factory/models/bagel/modeling/bagel/bagel.py
@@ -2,7 +2,7 @@
# SPDX-License-Identifier: Apache-2.0
import copy
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
import torch
import torch.nn.functional as F
@@ -299,15 +299,19 @@ def forward_cache_update_text(
packed_text_indexes: torch.LongTensor,
packed_key_value_indexes: torch.LongTensor,
key_values_lens: torch.IntTensor,
+ language_model_forward: Optional[Callable[..., Any]] = None,
):
- packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
+ # Flow-Factory injects its prepared component route here. Standalone Bagel
+ # keeps the original physical language-model call as the default.
+ if language_model_forward is None:
+ language_model_forward = self.language_model
extra_inputs = {}
if self.use_moe:
extra_inputs = {"mode": "und"}
- output = self.language_model.forward_inference(
- packed_query_sequence=packed_text_embedding,
+ output = language_model_forward(
+ packed_query_sequence=None,
query_lens=text_token_lens,
packed_query_position_ids=packed_text_position_ids,
packed_query_indexes=packed_text_indexes,
@@ -316,6 +320,7 @@ def forward_cache_update_text(
key_values_lens=key_values_lens,
update_past_key_values=True,
is_causal=True,
+ packed_text_ids=packed_text_ids,
**extra_inputs,
)
past_key_values = output.past_key_values
@@ -412,10 +417,10 @@ def forward_cache_update_vit(
packed_indexes: torch.LongTensor,
packed_key_value_indexes: torch.LongTensor,
key_values_lens: torch.IntTensor,
+ language_model_forward: Optional[Callable[..., Any]] = None,
):
- packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
- packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
- packed_sequence[packed_text_indexes] = packed_text_embedding
+ if language_model_forward is None:
+ language_model_forward = self.language_model
cu_seqlens = torch.nn.functional.pad(torch.cumsum(vit_token_seqlens, dim=0), (1, 0))
cu_seqlens = cu_seqlens.to(torch.int32)
@@ -429,6 +434,9 @@ def forward_cache_update_vit(
packed_vit_token_embed = self.connector(packed_vit_token_embed)
pos_emb = self.vit_pos_embed(packed_vit_position_ids)
packed_vit_token_embed = packed_vit_token_embed + pos_emb
+ packed_sequence = packed_vit_token_embed.new_zeros(
+ (sum(packed_seqlens), self.hidden_size), dtype=self.language_model.dtype
+ )
if packed_vit_token_embed.dtype != packed_sequence.dtype:
packed_vit_token_embed = packed_vit_token_embed.to(packed_sequence.dtype)
packed_sequence[packed_vit_token_indexes] = packed_vit_token_embed
@@ -437,7 +445,7 @@ def forward_cache_update_vit(
if self.use_moe:
extra_inputs = {"mode": "und"}
- output = self.language_model.forward_inference(
+ output = language_model_forward(
packed_query_sequence=packed_sequence,
query_lens=packed_seqlens,
packed_query_position_ids=packed_position_ids,
@@ -447,6 +455,8 @@ def forward_cache_update_vit(
key_values_lens=key_values_lens,
update_past_key_values=True,
is_causal=False,
+ packed_text_ids=packed_text_ids,
+ packed_text_indexes=packed_text_indexes,
**extra_inputs,
)
past_key_values = output.past_key_values
@@ -558,10 +568,10 @@ def forward_cache_update_vae(
packed_indexes: torch.LongTensor,
key_values_lens: torch.IntTensor,
packed_key_value_indexes: torch.Tensor,
+ language_model_forward: Optional[Callable[..., Any]] = None,
):
- packed_text_embedding = self.language_model.model.embed_tokens(packed_text_ids)
- packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size))
- packed_sequence[packed_text_indexes] = packed_text_embedding
+ if language_model_forward is None:
+ language_model_forward = self.language_model
padded_latent = vae_model.encode(padded_images)
@@ -575,6 +585,9 @@ def forward_cache_update_vae(
packed_pos_embed = self.latent_pos_embed(packed_vae_position_ids)
packed_timestep_embeds = self.time_embedder(packed_timesteps)
packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + packed_pos_embed
+ packed_sequence = packed_latent.new_zeros(
+ (sum(packed_seqlens), self.hidden_size), dtype=self.language_model.dtype
+ )
if packed_latent.dtype != packed_sequence.dtype:
packed_latent = packed_latent.to(packed_sequence.dtype)
packed_sequence[packed_vae_token_indexes] = packed_latent
@@ -584,10 +597,9 @@ def forward_cache_update_vae(
extra_inputs = {
"mode": "gen",
"packed_vae_token_indexes": packed_vae_token_indexes,
- "packed_text_indexes": packed_text_indexes,
}
- output = self.language_model.forward_inference(
+ output = language_model_forward(
packed_query_sequence=packed_sequence,
query_lens=packed_seqlens,
packed_query_position_ids=packed_position_ids,
@@ -597,6 +609,8 @@ def forward_cache_update_vae(
packed_key_value_indexes=packed_key_value_indexes,
update_past_key_values=True,
is_causal=False,
+ packed_text_ids=packed_text_ids,
+ packed_text_indexes=packed_text_indexes,
**extra_inputs,
)
past_key_values = output.past_key_values
diff --git a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
index d81e96f3c..9d0ad9328 100644
--- a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
+++ b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
@@ -1312,7 +1312,7 @@ def forward_train(
def forward_inference(
self,
- packed_query_sequence: torch.Tensor,
+ packed_query_sequence: Optional[torch.Tensor],
query_lens: torch.Tensor,
packed_query_position_ids: torch.Tensor,
packed_query_indexes: torch.Tensor,
@@ -1324,8 +1324,39 @@ def forward_inference(
mode="und",
packed_vae_token_indexes=None,
packed_text_indexes=None,
+ packed_text_ids: Optional[torch.LongTensor] = None,
) -> BaseNavitOutputWithPast:
+ # Keep token embedding inside the outer language-model forward. Distributed
+ # wrappers attach their unshard hooks to this boundary, so callers must not
+ # reach through to ``model.embed_tokens`` while the parameters are sharded.
+ if packed_text_ids is not None:
+ packed_text_embedding = self.model.embed_tokens(packed_text_ids)
+ if packed_query_sequence is None:
+ packed_query_sequence = packed_text_embedding
+ else:
+ if packed_text_indexes is None:
+ raise ValueError(
+ "packed_text_indexes is required when inserting packed_text_ids "
+ "into an existing packed_query_sequence"
+ )
+ if packed_text_ids.numel() != packed_text_indexes.numel():
+ raise ValueError(
+ "packed_text_ids and packed_text_indexes must contain the same "
+ f"number of tokens, got {packed_text_ids.numel()} and "
+ f"{packed_text_indexes.numel()}"
+ )
+ packed_query_sequence = packed_query_sequence.index_copy(
+ 0,
+ packed_text_indexes,
+ packed_text_embedding.to(packed_query_sequence.dtype),
+ )
+ elif packed_query_sequence is None:
+ raise ValueError(
+ "Qwen2ForCausalLM.forward_inference requires packed_query_sequence "
+ "or packed_text_ids"
+ )
+
outputs = self.model(
packed_query_sequence=packed_query_sequence,
query_lens=query_lens,
diff --git a/tests/models/test_bagel_tdm_contracts.py b/tests/models/test_bagel_tdm_contracts.py
index 9d6bf1010..627f925ab 100644
--- a/tests/models/test_bagel_tdm_contracts.py
+++ b/tests/models/test_bagel_tdm_contracts.py
@@ -18,7 +18,7 @@
import importlib.machinery
import sys
import types
-from types import MethodType
+from types import MethodType, SimpleNamespace
from typing import Any
import pytest
@@ -27,7 +27,7 @@
from peft import LoraConfig, get_peft_model
import flow_factory.utils.imports as import_utils
-from flow_factory.models.model_bundle import ModelBundle
+from flow_factory.models.model_bundle import ModelBundle, RoutedComponentProxy
from flow_factory.samples import BaseSample
from flow_factory.trainers.distillation.distillation_runtime import (
validate_media_free_rollout,
@@ -58,6 +58,13 @@ def _load_bagel_qwen_types(monkeypatch: pytest.MonkeyPatch) -> tuple[type, type]
return module.Qwen2Config, module.Qwen2ForCausalLM
+def _load_bagel_model_type(monkeypatch: pytest.MonkeyPatch) -> type:
+ """Load the vendored Bagel container behind the optional-kernel seam."""
+ _load_bagel_types(monkeypatch)
+ module = importlib.import_module("flow_factory.models.bagel.modeling.bagel.bagel")
+ return module.Bagel
+
+
class _FakeRotaryEmbedding(torch.nn.Module):
def forward(
self,
@@ -115,6 +122,38 @@ def _checkpointing_qwen_model(
return model, decoder
+def _raw_token_qwen_proxy(
+ monkeypatch: pytest.MonkeyPatch,
+) -> tuple[RoutedComponentProxy, torch.nn.Module]:
+ """Build the real PEFT/bundle route with a lightweight decoder body."""
+ qwen_config_type, qwen_causal_lm_type = _load_bagel_qwen_types(monkeypatch)
+ config = qwen_config_type(
+ vocab_size=8,
+ hidden_size=8,
+ intermediate_size=16,
+ num_hidden_layers=1,
+ num_attention_heads=2,
+ num_key_value_heads=2,
+ max_position_embeddings=16,
+ layer_module="Qwen2DecoderLayer",
+ qk_norm=False,
+ _attn_implementation="eager",
+ pad_token_id=0,
+ )
+ base_model = qwen_causal_lm_type(config)
+ model = get_peft_model(
+ base_model,
+ LoraConfig(r=2, lora_alpha=2, target_modules=["q_proj"]),
+ )
+ routed_base = model.get_base_model()
+ routed_base.model.layers = torch.nn.ModuleList([_CheckpointedDecoderLayer()])
+ routed_base.model.rotary_emb = _FakeRotaryEmbedding()
+ routed_base.model.norm = torch.nn.Identity()
+ routed_base.model.use_moe = False
+ bundle = ModelBundle({"transformer": model})
+ return RoutedComponentProxy(bundle, "transformer", model), routed_base
+
+
def _adapter(adapter_type: type, decoder: Any) -> Any:
adapter = object.__new__(adapter_type)
adapter.decode_latents = MethodType(decoder, adapter)
@@ -184,6 +223,175 @@ def test_bagel_qwen_checkpointing_does_not_replay_cache_updates(
assert decoder.calls == 1
+def test_bagel_qwen_routes_raw_text_embedding_through_outer_forward(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ transformer, base_model = _raw_token_qwen_proxy(monkeypatch)
+ with torch.no_grad():
+ base_model.model.embed_tokens.weight.copy_(
+ torch.arange(64, dtype=torch.float32).reshape(8, 8) / 10
+ )
+
+ text_ids = torch.tensor([1, 2])
+ output = transformer(
+ packed_query_sequence=None,
+ query_lens=torch.tensor([2]),
+ packed_query_position_ids=torch.tensor([0, 1]),
+ packed_query_indexes=torch.tensor([4, 5]),
+ past_key_values=object(),
+ key_values_lens=torch.tensor([4]),
+ packed_key_value_indexes=torch.arange(4),
+ update_past_key_values=True,
+ is_causal=True,
+ packed_text_ids=text_ids,
+ # These are merged-cache indexes for text prefill, not query-local
+ # insertion indexes. The sequence=None path must ignore them.
+ packed_text_indexes=torch.tensor([4, 5]),
+ )
+
+ expected = base_model.model.embed_tokens(text_ids).square()
+ torch.testing.assert_close(output.packed_query_sequence, expected)
+
+
+def test_bagel_qwen_inserts_raw_text_without_detaching_aux_sequence(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ transformer, base_model = _raw_token_qwen_proxy(monkeypatch)
+ base_model.model.embed_tokens.weight.requires_grad_(True)
+ with torch.no_grad():
+ base_model.model.embed_tokens.weight.copy_(
+ torch.arange(64, dtype=torch.float32).reshape(8, 8) / 10
+ )
+
+ aux_sequence = torch.full((3, 8), 2.0, requires_grad=True)
+ text_ids = torch.tensor([1])
+ output = transformer(
+ packed_query_sequence=aux_sequence,
+ query_lens=torch.tensor([3]),
+ packed_query_position_ids=torch.tensor([0, 0, 0]),
+ packed_query_indexes=torch.tensor([0, 1, 2]),
+ past_key_values=object(),
+ key_values_lens=torch.tensor([0]),
+ packed_key_value_indexes=torch.tensor([], dtype=torch.long),
+ update_past_key_values=False,
+ is_causal=False,
+ packed_text_ids=text_ids,
+ packed_text_indexes=torch.tensor([1]),
+ )
+ output.packed_query_sequence.sum().backward()
+
+ expected_input = aux_sequence.detach().clone()
+ expected_input[1] = base_model.model.embed_tokens(text_ids).detach()[0]
+ torch.testing.assert_close(output.packed_query_sequence, expected_input.square())
+ torch.testing.assert_close(aux_sequence.grad[0], torch.full((8,), 4.0))
+ torch.testing.assert_close(aux_sequence.grad[1], torch.zeros(8))
+ torch.testing.assert_close(aux_sequence.grad[2], torch.full((8,), 4.0))
+ assert base_model.model.embed_tokens.weight.grad is not None
+ assert base_model.model.embed_tokens.weight.grad[1].abs().sum() > 0
+
+
+def test_bagel_text_cache_uses_injected_language_model_forward(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ bagel_type = _load_bagel_model_type(monkeypatch)
+ bagel = object.__new__(bagel_type)
+ torch.nn.Module.__init__(bagel)
+ bagel.use_moe = False
+
+ class ExplodingPhysicalLanguageModel(torch.nn.Module):
+ def forward(self, *args: Any, **kwargs: Any) -> Any:
+ del args, kwargs
+ raise AssertionError("the physical language model must not be called")
+
+ @property
+ def model(self) -> Any:
+ raise AssertionError("the physical embedding must not be accessed")
+
+ bagel.language_model = ExplodingPhysicalLanguageModel()
+ returned_cache = object()
+ calls: list[dict[str, Any]] = []
+
+ def routed_forward(**kwargs: Any) -> Any:
+ calls.append(kwargs)
+ return SimpleNamespace(past_key_values=returned_cache)
+
+ result = bagel.forward_cache_update_text(
+ past_key_values=object(),
+ packed_text_ids=torch.tensor([1, 2]),
+ packed_text_position_ids=torch.tensor([0, 1]),
+ text_token_lens=torch.tensor([2]),
+ packed_text_indexes=torch.tensor([0, 1]),
+ packed_key_value_indexes=torch.tensor([], dtype=torch.long),
+ key_values_lens=torch.tensor([0]),
+ language_model_forward=routed_forward,
+ )
+
+ assert result is returned_cache
+ assert len(calls) == 1
+ assert calls[0]["packed_query_sequence"] is None
+ assert torch.equal(calls[0]["packed_text_ids"], torch.tensor([1, 2]))
+
+
+def test_bagel_adapter_injects_prepared_route_into_all_cache_updates(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ adapter_type, _ = _load_bagel_types(monkeypatch)
+ routed_transformer = object()
+ forwarded_routes: list[Any] = []
+
+ class FakeBagel:
+ def prepare_prompts(self, **kwargs: Any) -> tuple[dict[str, Any], list[int], list[int]]:
+ del kwargs
+ return {}, [1], [1]
+
+ def forward_cache_update_text(
+ self, past_key_values: Any, *, language_model_forward: Any
+ ) -> Any:
+ del past_key_values
+ forwarded_routes.append(language_model_forward)
+ return "text-cache"
+
+ def prepare_vae_images(self, **kwargs: Any) -> tuple[dict[str, Any], list[int], list[int]]:
+ del kwargs
+ return {}, [2], [2]
+
+ def forward_cache_update_vae(
+ self, vae_model: Any, past_key_values: Any, *, language_model_forward: Any
+ ) -> Any:
+ del vae_model, past_key_values
+ forwarded_routes.append(language_model_forward)
+ return "vae-cache"
+
+ def prepare_vit_images(self, **kwargs: Any) -> tuple[dict[str, Any], list[int], list[int]]:
+ del kwargs
+ return {}, [3], [3]
+
+ def forward_cache_update_vit(
+ self, past_key_values: Any, *, language_model_forward: Any
+ ) -> Any:
+ del past_key_values
+ forwarded_routes.append(language_model_forward)
+ return "vit-cache"
+
+ holder = SimpleNamespace(
+ pipeline=SimpleNamespace(bagel=FakeBagel(), vae=object()),
+ device=torch.device("cpu"),
+ transformer=routed_transformer,
+ _tokenizer=object(),
+ new_token_ids={},
+ vae_transform=object(),
+ vit_transform=object(),
+ )
+ context = {"kv_lens": [0], "ropes": [0], "past_key_values": object()}
+
+ text_context = adapter_type._update_context_text(holder, ["prompt"], context)
+ image_context = adapter_type._update_context_image(holder, [torch.zeros(3, 8, 8)], context)
+
+ assert text_context["past_key_values"] == "text-cache"
+ assert image_context["past_key_values"] == "vit-cache"
+ assert forwarded_routes == [routed_transformer] * 3
+
+
@pytest.mark.parametrize(
"layer_module",
["Qwen2DecoderLayer", "Qwen2MoEDecoderLayer", "Qwen2MoTDecoderLayer"],
From 82018db367afe6867d0bed98538373b630efeee1 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:24:03 +0800
Subject: [PATCH 38/76] [trainer] fix: align FSDP2 checkpoint ownership
---
.agents/knowledge/topics/fix_patterns.md | 18 +++++
guidance/new_model.md | 15 +++--
src/flow_factory/trainers/abc.py | 35 +++++++---
tests/models/test_variant_checkpointing.py | 18 ++++-
.../test_distributed_plan_validation.py | 65 +++++++++++++++++--
5 files changed, 128 insertions(+), 23 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 1f2d9efb1..a54a5045e 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -395,6 +395,24 @@ Based on the fix type, write the fix entry to the appropriate document:
masks the first failure and breaks distributed lifecycle semantics.
- **Related Constraint**: #9
+### FSDP2 activation checkpoints must replay inside the mixed-precision boundary
+- **Date**: 2026-08-30
+- **Symptom**: All four Wan FSDP2 trainers failed on their first backward because checkpointed
+ tensors were saved as BF16 but recomputed as FP32.
+- **Root Cause**: Model-level checkpointing captured FP32 block inputs before FSDP2's forward-input
+ cast, while backward replay re-entered a block in `PRE_BACKWARD` state where PyTorch deliberately
+ skips that cast.
+- **Fix**: When full model checkpointing and FSDP2 activation checkpointing are both requested, the
+ trainer now disables model-level boundaries and keeps Accelerate's backend checkpoint wrappers,
+ which replay inside the fully-sharded mixed-precision boundary. Selective policies fail closed
+ because backend checkpointing cannot preserve their exact selection, while FSDP1 retains its
+ existing owner. Wan FSDP2 GRPO, TDM, SFT, and offline DPO plus SD3.5 and Bagel regressions verify
+ the shared path.
+- **Lesson**: Checkpoint placement is part of distributed precision semantics. A recompute boundary
+ outside a sharded module may not replay its forward hooks, so backend-aligned checkpoint wrappers
+ must own FSDP2 full checkpointing instead of nesting model-level boundaries around sharded blocks.
+- **Related Constraint**: #9, #20
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/guidance/new_model.md b/guidance/new_model.md
index 71367702d..8c86bd80c 100644
--- a/guidance/new_model.md
+++ b/guidance/new_model.md
@@ -69,12 +69,15 @@ Diffusers model's `_repeated_blocks` declaration. Adapters with multiple forward
stacks should override `_gradient_checkpointing_units()` and return their blocks
in execution order.
-Checkpointing has one owner. An explicit train-level model policy disables FSDP
-activation checkpointing to avoid nested recomputation; when the train-level
-policy is disabled, the backend may apply full FSDP activation checkpointing.
-Transformers-style components support full checkpointing through
-`gradient_checkpointing_enable()`, but must expose the Diffusers callback API to
-support selective modes.
+Checkpointing has one owner. When FSDP2 full model checkpointing and backend
+activation checkpointing are both enabled, the model policy yields ownership to
+the backend so recomputation stays inside the sharded mixed-precision boundary.
+FSDP1 keeps model-level ownership. FSDP2 rejects a selective train-level policy
+combined with backend activation checkpointing, because the backend cannot
+preserve the requested `fraction`, `every_n`, or `layers` boundary; disable
+backend activation checkpointing when using those policies. Transformers-style
+components support full checkpointing through `gradient_checkpointing_enable()`,
+but must expose the Diffusers callback API to support selective modes.
## Step-by-Step Implementation
diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py
index d84e9b08d..0ecb77551 100644
--- a/src/flow_factory/trainers/abc.py
+++ b/src/flow_factory/trainers/abc.py
@@ -1175,10 +1175,8 @@ def _apply_backend_checkpointing_constraints(self) -> None:
)
)
fsdp_checkpointing = bool(getattr(fsdp_plugin, "activation_checkpointing", False))
- if (
- self.training_args.trainer_type == "tdm-r1"
- and getattr(fsdp_plugin, "fsdp_version", 1) < 2
- ):
+ fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) or 1
+ if self.training_args.trainer_type == "tdm-r1" and fsdp_version < 2:
if not model_checkpointing and not fsdp_checkpointing:
return
self.adapter.disable_gradient_checkpointing()
@@ -1194,11 +1192,30 @@ def _apply_backend_checkpointing_constraints(self) -> None:
return
if model_checkpointing and fsdp_checkpointing:
- fsdp_plugin.activation_checkpointing = False
- logger.info(
- "Disabled FSDP activation checkpointing because train-level model "
- "checkpointing is enabled; nested checkpoint boundaries duplicate recompute."
- )
+ if fsdp_version >= 2:
+ checkpoint_policy = self.training_args.enable_gradient_checkpointing
+ full_checkpointing = checkpoint_policy is True or (
+ getattr(checkpoint_policy, "mode", None) == "full"
+ )
+ if not full_checkpointing:
+ raise ValueError(
+ "FSDP2 activation checkpointing cannot preserve selective model "
+ "checkpointing boundaries. Disable fsdp_activation_checkpointing or "
+ "use train.enable_gradient_checkpointing=true/mode=full."
+ )
+ self.adapter.disable_gradient_checkpointing()
+ self.training_args.enable_gradient_checkpointing = False
+ logger.info(
+ "Disabled model gradient checkpointing because FSDP2 activation "
+ "checkpointing is enabled; checkpoint recomputation must stay inside "
+ "the FSDP2 mixed-precision boundary."
+ )
+ else:
+ fsdp_plugin.activation_checkpointing = False
+ logger.info(
+ "Disabled FSDP activation checkpointing because train-level model "
+ "checkpointing is enabled; nested checkpoint boundaries duplicate recompute."
+ )
def _initialization(self):
self._validate_paradigm_dynamics()
diff --git a/tests/models/test_variant_checkpointing.py b/tests/models/test_variant_checkpointing.py
index 3540c334f..797543267 100644
--- a/tests/models/test_variant_checkpointing.py
+++ b/tests/models/test_variant_checkpointing.py
@@ -27,9 +27,7 @@
from flow_factory.models.abc import BaseAdapter
from flow_factory.models.model_bundle import ModelBundle, RoutedComponentProxy
from flow_factory.models.variants import DEFAULT_BASE_VARIANT as BASE_VARIANT
-from flow_factory.models.variants import (
- ComponentVariantRegistry,
-)
+from flow_factory.models.variants import ComponentVariantRegistry
from flow_factory.trainers.abc import BaseTrainer
from flow_factory.trainers.common.runtime_state import TrainerRuntimeState
from flow_factory.trainers.distillation.tdm_r1 import (
@@ -257,6 +255,20 @@ def _trainer_runtime(
return trainer
+def test_disabling_gradient_checkpointing_visits_every_materialized_variant() -> None:
+ trainer = _trainer_runtime("full")
+ members = trainer.adapter.component_variant_registry.bundle_members()
+ disabled_routes = []
+ for route_name, component in members.items():
+ component.disable_gradient_checkpointing = (
+ lambda route_name=route_name: disabled_routes.append(route_name)
+ )
+
+ trainer.adapter.disable_gradient_checkpointing()
+
+ assert disabled_routes == list(members)
+
+
def _step_fake_role(trainer: TinyTrainer) -> None:
coordinator = trainer.role_optimization
fake_parameter = trainer.optimization_roles["fake"].parameters[0]
diff --git a/tests/trainers/test_distributed_plan_validation.py b/tests/trainers/test_distributed_plan_validation.py
index 28cf3ded9..0766c89f3 100644
--- a/tests/trainers/test_distributed_plan_validation.py
+++ b/tests/trainers/test_distributed_plan_validation.py
@@ -233,9 +233,7 @@ def test_tdm_r1_fsdp1_disables_incompatible_activation_checkpointing() -> None:
trainer_type="tdm-r1",
enable_gradient_checkpointing=True,
),
- adapter=SimpleNamespace(
- disable_gradient_checkpointing=lambda: disabled.append(True)
- ),
+ adapter=SimpleNamespace(disable_gradient_checkpointing=lambda: disabled.append(True)),
)
BaseTrainer._apply_backend_checkpointing_constraints(trainer)
@@ -245,10 +243,17 @@ def test_tdm_r1_fsdp1_disables_incompatible_activation_checkpointing() -> None:
assert plugin.activation_checkpointing is False
-def test_fsdp2_keeps_model_checkpointing_and_disables_nested_backend_checkpointing() -> None:
+@pytest.mark.parametrize(
+ "checkpoint_policy",
+ [True, SimpleNamespace(mode="full")],
+)
+def test_fsdp2_disables_model_checkpointing_and_keeps_backend_checkpointing(
+ checkpoint_policy: object,
+) -> None:
from flow_factory.trainers.abc import BaseTrainer
plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
+ disabled = []
trainer = SimpleNamespace(
accelerator=SimpleNamespace(
distributed_type=DistributedType.FSDP,
@@ -256,11 +261,34 @@ def test_fsdp2_keeps_model_checkpointing_and_disables_nested_backend_checkpointi
),
training_args=SimpleNamespace(
trainer_type="tdm-r1",
+ enable_gradient_checkpointing=checkpoint_policy,
+ ),
+ adapter=SimpleNamespace(disable_gradient_checkpointing=lambda: disabled.append(True)),
+ )
+
+ BaseTrainer._apply_backend_checkpointing_constraints(trainer)
+
+ assert disabled == [True]
+ assert trainer.training_args.enable_gradient_checkpointing is False
+ assert plugin.activation_checkpointing is True
+
+
+def test_fsdp1_keeps_model_checkpointing_and_disables_nested_backend_checkpointing() -> None:
+ from flow_factory.trainers.abc import BaseTrainer
+
+ plugin = SimpleNamespace(fsdp_version=1, activation_checkpointing=True)
+ trainer = SimpleNamespace(
+ accelerator=SimpleNamespace(
+ distributed_type=DistributedType.FSDP,
+ state=SimpleNamespace(fsdp_plugin=plugin),
+ ),
+ training_args=SimpleNamespace(
+ trainer_type="grpo",
enable_gradient_checkpointing=True,
),
adapter=SimpleNamespace(
disable_gradient_checkpointing=lambda: pytest.fail(
- "FSDP2 must keep model checkpointing enabled"
+ "FSDP1 must keep model checkpointing enabled"
)
),
)
@@ -271,6 +299,33 @@ def test_fsdp2_keeps_model_checkpointing_and_disables_nested_backend_checkpointi
assert plugin.activation_checkpointing is False
+def test_fsdp2_rejects_selective_model_and_backend_checkpointing() -> None:
+ from flow_factory.trainers.abc import BaseTrainer
+
+ plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
+ trainer = SimpleNamespace(
+ accelerator=SimpleNamespace(
+ distributed_type=DistributedType.FSDP,
+ state=SimpleNamespace(fsdp_plugin=plugin),
+ ),
+ training_args=SimpleNamespace(
+ trainer_type="grpo",
+ enable_gradient_checkpointing=SimpleNamespace(mode="every_n"),
+ ),
+ adapter=SimpleNamespace(
+ disable_gradient_checkpointing=lambda: pytest.fail(
+ "selective checkpointing must fail before mutating the adapter"
+ )
+ ),
+ )
+
+ with pytest.raises(ValueError, match="cannot preserve selective"):
+ BaseTrainer._apply_backend_checkpointing_constraints(trainer)
+
+ assert trainer.training_args.enable_gradient_checkpointing.mode == "every_n"
+ assert plugin.activation_checkpointing is True
+
+
def test_fsdp2_keeps_backend_checkpointing_when_model_policy_is_disabled() -> None:
from flow_factory.trainers.abc import BaseTrainer
From 7b3295d388790f9d801127d9052ae53e0110149e Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:34:05 +0800
Subject: [PATCH 39/76] [models] fix: allow optional dtype manifest selectors
---
.agents/knowledge/topics/fix_patterns.md | 17 ++++++++
src/flow_factory/models/abc.py | 12 ++++++
src/flow_factory/models/precision.py | 12 +++++-
tests/models/test_model_precision_policy.py | 45 ++++++++++++++++++++-
4 files changed, 82 insertions(+), 4 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index a54a5045e..dd414c97a 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -413,6 +413,23 @@ Based on the fix type, write the fix entry to the appropriate document:
must own FSDP2 full checkpointing instead of nesting model-level boundaries around sharded blocks.
- **Related Constraint**: #9, #20
+### Adapter dtype manifests may span optional checkpoint components
+- **Date**: 2026-08-30
+- **Symptom**: Every Wan2.2 TI2V trainer rejected the Wan I2V adapter's `image_encoder` load-dtype
+ default even though the checkpoint legitimately omits that optional component.
+- **Root Cause**: The eager loader validated an adapter-wide dtype manifest only against components
+ present in one checkpoint, conflating the checkpoint instance with the pipeline class's wider
+ optional-component contract.
+- **Fix**: Eager pipeline loading now validates adapter manifest selectors against the union of the
+ checkpoint components and the pipeline class's declared optional components, while resolving
+ dtype arguments only for components actually present. User overrides remain strict against the
+ selected checkpoint. Regressions cover absent and present optional components, invalid manifest
+ selectors, and explicit user selection of an absent component.
+- **Lesson**: Adapter defaults may intentionally cover several checkpoint variants of one pipeline
+ class. Optional class-level declarations belong to manifest validation, but they must not create
+ components or weaken the fail-fast contract for checkpoint-specific user overrides.
+- **Related Constraint**: #20
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index d110e78a1..e2abfef2d 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -963,12 +963,24 @@ def _load_diffusers_pipeline(
for name, value in pipeline_config.items()
if isinstance(value, (list, tuple)) and len(value) >= 2
]
+ # Adapter defaults may cover several checkpoint variants of one pipeline
+ # class. Keep absent class-declared optional components valid for manifest
+ # validation, but resolve the actual dtype mapping only for this checkpoint.
+ manifest_declared_names = list(
+ dict.fromkeys(
+ [
+ *component_names,
+ *getattr(pipeline_class, "_optional_components", ()),
+ ]
+ )
+ )
load_dtype_kwargs = build_component_load_dtype_kwargs(
user_policy=user_policy,
manifest_policy=manifest_policy,
component_names=component_names,
transformer_names=[name for name in component_names if "transformer" in name],
text_encoder_names=[name for name in component_names if "text_encoder" in name],
+ manifest_declared_names=manifest_declared_names,
preserve_unselected=True,
)
kwargs.update(
diff --git a/src/flow_factory/models/precision.py b/src/flow_factory/models/precision.py
index e75a133be..4537b7fe8 100644
--- a/src/flow_factory/models/precision.py
+++ b/src/flow_factory/models/precision.py
@@ -123,10 +123,16 @@ def component_dtype_mapping(
component_names: Sequence[str],
transformer_names: Sequence[str],
text_encoder_names: Sequence[str],
+ manifest_declared_names: Sequence[str] | None = None,
) -> dict[str, torch.dtype]:
- """Resolve a policy to the concrete non-null mapping accepted by loaders."""
+ """Resolve policies for concrete components, with wider adapter-manifest declarations."""
validate_dtype_policy_selectors(user_policy, declared_names=component_names)
- validate_dtype_policy_selectors(manifest_policy, declared_names=component_names)
+ validate_dtype_policy_selectors(
+ manifest_policy,
+ declared_names=(
+ component_names if manifest_declared_names is None else manifest_declared_names
+ ),
+ )
return {
name: dtype
for name in component_names
@@ -150,6 +156,7 @@ def build_component_load_dtype_kwargs(
component_names: Sequence[str],
transformer_names: Sequence[str],
text_encoder_names: Sequence[str],
+ manifest_declared_names: Sequence[str] | None = None,
requested_names: Sequence[str] | None = None,
preserve_unselected: bool = False,
) -> Dict[str, object]:
@@ -165,6 +172,7 @@ def build_component_load_dtype_kwargs(
component_names=component_names,
transformer_names=transformer_names,
text_encoder_names=text_encoder_names,
+ manifest_declared_names=manifest_declared_names,
)
if requested_names is not None:
requested = set(requested_names)
diff --git a/tests/models/test_model_precision_policy.py b/tests/models/test_model_precision_policy.py
index 2047d375e..91036c697 100644
--- a/tests/models/test_model_precision_policy.py
+++ b/tests/models/test_model_precision_policy.py
@@ -14,8 +14,8 @@
import pytest
import torch
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline
from flow_factory.models.abc import BaseAdapter
from flow_factory.models.precision import (
cast_module_role_dtypes,
@@ -44,6 +44,7 @@ def __init__(self) -> None:
class _PipelineFake(DiffusionPipeline):
load_call = None
+ _optional_components = ["image_encoder"]
@classmethod
def load_config(cls, pretrained_model_name_or_path: str, **kwargs):
@@ -86,6 +87,17 @@ def test_user_default_null_disables_concrete_manifest_defaults() -> None:
)
+def test_present_optional_manifest_selector_resolves_its_dtype() -> None:
+ assert component_dtype_mapping(
+ user_policy=None,
+ manifest_policy={"image_encoder": torch.float32},
+ component_names=["transformer", "image_encoder"],
+ transformer_names=["transformer"],
+ text_encoder_names=[],
+ manifest_declared_names=["transformer", "image_encoder"],
+ ) == {"image_encoder": torch.float32}
+
+
def test_transformers_group_includes_reference_transformer() -> None:
assert component_dtype_mapping(
user_policy={"transformers": torch.bfloat16},
@@ -101,7 +113,10 @@ def test_eager_pipeline_loader_expands_role_selectors() -> None:
"AdapterStub",
(),
{
- "_component_load_dtype_manifest": {"transformers": torch.bfloat16},
+ "_component_load_dtype_manifest": {
+ "transformers": torch.bfloat16,
+ "image_encoder": torch.float32,
+ },
"_component_load_dtype_overrides": None,
"_resolve_component_load_dtype_mapping": (
BaseAdapter._resolve_component_load_dtype_mapping
@@ -130,6 +145,20 @@ def test_eager_pipeline_loader_expands_role_selectors() -> None:
)
+def test_user_policy_cannot_select_an_absent_optional_component() -> None:
+ adapter = type(
+ "AdapterStub",
+ (),
+ {
+ "_component_load_dtype_manifest": None,
+ "_component_load_dtype_overrides": {"image_encoder": torch.float32},
+ },
+ )()
+
+ with pytest.raises(ValueError, match=r"unknown=.*image_encoder"):
+ BaseAdapter._load_diffusers_pipeline(adapter, _PipelineFake, "model")
+
+
def test_unknown_load_policy_selector_fails_with_runtime_context() -> None:
with pytest.raises(ValueError, match=r"unknown=.*missing.*declared=.*transformer"):
component_dtype_mapping(
@@ -141,6 +170,18 @@ def test_unknown_load_policy_selector_fails_with_runtime_context() -> None:
)
+def test_unknown_manifest_selector_remains_strict_with_optional_declarations() -> None:
+ with pytest.raises(ValueError, match=r"unknown=.*missing"):
+ component_dtype_mapping(
+ user_policy=None,
+ manifest_policy={"missing": torch.bfloat16},
+ component_names=["transformer"],
+ transformer_names=["transformer"],
+ text_encoder_names=[],
+ manifest_declared_names=["transformer", "image_encoder"],
+ )
+
+
def test_role_cast_preserves_diffusers_fp32_islands_but_not_lora_dtype() -> None:
model = _ProtectedDiffusersModel()
model.requires_grad_(False)
From 589ae026023fb127897c90ac75ffa8bf2a6d82ad Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:40:36 +0800
Subject: [PATCH 40/76] [models,loading] fix: separate absent component roots
---
.agents/knowledge/topics/fix_patterns.md | 16 ++++++
src/flow_factory/loading/coordinator.py | 13 ++++-
src/flow_factory/models/runtime/classic.py | 2 +
tests/loading/test_coordinator.py | 61 +++++++++++++++++++++-
4 files changed, 89 insertions(+), 3 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index dd414c97a..2758caa57 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -430,6 +430,22 @@ Based on the fix type, write the fix entry to the appropriate document:
components or weaken the fail-fast contract for checkpoint-specific user overrides.
- **Related Constraint**: #20
+### Absent optional components need distinct physical roots
+- **Date**: 2026-08-30
+- **Symptom**: Every Wan2.2 TI2V trainer rejected its load plan because the physical
+ `image_encoder` root appeared to combine incompatible auxiliary and host roles.
+- **Root Cause**: The eager runtime used object identity to collapse logical aliases, so multiple
+ optional components whose value was the singleton `None` were mistaken for one shared physical
+ object.
+- **Fix**: Classic pipelines now preserve a declared optional `None` under its own logical root
+ while retaining identity aliasing for real objects. The load coordinator finalizes only
+ replicated roots that actually materialized, preventing FSDP replica checks from resolving an
+ allowed absent component.
+- **Lesson**: Object identity establishes physical aliasing only for materialized objects. Optional
+ declarations retain distinct lifecycle identities, and backend finalization must follow observed
+ materialization rather than the requested name set.
+- **Related Constraint**: #9
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/loading/coordinator.py b/src/flow_factory/loading/coordinator.py
index 770f8327e..bfec60399 100644
--- a/src/flow_factory/loading/coordinator.py
+++ b/src/flow_factory/loading/coordinator.py
@@ -113,9 +113,18 @@ def load_components(
)
if not replicated and not remainder_roots:
return
+ loaded_replicated: List[str] = []
with self.load_scope(ComponentRole.AUXILIARY):
if replicated:
self.adapter.on_load_components(components=replicated, device=device)
+ materialized_roots = set(
+ self.adapter.component_runtime.materialized_component_names
+ )
+ loaded_replicated = [
+ name
+ for name in replicated
+ if self.plan.descriptors[name].root in materialized_roots
+ ]
for root in remainder_roots:
request = self.plan.request_for_root(root)
excluded_paths = [
@@ -128,8 +137,8 @@ def load_components(
excluded_paths=excluded_paths,
device=device,
)
- if replicated:
- self.components_loaded(replicated)
+ if loaded_replicated:
+ self.components_loaded(loaded_replicated)
def prepare(self, *objects: Any) -> Any:
return self.backend.prepare(*objects)
diff --git a/src/flow_factory/models/runtime/classic.py b/src/flow_factory/models/runtime/classic.py
index 5ea7fca7e..897d35a8c 100644
--- a/src/flow_factory/models/runtime/classic.py
+++ b/src/flow_factory/models/runtime/classic.py
@@ -56,6 +56,8 @@ def physical_route(self, name: str) -> tuple[str, tuple[str, ...]]:
"""Collapse logical aliases that reference one canonical module object."""
self._validate_declared_names([name])
component = self.declared_components[name]
+ if component is None:
+ return name, ()
for canonical_name, canonical_component in self.canonical_components.items():
if component is canonical_component:
return canonical_name, ()
diff --git a/tests/loading/test_coordinator.py b/tests/loading/test_coordinator.py
index db2977c36..1de700ed8 100644
--- a/tests/loading/test_coordinator.py
+++ b/tests/loading/test_coordinator.py
@@ -45,6 +45,27 @@ def test_classic_runtime_collapses_same_object_aliases() -> None:
assert runtime.physical_route("text_encoder") == ("text_encoder_2", ())
+def test_classic_runtime_keeps_absent_optional_components_as_distinct_roots() -> None:
+ pipeline = SimpleNamespace(
+ components={"image_encoder": None, "image_processor": None},
+ image_encoder=None,
+ image_processor=None,
+ )
+ runtime = ClassicPipelineRuntime(pipeline)
+ adapter = SimpleNamespace(
+ component_runtime=runtime,
+ model_args=SimpleNamespace(target_components=[]),
+ _resolve_component_names=runtime.resolve_component_names,
+ )
+
+ plan = build_adapter_load_plan(adapter)
+
+ assert runtime.physical_route("image_encoder") == ("image_encoder", ())
+ assert runtime.physical_route("image_processor") == ("image_processor", ())
+ assert plan.request_for_component("image_encoder").role is ComponentRole.AUXILIARY
+ assert plan.request_for_component("image_processor").role is ComponentRole.HOST
+
+
def test_coordinator_expands_target_component_groups() -> None:
transformer = torch.nn.Linear(2, 2)
runtime = PseudoPipelineRuntime(
@@ -122,9 +143,10 @@ def test_coordinator_moves_only_auxiliary_remainder_of_target_owned_root() -> No
_resolve_component_names=lambda components: list(components),
on_load_components=lambda components, device: calls.append(("load", components, device)),
component_runtime=SimpleNamespace(
+ materialized_component_names=["bagel", "vae"],
load_root_remainder=lambda root, excluded_paths, device: calls.append(
("remainder", root, excluded_paths, device)
- )
+ ),
),
)
coordinator.load_scope = lambda role: nullcontext()
@@ -147,6 +169,43 @@ def test_coordinator_moves_only_auxiliary_remainder_of_target_owned_root() -> No
]
+def test_coordinator_finalizes_only_materialized_replicas() -> None:
+ plan = LoadPlanner().build(
+ [
+ ComponentDescriptor(
+ name="image_encoder",
+ root="image_encoder",
+ role=ComponentRole.AUXILIARY,
+ ),
+ ComponentDescriptor(
+ name="vae",
+ root="vae",
+ role=ComponentRole.AUXILIARY,
+ ),
+ ]
+ )
+ calls = []
+ coordinator = object.__new__(ModelLoadCoordinator)
+ coordinator.plan = plan
+ coordinator.adapter = SimpleNamespace(
+ _resolve_component_names=lambda components: list(components),
+ on_load_components=lambda components, device: calls.append(("load", components, device)),
+ component_runtime=SimpleNamespace(materialized_component_names=["vae"]),
+ )
+ coordinator.load_scope = lambda role: nullcontext()
+ coordinator.components_loaded = lambda components: calls.append(("finalize", components))
+
+ coordinator.load_components(
+ ["image_encoder", "vae"],
+ device=torch.device("cpu"),
+ )
+
+ assert calls == [
+ ("load", ["image_encoder", "vae"], torch.device("cpu")),
+ ("finalize", ["vae"]),
+ ]
+
+
def test_pseudo_runtime_does_not_move_excluded_target_submodule() -> None:
class TrackingModule(torch.nn.Module):
def __init__(self):
From e6fcb360557f9bb1400b31572fc79780442bbb44 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:45:14 +0800
Subject: [PATCH 41/76] [dependencies] fix: install Wan prompt normalizer
---
.agents/knowledge/topics/fix_patterns.md | 14 ++++++++++++++
pyproject.toml | 3 ++-
tests/dependencies/test_diffusers_pin.py | 12 ++++++++++++
3 files changed, 28 insertions(+), 1 deletion(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 2758caa57..9b9b2c982 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -446,6 +446,20 @@ Based on the fix type, write the fix entry to the appropriate document:
materialization rather than the requested name set.
- **Related Constraint**: #9
+### Supported adapter paths require their upstream optional dependencies
+- **Date**: 2026-08-30
+- **Symptom**: Every Wan I2V trainer reached prompt preprocessing and then failed with
+ `NameError: name 'ftfy' is not defined` inside Diffusers prompt normalization.
+- **Root Cause**: Diffusers imports `ftfy` conditionally and declares it only in development/test
+ extras, while its Wan I2V prompt helper calls the package unconditionally. Flow-Factory exposes
+ Wan I2V as a core adapter but did not close that runtime dependency gap.
+- **Fix**: `ftfy` is now a core project dependency, with a metadata regression that keeps exactly
+ one install requirement for Wan prompt normalization.
+- **Lesson**: A framework that promotes an upstream optional code path to a supported core feature
+ also owns the path's transitive runtime dependencies; successful import of the upstream module
+ does not prove its conditionally imported helpers are callable.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/pyproject.toml b/pyproject.toml
index 6bdd6a172..7957a5024 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,6 +41,7 @@ dependencies = [
# Critical for T5/Flux
"protobuf>=6.33.2",
"sentencepiece>=0.2.1",
+ "ftfy", # Required by Diffusers Wan prompt normalization
# Basic Utils
"av>=18.0.0",
@@ -118,4 +119,4 @@ target-version = ['py310', 'py311', 'py312']
[tool.isort]
profile = "black"
-line_length = 100
\ No newline at end of file
+line_length = 100
diff --git a/tests/dependencies/test_diffusers_pin.py b/tests/dependencies/test_diffusers_pin.py
index 4faa9e982..c48156c78 100644
--- a/tests/dependencies/test_diffusers_pin.py
+++ b/tests/dependencies/test_diffusers_pin.py
@@ -17,6 +17,7 @@
ROOT = Path(__file__).resolve().parents[2]
EXPECTED_DIFFUSERS_REQUIREMENT = "diffusers>=0.40.0"
+EXPECTED_WAN_PROMPT_DEPENDENCY = "ftfy"
def test_project_metadata_requires_released_diffusers_with_h3_support() -> None:
@@ -28,3 +29,14 @@ def test_project_metadata_requires_released_diffusers_with_h3_support() -> None:
f"MiniMax H3 support; expected={EXPECTED_DIFFUSERS_REQUIREMENT!r}, "
f"observed={diffusers_requirements!r}"
)
+
+
+def test_project_metadata_requires_wan_prompt_normalization_dependency() -> None:
+ pyproject = (ROOT / "pyproject.toml").read_text(encoding="utf-8")
+ ftfy_requirements = re.findall(r'"(ftfy[^"]*)"', pyproject)
+
+ assert ftfy_requirements == [EXPECTED_WAN_PROMPT_DEPENDENCY], (
+ "pyproject.toml must contain exactly one core ftfy requirement for Wan prompt "
+ f"normalization; expected={EXPECTED_WAN_PROMPT_DEPENDENCY!r}, "
+ f"observed={ftfy_requirements!r}"
+ )
From a1e01fc0893ec603be60ea9be81de710365f881b Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:51:48 +0800
Subject: [PATCH 42/76] [models] fix: validate Wan condition temporal geometry
---
.agents/knowledge/topics/fix_patterns.md | 15 +++++++++++++++
src/flow_factory/models/wan/_conditioning.py | 3 ++-
tests/models/test_wan_output_codec.py | 14 ++++++++++++--
3 files changed, 29 insertions(+), 3 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 9b9b2c982..e35b4824c 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -460,6 +460,21 @@ Based on the fix type, write the fix entry to the appropriate document:
does not prove its conditionally imported helpers are callable.
- **Related Constraint**: N/A
+### Condition-latent geometry follows the encoded source, not the rollout target
+- **Date**: 2026-08-30
+- **Symptom**: Wan2.2 TI2V rejected a one-frame VAE condition latent with temporal size one because
+ the rollout noise and target video had temporal latent size two.
+- **Root Cause**: The condition validator derived its expected temporal size from configured output
+ frames. In the expanded-timestep pipeline, Diffusers intentionally encodes only the first input
+ frame and broadcasts that one-frame condition through a full-length first-frame mask.
+- **Fix**: Wan condition validation now derives temporal geometry from the actual video tensor sent
+ to the VAE. The regression models one-frame encoding and proves the official broadcast produces
+ the full rollout shape while non-expanded conditions keep their full temporal encoding.
+- **Lesson**: Conditioning and generated states can share channels and spatial geometry without
+ sharing sequence length. Validate each representation against its own source transform before
+ relying on an explicitly defined broadcast or mask contract.
+- **Related Constraint**: #7
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/wan/_conditioning.py b/src/flow_factory/models/wan/_conditioning.py
index 71556ab1d..47503ab55 100644
--- a/src/flow_factory/models/wan/_conditioning.py
+++ b/src/flow_factory/models/wan/_conditioning.py
@@ -319,10 +319,11 @@ def prepare_wan_i2v_condition_tensors(
).to(device=device, dtype=dtype)
latent_condition = normalize_wan_video_latents(adapter, latent_condition)
+ condition_latent_frames = (video_condition.shape[2] - 1) // temporal_scale + 1
expected_latents = (
batch_size,
vae.config.z_dim,
- num_latent_frames,
+ condition_latent_frames,
latent_height,
latent_width,
)
diff --git a/tests/models/test_wan_output_codec.py b/tests/models/test_wan_output_codec.py
index e3dab0ae4..ab0a935f4 100644
--- a/tests/models/test_wan_output_codec.py
+++ b/tests/models/test_wan_output_codec.py
@@ -72,11 +72,15 @@ def __init__(self) -> None:
latents_std=[2.0, 4.0, 5.0],
)
raw_channels = torch.tensor([3.0, 6.0, 8.0]).view(1, 3, 1, 1, 1)
- self.posterior = _Posterior(raw_channels.expand(1, 3, 2, 2, 2).clone())
+ self.full_latents = raw_channels.expand(1, 3, 2, 2, 2).clone()
+ self.posterior = _Posterior(self.full_latents)
self.encoded_pixels: list[torch.Tensor] = []
def encode(self, pixels: torch.Tensor) -> Any:
self.encoded_pixels.append(pixels)
+ self.posterior.latents = (
+ self.full_latents[:, :, :1] if pixels.shape[2] == 1 else self.full_latents
+ )
return SimpleNamespace(latent_dist=self.posterior)
@@ -384,8 +388,14 @@ def test_wan_i2v_expand_condition_binds_target_active_mask() -> None:
condition = prepared.forward_context["latent_condition"]
mask = prepared.forward_context["first_frame_mask"]
- assert condition.shape == (1, 3, 2, 2, 2)
+ assert condition.shape == (1, 3, 1, 2, 2)
assert mask.shape == (1, 1, 2, 2, 2)
+ assert adapter.vae.encoded_pixels[0].shape[2] == 1
+ rollout_latents = torch.full((1, 3, 2, 2, 2), 2.0)
+ blended = (1 - mask) * condition + mask * rollout_latents
+ assert blended.shape == rollout_latents.shape
+ torch.testing.assert_close(blended[:, :, 0], condition[:, :, 0])
+ torch.testing.assert_close(blended[:, :, 1], rollout_latents[:, :, 1])
torch.testing.assert_close(mask[:, :, 0], torch.zeros(1, 1, 2, 2))
torch.testing.assert_close(mask[:, :, 1], torch.ones(1, 1, 2, 2))
assert prepared.output_context["first_frame_mask"] is mask
From 140a7a3914cc45c92aeb1e6c7b6925d3f8ed1bd9 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 20:56:22 +0800
Subject: [PATCH 43/76] [models] fix: normalize Wan sample condition rows
---
.agents/knowledge/topics/fix_patterns.md | 14 ++++++++++++++
src/flow_factory/models/wan/wan2_i2v.py | 2 +-
tests/models/test_wan_output_codec.py | 17 +++++++++++++++--
3 files changed, 30 insertions(+), 3 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index e35b4824c..4de65f872 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -475,6 +475,20 @@ Based on the fix type, write the fix entry to the appropriate document:
relying on an explicitly defined broadcast or mask contract.
- **Related Constraint**: #7
+### Internal immutable media rows must cross sample boundaries as public batch types
+- **Date**: 2026-08-30
+- **Symptom**: Wan I2V rollouts finished denoising but failed while constructing each sample because
+ image canonicalization rejected a tuple of ordered first/last frames.
+- **Root Cause**: Wan's internal condition normalizer intentionally returns immutable tuple rows,
+ while `ImageConditionSample.condition_images` follows the public `ImageBatch` contract of lists,
+ tensors, or arrays. The adapter passed the internal representation across that boundary unchanged.
+- **Fix**: Wan now converts each ordered condition row to a list at sample construction. The
+ regression proves first/last color order survives sample canonicalization and replay stacking.
+- **Lesson**: Model-internal containers may enforce stronger invariants than shared sample APIs, but
+ adapters must translate them explicitly at the ownership boundary instead of broadening a common
+ media utility for one private representation.
+- **Related Constraint**: #5
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/wan/wan2_i2v.py b/src/flow_factory/models/wan/wan2_i2v.py
index 297d97c59..f7049c21b 100644
--- a/src/flow_factory/models/wan/wan2_i2v.py
+++ b/src/flow_factory/models/wan/wan2_i2v.py
@@ -769,7 +769,7 @@ def inference(
height=height,
width=width,
# Conditions
- condition_images=condition_rows[b],
+ condition_images=list(condition_rows[b]),
latent_condition=condition[b],
first_frame_mask=(first_frame_mask[b] if first_frame_mask is not None else None),
image_embeds=per_sample_image_embeds[b],
diff --git a/tests/models/test_wan_output_codec.py b/tests/models/test_wan_output_codec.py
index ab0a935f4..4de48de78 100644
--- a/tests/models/test_wan_output_codec.py
+++ b/tests/models/test_wan_output_codec.py
@@ -442,12 +442,25 @@ def test_wan_i2v_expand_target_requires_prepared_active_mask_before_vae() -> Non
def test_wan_i2v_normalization_never_truncates_optional_last_frame() -> None:
- first = Image.new("RGB", (4, 4))
- last = Image.new("RGB", (4, 4))
+ first = Image.new("RGB", (4, 4), color="red")
+ last = Image.new("RGB", (4, 4), color="blue")
rows = normalize_wan_i2v_image_rows([[first, last]], expected_batch_size=1)
+ sample = WanI2VSample(condition_images=list(rows[0]))
+ stacked = WanI2VSample.stack([sample])
assert rows == ((first, last),)
+ assert len(sample.condition_images) == 2
+ torch.testing.assert_close(sample.condition_images[0][:, 0, 0], torch.tensor([1.0, 0.0, 0.0]))
+ torch.testing.assert_close(sample.condition_images[1][:, 0, 0], torch.tensor([0.0, 0.0, 1.0]))
+ assert len(stacked["condition_images"]) == 1
+ assert len(stacked["condition_images"][0]) == 2
+ torch.testing.assert_close(
+ stacked["condition_images"][0][0][:, 0, 0], torch.tensor([1.0, 0.0, 0.0])
+ )
+ torch.testing.assert_close(
+ stacked["condition_images"][0][1][:, 0, 0], torch.tensor([0.0, 0.0, 1.0])
+ )
def test_wan_i2v_online_prepare_latents_reuses_condition_mode_path() -> None:
From 96228defb105ebacb0a7d58261584386e5d72f06 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 21:12:35 +0800
Subject: [PATCH 44/76] [dataset,docs] fix: use Wan FLF2V checkpoint
---
.agents/knowledge/topics/fix_patterns.md | 14 ++++++++++++++
README.md | 1 +
dataset/offline_smoke/profiles.py | 2 +-
guidance/gpu_validation.md | 15 ++++++++-------
tests/dataset/test_offline_smoke_profiles.py | 1 +
tests/docs/test_minimax_h3_docs.py | 3 +++
6 files changed, 28 insertions(+), 8 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 4de65f872..8e9473426 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -489,6 +489,20 @@ Based on the fix type, write the fix entry to the appropriate document:
media utility for one private representation.
- **Related Constraint**: #5
+### Endpoint-conditioned checkpoints are not interchangeable with I2V checkpoints
+- **Date**: 2026-08-30
+- **Symptom**: Every Wan first/last-frame smoke reached the transformer but failed while concatenating
+ image and text states because their leading dimensions were two and one.
+- **Root Cause**: The FLF2V matrix profile selected a standard Wan2.1 I2V checkpoint, whose image
+ projection lacks the learned endpoint positional embedding that folds two ordered CLIP image rows
+ back into one logical sample.
+- **Fix**: The public smoke profile and GPU validation plan now select the dedicated Wan2.1 FLF2V
+ checkpoint, and the supported-model table documents that checkpoint explicitly.
+- **Lesson**: Checkpoints that share a Diffusers pipeline class can still implement distinct
+ conditioning contracts. Validation profiles must bind semantic modes to weights whose trained
+ embedding layout realizes that mode instead of relying only on adapter-class compatibility.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/README.md b/README.md
index 6e6ca2788..c62bb489f 100644
--- a/README.md
+++ b/README.md
@@ -83,6 +83,7 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
| Wan2.1-I2V-14B-720P | 14B | wan2_i2v |
| Wan2.2-TI2V-5B | 5B | wan2_i2v |
| Wan2.2-I2V-A14B | A14B | wan2_i2v |
+ | First/Last-Frame-to-Video | Wan2.1-FLF2V-14B-720P | 14B | wan2_i2v |
| Text-to-Audio-Video | LTX-2 | 19B | ltx2_t2av |
| LTX-2.3 | 22B | ltx2_t2av |
diff --git a/dataset/offline_smoke/profiles.py b/dataset/offline_smoke/profiles.py
index d273719ee..bc2073e60 100644
--- a/dataset/offline_smoke/profiles.py
+++ b/dataset/offline_smoke/profiles.py
@@ -209,7 +209,7 @@ def _image_rule(
"bagel-mri2i": "ByteDance-Seed/BAGEL-7B-MoT",
"wan-t2v": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
"wan-i2v-first": "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
- "wan-flf2v": "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers",
+ "wan-flf2v": "Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers",
"ltx2-t2av": "Lightricks/LTX-2",
"ltx2-i2av": "Lightricks/LTX-2",
"h3-t2va": "MiniMaxAI/MiniMax-H3",
diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md
index 92935b7a3..272b19f1c 100644
--- a/guidance/gpu_validation.md
+++ b/guidance/gpu_validation.md
@@ -37,7 +37,7 @@ different condition layouts and active masks.
| `bagel-mri2i` | Bagel ordered multi-reference-images-to-image regression anchor | `ByteDance-Seed/BAGEL-7B-MoT` | prompt plus exactly two ordered images per sample | image |
| `wan-t2v` | Wan text-to-video | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | prompt | video |
| `wan-i2v-first` | Wan first-frame-to-video | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | exactly one first-frame image | video |
-| `wan-flf2v` | Wan first/last-frame-to-video | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | ordered first and last images | video |
+| `wan-flf2v` | Wan first/last-frame-to-video | `Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers` | ordered first and last images | video |
| `ltx2-t2av` | LTX2 text-to-audio-video | `Lightricks/LTX-2` | prompt | ordered video and audio |
| `ltx2-i2av` | LTX2 image-to-audio-video | `Lightricks/LTX-2` | prompt plus one image | ordered video and audio |
| `h3-t2va` | MiniMax H3 text-to-video-audio | `MiniMaxAI/MiniMax-H3` | prompt | ordered video and audio |
@@ -94,7 +94,7 @@ recorded in the resolved YAML; do not silently return to a large quality recipe.
| `bagel-mri2i` | `examples/grpo/lora/bagel/i2i.yaml` | `resolution: 256` | `num_inference_steps: 2` | Every row has exactly two ordered references; keep `shuffle_samples: false`. |
| `wan-t2v` | `examples/grpo/lora/wan21/t2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Target video has at least five frames and carries its source `fps`. |
| `wan-i2v-first` | `examples/grpo/lora/wan22/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use TI2V-5B and exactly one condition image. |
-| `wan-flf2v` | `examples/grpo/lora/wan21/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use two ordered condition images; do not use an expanded-timestep checkpoint. |
+| `wan-flf2v` | `examples/grpo/lora/wan21/i2v.yaml` | `resolution: 240`, `num_frames: 5` | `num_inference_steps: 2` | Use the dedicated FLF2V checkpoint with two ordered condition images. |
| `ltx2-t2av` | `examples/grpo/lora/ltx2/t2av.yaml` | `resolution: [128, 192]`, `num_frames: 9`, `frame_rate: 24.0` | `num_inference_steps: 2` | AV targets cover the exact 9-frame clock; audio carries `sample_rate`. |
| `ltx2-i2av` | `examples/grpo/lora/ltx2/i2av.yaml` | `resolution: [128, 192]`, `num_frames: 9`, `frame_rate: 24.0` | `num_inference_steps: 2` | One condition image; verify the first latent frame is inactive in the loss. |
| `h3-t2va` | `examples/grpo/lora/minimax_h3_t2va/debug.yaml` | `resolution: [64, 96]`, `num_frames: 124`, `frame_rate: 24.0` | `num_inference_steps: 2` | Preserve the released five-second minimum and neutral guidance. |
@@ -163,14 +163,15 @@ backends and all four algorithms.
|---|---|
| Wan T2V | `Wan2.1-T2V-14B-Diffusers`, `Wan2.2-TI2V-5B-Diffusers`, `Wan2.2-T2V-A14B-Diffusers` |
| Wan I2V first-only | `Wan2.1-I2V-14B-480P-Diffusers`, `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
-| Wan first/last | `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
+| Wan first/last | `Wan2.2-I2V-A14B-Diffusers` |
| LTX2 T2AV and I2AV | `dg845/LTX-2.3-Diffusers` |
| MiniMax H3 | The same checkpoint is covered separately by T2VA, FL2VA, and Ref2VA inputs. Add an FL2VA first-plus-last fixture to complement the first-only/last-only main jobs. |
-Wan2.2 TI2V-5B uses expanded timesteps. Official Diffusers ignores a supplied
-last image in that mode, so its effective input contract is first-frame only;
-do not count it as a first/last checkpoint. Other Wan I2V checkpoints must prove
-both first-only and first/last execution.
+Wan2.2 TI2V-5B uses expanded timesteps, and standard CLIP-conditioned Wan2.1 I2V
+checkpoints lack endpoint positional embeddings; both are first-frame only. The
+dedicated Wan2.1 FLF2V checkpoint requires both endpoints. Wan2.2 I2V-A14B does
+not use CLIP image embeddings, so its VAE-only optional-last path remains a
+separate first/last variant gate.
For the Wan2.2 A14B dual-transformer gate, force or instrument two offline
timestep samples so that one routes below the transformer boundary and one
diff --git a/tests/dataset/test_offline_smoke_profiles.py b/tests/dataset/test_offline_smoke_profiles.py
index d615cd1b0..b9efb8b3f 100644
--- a/tests/dataset/test_offline_smoke_profiles.py
+++ b/tests/dataset/test_offline_smoke_profiles.py
@@ -106,6 +106,7 @@ def test_gpu_variants_keep_model_specific_clocks() -> None:
assert (wan.height, wan.num_frames) == (240, 5)
assert (ltx.height, ltx.width, ltx.num_frames, ltx.sample_rate) == (128, 192, 9, 16000)
assert (h3.height, h3.width, h3.num_frames, h3.sample_rate) == (64, 96, 124, 32000)
+ assert cases["wan-flf2v"].checkpoint == "Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers"
assert p.get_profile("image-i2i") is p.CANONICAL_PROFILES["image_to_image"]
diff --git a/tests/docs/test_minimax_h3_docs.py b/tests/docs/test_minimax_h3_docs.py
index bbfc174a4..cb4f5e0f2 100644
--- a/tests/docs/test_minimax_h3_docs.py
+++ b/tests/docs/test_minimax_h3_docs.py
@@ -106,6 +106,9 @@ def test_gpu_validation_plan_declares_the_complete_smoke_matrix() -> None:
assert "two training epochs" in text
assert "eval.eval_freq: 0" in text
assert "DistributedSampler" in text
+ assert "Wan-AI/Wan2.1-FLF2V-14B-720P-diffusers" in text
+ assert "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | ordered first and last" not in text
+ assert "| Wan first/last | `Wan2.2-I2V-A14B-Diffusers` |" in text
def test_install_docs_use_the_released_diffusers_runtime() -> None:
From e31ed67ec5d5042b0ee4972dde44e21efbe4b092 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 21:16:43 +0800
Subject: [PATCH 45/76] [models,dataset] fix: specialize Wan endpoint contracts
---
.agents/knowledge/topics/fix_patterns.md | 14 +++++++++
dataset/offline_smoke/profiles.py | 4 +--
guidance/gpu_validation.md | 11 ++++---
src/flow_factory/models/wan/wan2_i2v.py | 33 ++++++++++++++++++--
tests/dataset/test_offline_smoke_profiles.py | 3 +-
tests/models/test_wan_output_codec.py | 23 ++++++++++++--
6 files changed, 76 insertions(+), 12 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 8e9473426..78f40ed25 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -503,6 +503,20 @@ Based on the fix type, write the fix entry to the appropriate document:
embedding layout realizes that mode instead of relying only on adapter-class compatibility.
- **Related Constraint**: N/A
+### Wan endpoint cardinality follows the checkpoint embedding path
+- **Date**: 2026-08-30
+- **Symptom**: The Wan I2V adapter advertised one optional last frame for every checkpoint even
+ though standard Wan2.1 I2V and dedicated FLF2V weights require different exact image counts.
+- **Root Cause**: The effective pipeline contract specialized only expanded-timestep checkpoints
+ and ignored whether the loaded transformer used no CLIP image states, ordinary CLIP states, or
+ learned first/last endpoint positional embeddings.
+- **Fix**: Wan now resolves exact-one for expanded or ordinary CLIP-conditioned checkpoints,
+ exact-two for endpoint-positioned FLF2V weights, and preserves one-or-two for Wan2.2's VAE-only
+ condition path. The public FLF smoke profile independently requires both endpoint slots.
+- **Lesson**: A shared adapter's public superset contract must be narrowed from realized checkpoint
+ structure before dataset validation, cache identity, or exact-resume identity is derived.
+- **Related Constraint**: #5
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/dataset/offline_smoke/profiles.py b/dataset/offline_smoke/profiles.py
index bc2073e60..cea883b1e 100644
--- a/dataset/offline_smoke/profiles.py
+++ b/dataset/offline_smoke/profiles.py
@@ -172,10 +172,10 @@ def _image_rule(
)
_FL2V = video_output_contract(
negative_prompt=_NO_NEGATIVE,
- input_image_min_count=1,
+ input_image_min_count=2,
input_image_max_count=2,
input_image_slots=("first_frame", "last_frame"),
- required_input_image_slots=("first_frame",),
+ required_input_image_slots=("first_frame", "last_frame"),
output_fps=RateRequirement.REQUIRED,
)
_T2AV = audio_video_output_contract(negative_prompt=_NO_NEGATIVE)
diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md
index 272b19f1c..23b8a6d37 100644
--- a/guidance/gpu_validation.md
+++ b/guidance/gpu_validation.md
@@ -47,11 +47,12 @@ different condition layouts and active masks.
The public V2 schema uses the `type` discriminator and an optional input-only
`slot`. An explicit slot reserves its adapter-declared semantic argument;
unslotted media is only a positional shorthand that fills remaining slots in
-declaration order. Wan requires `first_frame` and optionally accepts
-`last_frame`; H3 FL accepts either slot or both. Supervision media does not
-repeat independent condition-image objects. Its video is nevertheless the full
-configured output sequence, whose first and/or last endpoint must correspond to
-the supplied endpoint conditions.
+declaration order. Standard Wan2.1 I2V checkpoints require `first_frame`, the
+dedicated FLF2V checkpoint requires both endpoints, and Wan2.2 I2V-A14B retains
+an optional VAE-only `last_frame`; H3 FL accepts either slot or both. Supervision
+media does not repeat independent condition-image objects. Its video is
+nevertheless the full configured output sequence, whose first and/or last
+endpoint must correspond to the supplied endpoint conditions.
### Backends
diff --git a/src/flow_factory/models/wan/wan2_i2v.py b/src/flow_factory/models/wan/wan2_i2v.py
index f7049c21b..a08ddf66c 100644
--- a/src/flow_factory/models/wan/wan2_i2v.py
+++ b/src/flow_factory/models/wan/wan2_i2v.py
@@ -122,9 +122,38 @@ def load_pipeline(self) -> WanImageToVideoPipeline:
)
def _resolve_pipeline_io_contract(self) -> PipelineIOContract:
- """Narrow expand-timestep checkpoints to their first-frame-only semantics."""
+ """Resolve checkpoint-specific first/last-frame cardinality."""
+ supports_endpoint_pair = False
if not self.pipeline.config.expand_timesteps:
- return type(self).pipeline_io_contract
+ transformer_configs = tuple(
+ transformer.config
+ for transformer in (
+ getattr(self.pipeline, "transformer", None),
+ getattr(self.pipeline, "transformer_2", None),
+ )
+ if transformer is not None
+ )
+ clip_configs = tuple(
+ config
+ for config in transformer_configs
+ if getattr(config, "image_dim", None) is not None
+ )
+ if not clip_configs:
+ return type(self).pipeline_io_contract
+ supports_endpoint_pair = any(
+ getattr(config, "pos_embed_seq_len", None) is not None for config in clip_configs
+ )
+ if supports_endpoint_pair:
+ return video_output_contract(
+ negative_prompt=NegativePromptPolicy.OPTIONAL,
+ input_image_min_count=2,
+ input_image_max_count=2,
+ input_image_slots=("first_frame", "last_frame"),
+ required_input_image_slots=("first_frame", "last_frame"),
+ output_fps=RateRequirement.REQUIRED,
+ geometry_source=GeometrySource.CONFIGURED,
+ batch_capability=BatchCapability.SINGLE_SAMPLE,
+ )
return video_output_contract(
negative_prompt=NegativePromptPolicy.OPTIONAL,
input_image_min_count=1,
diff --git a/tests/dataset/test_offline_smoke_profiles.py b/tests/dataset/test_offline_smoke_profiles.py
index b9efb8b3f..f7a3ae0cd 100644
--- a/tests/dataset/test_offline_smoke_profiles.py
+++ b/tests/dataset/test_offline_smoke_profiles.py
@@ -79,8 +79,9 @@ def test_image_and_endpoint_profiles_keep_cardinality_order_and_slots() -> None:
assert (first_video.slots, first_video.required_slots) == (("first_frame",), ("first_frame",))
assert (fl_video.slots, fl_video.required_slots) == (
("first_frame", "last_frame"),
- ("first_frame",),
+ ("first_frame", "last_frame"),
)
+ assert (fl_video.min_count, fl_video.max_count) == (2, 2)
assert (first_av.slots, first_av.required_slots) == (("first_frame",), ("first_frame",))
assert (fl_av.slots, fl_av.required_slots) == (("first_frame", "last_frame"), ())
diff --git a/tests/models/test_wan_output_codec.py b/tests/models/test_wan_output_codec.py
index 4de48de78..16e94a891 100644
--- a/tests/models/test_wan_output_codec.py
+++ b/tests/models/test_wan_output_codec.py
@@ -299,7 +299,7 @@ def test_wan_geometry_validator_rejects_output_context_drift() -> None:
)
-def test_wan_i2v_declares_ordered_first_optional_last_offline_contract() -> None:
+def test_wan_i2v_resolves_checkpoint_specific_endpoint_contract() -> None:
contract = Wan2_I2V_Adapter.pipeline_io_contract
assert contract.geometry_source is GeometrySource.CONFIGURED
@@ -322,8 +322,27 @@ def test_wan_i2v_declares_ordered_first_optional_last_offline_contract() -> None
runtime = object.__new__(Wan2_I2V_Adapter)
runtime.pipeline = _Adapter().pipeline
+ runtime.pipeline.transformer.config.image_dim = None
+ runtime.pipeline.transformer.config.pos_embed_seq_len = None
+ runtime.pipeline.transformer_2 = SimpleNamespace(
+ config=SimpleNamespace(image_dim=None, pos_embed_seq_len=None)
+ )
+ effective = Wan2_I2V_Adapter._resolve_pipeline_io_contract(runtime)
+ rule = effective.input_media.rules[0]
+ assert (rule.min_count, rule.max_count) == (1, 2)
+
+ runtime.pipeline.transformer_2 = None
+ runtime.pipeline.transformer.config.image_dim = 1280
effective = Wan2_I2V_Adapter._resolve_pipeline_io_contract(runtime)
- assert effective.input_media.rules[0].max_count == 2
+ rule = effective.input_media.rules[0]
+ assert (rule.min_count, rule.max_count) == (1, 1)
+
+ runtime.pipeline.transformer.config.pos_embed_seq_len = 514
+ effective = Wan2_I2V_Adapter._resolve_pipeline_io_contract(runtime)
+ rule = effective.input_media.rules[0]
+ assert (rule.min_count, rule.max_count) == (2, 2)
+ assert rule.required_slots == ("first_frame", "last_frame")
+
runtime.pipeline.config.expand_timesteps = True
effective = Wan2_I2V_Adapter._resolve_pipeline_io_contract(runtime)
assert effective.input_media.rules[0].max_count == 1
From 54f92efb7f9f34c62158e3909a8be196c6694839 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 21:50:02 +0800
Subject: [PATCH 46/76] [models] fix: expose repeated blocks to FSDP
---
.agents/knowledge/topics/component_runtime.md | 18 ++++++++++++++
src/flow_factory/models/model_bundle.py | 19 ++++++++-------
tests/models/test_component_variants.py | 24 +++++++++++++++++++
3 files changed, 53 insertions(+), 8 deletions(-)
diff --git a/.agents/knowledge/topics/component_runtime.md b/.agents/knowledge/topics/component_runtime.md
index 7bf833aff..06729b22d 100644
--- a/.agents/knowledge/topics/component_runtime.md
+++ b/.agents/knowledge/topics/component_runtime.md
@@ -70,6 +70,24 @@ instance attribute.
- A target-owned composite root may still contain frozen auxiliary siblings. Pseudo runtimes move
only that remainder and exclude every prepared target route.
+## Fix records
+
+### Repeated-block metadata must survive the distributed bundle boundary
+
+- **Date**: 2026-08-30
+- **Symptom**: Two-rank LTX2 FSDP2 runs exhausted a 95 GiB GPU while
+ `fully_shard()` initialized the prepared model root, before the first training step.
+- **Root Cause**: LTX2 declares its 48 transformer units through Diffusers'
+ `_repeated_blocks`, but `ModelBundle` surfaced only `_no_split_modules` to Accelerate;
+ the empty auto-wrap policy therefore sharded the complete 19B transformer as one unit.
+- **Fix**: `ModelBundle._no_split_modules` now falls back to `_repeated_blocks` for a
+ member that has no legacy no-split declaration, with a regression that resolves the
+ repeated block class through Accelerate's transformer-based FSDP policy.
+- **Lesson**: A distributed wrapper becomes the metadata boundary seen by backend
+ policy discovery. It must preserve both legacy and current model block declarations,
+ or a correct component graph can silently collapse into one memory-prohibitive shard.
+- **Related Constraint**: #9
+
## Failure modes
- Expanding omitted materialization to all declarations loads tokenizers, configs, and weights
diff --git a/src/flow_factory/models/model_bundle.py b/src/flow_factory/models/model_bundle.py
index 09dc43339..debf0089b 100644
--- a/src/flow_factory/models/model_bundle.py
+++ b/src/flow_factory/models/model_bundle.py
@@ -69,8 +69,7 @@ def __init__(self, members: Dict[str, nn.Module]):
@property
def _no_split_modules(self):
- """Aggregate the members' ``_no_split_modules`` so accelerate's FSDP
- ``TRANSFORMER_BASED_WRAP`` can discover the transformer block class(es).
+ """Aggregate member block metadata for Accelerate's FSDP wrap policy.
accelerate's ``set_auto_wrap_policy`` reads ``getattr(root, "_no_split_modules")``
off the single root passed to ``prepare`` (this bundle). Without this the
@@ -78,17 +77,21 @@ def _no_split_modules(self):
bundle is wrapped as ONE FSDP unit -> one monolithic flat param (the full
unsharded model materialized on every rank) -> OOM at init for large models
(e.g. Wan2.2 A14B: ~53GB flat param). We surface the block-class *names* that
- the underlying model classes already declare (e.g. diffusers
- ``WanTransformer3DModel._no_split_modules == ['WanTransformerBlock']``);
- accelerate then resolves each name to its class via ``get_module_class_from_name``
- over this root's submodule tree and wraps per-block. Returns ``None`` when no
- member declares any (accelerate keeps its prior fallback unchanged).
+ underlying model classes already declare. Older models use
+ ``_no_split_modules`` (for example ``WanTransformerBlock``), while newer
+ Diffusers models may expose only ``_repeated_blocks`` (for example
+ ``LTX2VideoTransformerBlock``). Accelerate then resolves each name to its
+ class over this root's submodule tree and wraps per block. Returns ``None``
+ when no member declares either form of metadata.
"""
names: list[str] = []
# Walk members' submodule tree (NOT self, to avoid recursing on this property);
- # collect every ``_no_split_modules`` a nested module declares.
+ # preserve the legacy no-split declaration when present and otherwise use
+ # Diffusers' repeated-block declaration.
for module in self.members.modules():
nsm = getattr(module, "_no_split_modules", None)
+ if not nsm:
+ nsm = getattr(module, "_repeated_blocks", None)
if nsm:
names.extend(nsm)
deduped = list(dict.fromkeys(names))
diff --git a/tests/models/test_component_variants.py b/tests/models/test_component_variants.py
index 3d64b2c9e..dda02ffbb 100644
--- a/tests/models/test_component_variants.py
+++ b/tests/models/test_component_variants.py
@@ -20,6 +20,7 @@
import pytest
import torch
+from accelerate import FullyShardedDataParallelPlugin
from peft import LoraConfig, get_peft_model
from flow_factory.models.abc import BaseAdapter
@@ -32,6 +33,29 @@
)
+def test_model_bundle_exposes_diffusers_repeated_blocks_to_fsdp() -> None:
+ class RepeatedBlock(torch.nn.Module):
+ pass
+
+ class RepeatedBlockModel(torch.nn.Module):
+ _repeated_blocks = ["RepeatedBlock"]
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.blocks = torch.nn.ModuleList([RepeatedBlock(), RepeatedBlock()])
+
+ bundle = ModelBundle({"transformer": RepeatedBlockModel()})
+ plugin = FullyShardedDataParallelPlugin(
+ fsdp_version=2,
+ auto_wrap_policy="transformer_based_wrap",
+ )
+
+ plugin.set_auto_wrap_policy(bundle)
+
+ assert bundle._no_split_modules == ["RepeatedBlock"]
+ assert plugin.auto_wrap_policy.keywords["transformer_layer_cls"] == {RepeatedBlock}
+
+
def _registry() -> ComponentVariantRegistry:
return ComponentVariantRegistry(
SimpleNamespace(trainable_component_names=["transformer", "transformer_2"])
From 7684fb6c952c5d07683c13025cc67eb4af6e1bb9 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 22:15:37 +0800
Subject: [PATCH 47/76] [models] fix: restore H3 layout coordinate precision
---
.agents/knowledge/topics/dtype_precision.md | 8 ++++++++
src/flow_factory/models/minimax_h3/workflow.py | 8 +++++++-
tests/models/minimax_h3/test_review_fixes.py | 10 +++++++++-
3 files changed, 24 insertions(+), 2 deletions(-)
diff --git a/.agents/knowledge/topics/dtype_precision.md b/.agents/knowledge/topics/dtype_precision.md
index 2913b9995..0280fdbd9 100644
--- a/.agents/knowledge/topics/dtype_precision.md
+++ b/.agents/knowledge/topics/dtype_precision.md
@@ -86,6 +86,14 @@ The round-trip ensures that the precision of stored latents matches what trainin
- **Lesson**: Compute precision and boundary representation are separate contracts. Restore representation at the producer boundary instead of weakening adapter validation or casting from global configuration.
- **Related Constraint**: N/A
+### H3 cached layout coordinates lost their semantic dtype
+- **Date**: 2026-08-30
+- **Symptom**: MiniMax H3 offline SFT and DPO rejected cached `position_ids` as float32 even though the preprocessing cache schema declared float64 coordinates.
+- **Root Cause**: Hugging Face Dataset's torch formatter defaults floating Arrow columns to float32, and H3 layout normalization preserved that runtime downcast.
+- **Fix**: H3 layout normalization now restores matrix coordinates to float64 before output encoding, with a regression covering a float32 cached tensor.
+- **Lesson**: An Arrow schema dtype does not guarantee the dtype returned by a runtime formatter. Restore model-semantic precision at the adapter normalization boundary before strict validation.
+- **Related Constraint**: #20
+
## Cross-refs
- `constraints.md` #18 (all-rank synchronization — precision errors may manifest differently per rank)
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index c0702da5d..e7d4665b4 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -959,7 +959,13 @@ def _normalize_layout(values: Mapping[str, Any]) -> Dict[str, Any]:
f"MiniMax H3 layout field={field!r} expected shape (N,D) or "
f"collated (B=1,N,D), received {tuple(value.shape)}"
)
- normalized[field] = value
+ if value.dtype not in (torch.float32, torch.float64):
+ raise ValueError(
+ f"MiniMax H3 layout field={field!r} expected dtype float32 or float64, "
+ f"received {value.dtype}"
+ )
+ # HF Dataset's torch formatter downcasts cached float64 coordinates.
+ normalized[field] = value.to(dtype=torch.float64)
for field in _LAYOUT_INDEX_FIELDS:
if field not in source:
continue
diff --git a/tests/models/minimax_h3/test_review_fixes.py b/tests/models/minimax_h3/test_review_fixes.py
index 91ffd9d76..41ce8d9ec 100644
--- a/tests/models/minimax_h3/test_review_fixes.py
+++ b/tests/models/minimax_h3/test_review_fixes.py
@@ -299,10 +299,11 @@ def test_pinned_preprocessing_runs_under_no_grad(monkeypatch) -> None:
def test_layout_normalization_uses_field_specific_shapes() -> None:
+ position_ids = torch.arange(21, dtype=torch.float32).reshape(1, 7, 3) / 3
layout = workflow._normalize_layout(
{
"layout": {
- "position_ids": torch.zeros(1, 7, 3),
+ "position_ids": position_ids,
"token_tags": torch.arange(7).unsqueeze(0),
"video_indices": torch.arange(2).unsqueeze(0),
"num_condition_video_rows": [1],
@@ -311,11 +312,18 @@ def test_layout_normalization_uses_field_specific_shapes() -> None:
)
assert layout["position_ids"].shape == (7, 3)
+ assert layout["position_ids"].dtype == torch.float64
+ assert torch.equal(layout["position_ids"], position_ids[0].to(torch.float64))
assert layout["token_tags"].shape == (7,)
assert layout["video_indices"].shape == (2,)
assert layout["num_condition_video_rows"] == 1
+def test_layout_normalization_rejects_non_floating_position_ids() -> None:
+ with pytest.raises(ValueError, match="expected dtype float32 or float64, received torch.int64"):
+ workflow._normalize_layout({"position_ids": torch.zeros(7, 3, dtype=torch.int64)})
+
+
def test_decode_materializes_exact_frozen_components(monkeypatch) -> None:
monkeypatch.setattr(
workflow,
From fcac4c5a5b4d0d8b084184fbd56eb503869c92e6 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 22:33:27 +0800
Subject: [PATCH 48/76] [dataset] fix: forward ordered reference manifests
---
.agents/knowledge/topics/fix_patterns.md | 15 +++++++++++++++
src/flow_factory/data_utils/dataset.py | 5 ++++-
tests/data_utils/test_ordered_references.py | 4 ++++
3 files changed, 23 insertions(+), 1 deletion(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 78f40ed25..b67fcadb0 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -517,6 +517,21 @@ Based on the fix type, write the fix entry to the appropriate document:
structure before dataset validation, cache identity, or exact-resume identity is derived.
- **Related Constraint**: #5
+### Ordered-reference preprocessors need the canonical manifest sidecar
+- **Date**: 2026-08-30
+- **Symptom**: MiniMax H3 Ref2VA failed during distributed dataset preprocessing because its
+ strict workflow received decoded `references` but `reference_manifest=None`.
+- **Root Cause**: `GeneralDataset` canonicalized and retained each ordered-reference manifest for
+ Arrow output, but omitted that same manifest from the arguments passed to the adapter
+ preprocessor.
+- **Fix**: Ordered-reference preprocessing now forwards the canonical manifest beside the decoded
+ transient media, and the real-media round-trip regression requires the preprocessor to receive
+ the same canonical ordering that is stored in the cache.
+- **Lesson**: Identity and reconstruction sidecars must cross the same preprocessing boundary as
+ the transient inputs they describe. Preserve strict batch validation at the adapter instead of
+ delaying a missing-sidecar failure until sample construction.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/data_utils/dataset.py b/src/flow_factory/data_utils/dataset.py
index f7af6eccd..cf7e09653 100644
--- a/src/flow_factory/data_utils/dataset.py
+++ b/src/flow_factory/data_utils/dataset.py
@@ -720,7 +720,10 @@ def _preprocess_batch(
for reference_index, entry in enumerate(json.loads(manifest))
]
)
- reference_args["references"] = loaded_reference_batch
+ reference_args.update(
+ references=loaded_reference_batch,
+ reference_manifest=canonical_manifests,
+ )
batch["reference_manifest"] = canonical_manifests
slot_args = {
diff --git a/tests/data_utils/test_ordered_references.py b/tests/data_utils/test_ordered_references.py
index 18a02ddf0..6cc984eac 100644
--- a/tests/data_utils/test_ordered_references.py
+++ b/tests/data_utils/test_ordered_references.py
@@ -38,11 +38,13 @@ class OrderedPreprocessor:
def __init__(self) -> None:
self.received: List[List[Dict[str, Any]]] = []
+ self.received_manifests: List[str] = []
def preprocess(
self,
prompt: List[str],
references: List[List[Dict[str, Any]]],
+ reference_manifest: List[str],
workflow: str,
width: int,
height: int = 512,
@@ -51,6 +53,7 @@ def preprocess(
model_name_or_path: str = "MiniMaxAI/MiniMax-H3",
) -> Dict[str, Any]:
self.received = references
+ self.received_manifests = reference_manifest
assert workflow == "ref2va"
return {"encoded": torch.tensor([[len(references[0])]], dtype=torch.float32)}
@@ -143,6 +146,7 @@ def test_ordered_references_round_trip_real_media_and_merged_cache(tmp_path: Pat
assert preprocessor.received[0][0]["media"].size == (4, 3)
assert preprocessor.received[0][1]["sample_rate"] == 22050
assert preprocessor.received[0][1]["media"].shape[0] == 1
+ assert json.loads(preprocessor.received_manifests[0]) == references
row = dataset[0]
assert json.loads(row["reference_manifest"]) == references
assert all(value is not None for value in row.values())
From 2b68c2f6a8a2374797b3911f7b52bdca4c641223 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 23:07:29 +0800
Subject: [PATCH 49/76] [models] fix: bound H3 feed-forward token memory
---
.agents/knowledge/topics/fix_patterns.md | 15 ++
.../models/minimax_h3/_chunking.py | 153 ++++++++++++++++++
.../models/minimax_h3/workflow.py | 6 +-
tests/models/minimax_h3/test_chunking.py | 151 +++++++++++++++++
.../minimax_h3/test_workflow_adapters.py | 42 ++++-
5 files changed, 362 insertions(+), 5 deletions(-)
create mode 100644 src/flow_factory/models/minimax_h3/_chunking.py
create mode 100644 tests/models/minimax_h3/test_chunking.py
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index b67fcadb0..1028f3853 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -532,6 +532,21 @@ Based on the fix type, write the fix entry to the appropriate document:
delaying a missing-sidecar failure until sample construction.
- **Related Constraint**: N/A
+### Token-local H3 feed-forward work should bound packed-sequence temporaries
+- **Date**: 2026-08-30
+- **Symptom**: MiniMax H3 Ref2VA ZeRO-2 offline DPO and FSDP2 trainers exhausted a 95 GiB
+ device while allocating one 380 MiB SwiGLU activation for a 13,889-token packed sequence.
+- **Root Cause**: Diffusers' H3 blocks evaluated every feed-forward projection over the complete
+ dynamic sequence, and its generic chunk helper rejected non-divisible lengths such as 13,889.
+- **Fix**: H3 runtime setup now reuses each existing `ff.net` parameter tree inside a remainder-safe
+ 4,096-token executor for both token-refiner and main transformer blocks. Installation precedes
+ Flow-Factory resume loading, LoRA, checkpointing, and distributed wrapping, so state-dict keys,
+ parameter identities, and execution policy remain consistent across rollout and training.
+- **Lesson**: A token-local operation does not need to inherit the peak allocation of a packed
+ attention sequence. Apply memory bounds at the operation boundary while preserving the original
+ parameter tree and accepting dynamic tail chunks.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
new file mode 100644
index 000000000..63dd64f7c
--- /dev/null
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -0,0 +1,153 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Bound peak memory for token-local MiniMax H3 feed-forward layers."""
+
+from __future__ import annotations
+
+from typing import Iterable
+
+import torch
+from torch import nn
+
+H3_MAX_FEED_FORWARD_TOKENS = 4096
+
+
+class _ChunkedFeedForward(nn.Module):
+ """Run one existing feed-forward network over remainder-safe token chunks.
+
+ The upstream Diffusers H3 blocks expose their feed-forward layers as a sole
+ ``net`` ModuleList. Registering that same ModuleList directly keeps parameter
+ identities and ``ff.net.*`` state-dict/LoRA paths unchanged.
+ """
+
+ def __init__(self, net: nn.ModuleList, *, max_tokens: int) -> None:
+ super().__init__()
+ self.net = net
+ self.max_tokens = _positive_int(max_tokens, "max_tokens")
+
+ def _forward_chunk(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for module in self.net:
+ hidden_states = module(hidden_states)
+ return hidden_states
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ """Apply the token-local network without requiring an even split."""
+ if hidden_states.ndim < 3:
+ raise ValueError(
+ "MiniMax H3 feed-forward chunking expected [batch, tokens, hidden] "
+ f"input, received shape={tuple(hidden_states.shape)}"
+ )
+ if hidden_states.shape[1] <= self.max_tokens:
+ return self._forward_chunk(hidden_states)
+ return torch.cat(
+ [self._forward_chunk(chunk) for chunk in hidden_states.split(self.max_tokens, dim=1)],
+ dim=1,
+ )
+
+
+def install_h3_feed_forward_chunking(
+ transformer: nn.Module,
+ *,
+ max_tokens: int = H3_MAX_FEED_FORWARD_TOKENS,
+) -> int:
+ """Install bounded feed-forward execution on both H3 repeated stacks.
+
+ Installation happens immediately after pretrained component materialization,
+ before Flow-Factory resume loading, LoRA injection, gradient checkpointing, or
+ distributed wrapping. The mutation is idempotent and preserves every parameter
+ object.
+
+ Returns:
+ Number of feed-forward layers configured across both stacks.
+ """
+ max_tokens = _positive_int(max_tokens, "max_tokens")
+ blocks = tuple(_h3_feed_forward_blocks(transformer))
+ configured = 0
+ for name, block in blocks:
+ feed_forward = getattr(block, "ff", None)
+ if isinstance(feed_forward, _ChunkedFeedForward):
+ if feed_forward.max_tokens != max_tokens:
+ raise ValueError(
+ f"MiniMax H3 {name}.ff already uses max_tokens="
+ f"{feed_forward.max_tokens}, received conflicting {max_tokens}"
+ )
+ configured += 1
+ continue
+ net = getattr(feed_forward, "net", None)
+ if not isinstance(net, nn.ModuleList):
+ raise TypeError(
+ f"MiniMax H3 {name}.ff expected a sole net ModuleList, received "
+ f"{type(feed_forward).__name__} with net={type(net).__name__}"
+ )
+ if tuple(feed_forward._modules) != ("net",):
+ raise TypeError(
+ f"MiniMax H3 {name}.ff expected only the net child module, received "
+ f"children={tuple(feed_forward._modules)}"
+ )
+ if not (
+ len(net) == 3
+ and type(net[0]).__name__ == "SwiGLU"
+ and isinstance(net[1], nn.Dropout)
+ and net[1].p == 0.0
+ and isinstance(net[2], nn.Linear)
+ ):
+ raise TypeError(
+ f"MiniMax H3 {name}.ff expected SwiGLU, Dropout(0), Linear; "
+ f"received={[type(module).__name__ for module in net]}"
+ )
+ if tuple(feed_forward.named_parameters(recurse=False)) or tuple(
+ feed_forward.named_buffers(recurse=False)
+ ):
+ raise TypeError(f"MiniMax H3 {name}.ff expected no direct parameters or buffers")
+ hook_fields = (
+ "_forward_pre_hooks",
+ "_forward_hooks",
+ "_backward_pre_hooks",
+ "_backward_hooks",
+ )
+ active_hooks = tuple(field for field in hook_fields if getattr(feed_forward, field, None))
+ if active_hooks or getattr(feed_forward, "_hf_hook", None) is not None:
+ raise TypeError(
+ f"MiniMax H3 {name}.ff must be configured before execution hooks; "
+ f"received hooks={active_hooks}, hf_hook="
+ f"{type(getattr(feed_forward, '_hf_hook', None)).__name__}"
+ )
+ replacement = _ChunkedFeedForward(net, max_tokens=max_tokens)
+ replacement.train(feed_forward.training)
+ block.ff = replacement
+ configured += 1
+ return configured
+
+
+def _h3_feed_forward_blocks(transformer: nn.Module) -> Iterable[tuple[str, nn.Module]]:
+ token_refiner = getattr(transformer, "token_refiner", None)
+ stacks = (
+ ("token_refiner.refiner_blocks", getattr(token_refiner, "refiner_blocks", None)),
+ ("transformer_blocks", getattr(transformer, "transformer_blocks", None)),
+ )
+ for stack_name, stack in stacks:
+ if not isinstance(stack, nn.ModuleList) or not stack:
+ raise TypeError(
+ f"MiniMax H3 expected non-empty {stack_name} ModuleList, received "
+ f"{type(stack).__name__}"
+ )
+ for index, block in enumerate(stack):
+ yield f"{stack_name}.{index}", block
+
+
+def _positive_int(value: int, field: str) -> int:
+ if not isinstance(value, int) or isinstance(value, bool) or value < 1:
+ raise ValueError(f"MiniMax H3 {field} expected a positive int, received {value!r}")
+ return value
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index e7d4665b4..0b3376247 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -30,13 +30,12 @@
)
from ...scheduler import MiniMaxH3SDEScheduler, SchedulerGroup
from ..runtime import ModularPipelineRuntime
+from ._chunking import install_h3_feed_forward_chunking
from ._common import (
build_structured_trajectories,
)
from ._common import build_training_component_times as build_h3_component_times
-from ._common import (
- validate_target_state,
-)
+from ._common import validate_target_state
from .blocks import encode_h3_workflow_inputs, prepare_h3_rollout_state
from .decoding import decode_h3_targets
from .denoise import forward_h3_state
@@ -98,6 +97,7 @@ def build_h3_component_runtime(adapter: Any) -> ModularPipelineRuntime:
validate_h3_target_components(adapter)
runtime = ModularPipelineRuntime.from_adapter(adapter, adapter.load_pipeline())
runtime.materialize_components([adapter.transformer_component_name])
+ install_h3_feed_forward_chunking(runtime.get_component(adapter.transformer_component_name))
return runtime
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
new file mode 100644
index 000000000..82f3fd66d
--- /dev/null
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -0,0 +1,151 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from copy import deepcopy
+
+import pytest
+import torch
+from torch import nn
+
+from flow_factory.models.minimax_h3._chunking import (
+ _ChunkedFeedForward,
+ install_h3_feed_forward_chunking,
+)
+
+
+class SwiGLU(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.proj = nn.Linear(5, 28, bias=False)
+ self.activation = nn.SiLU()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states, gate = self.proj(hidden_states).chunk(2, dim=-1)
+ return hidden_states * self.activation(gate)
+
+
+class FeedForwardFake(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.net = nn.ModuleList(
+ [
+ SwiGLU(),
+ nn.Dropout(0.0),
+ nn.Linear(14, 5, bias=False),
+ ]
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for module in self.net:
+ hidden_states = module(hidden_states)
+ return hidden_states
+
+
+class BlockFake(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.ff = FeedForwardFake()
+
+
+class TransformerFake(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.token_refiner = nn.Module()
+ self.token_refiner.refiner_blocks = nn.ModuleList([BlockFake(), BlockFake()])
+ self.transformer_blocks = nn.ModuleList([BlockFake(), BlockFake(), BlockFake()])
+
+
+def _feed_forwards(transformer: TransformerFake) -> list[nn.Module]:
+ return [
+ *[block.ff for block in transformer.token_refiner.refiner_blocks],
+ *[block.ff for block in transformer.transformer_blocks],
+ ]
+
+
+def test_chunking_preserves_parameter_tree_and_is_idempotent() -> None:
+ transformer = TransformerFake()
+ state = deepcopy(transformer.state_dict())
+ keys_before = tuple(state)
+ parameter_ids_before = {
+ name: id(parameter) for name, parameter in transformer.named_parameters()
+ }
+
+ assert install_h3_feed_forward_chunking(transformer, max_tokens=4) == 5
+ assert install_h3_feed_forward_chunking(transformer, max_tokens=4) == 5
+
+ assert all(isinstance(module, _ChunkedFeedForward) for module in _feed_forwards(transformer))
+ assert tuple(transformer.state_dict()) == keys_before
+ assert {
+ name: id(parameter) for name, parameter in transformer.named_parameters()
+ } == parameter_ids_before
+ transformer.load_state_dict(state, strict=True)
+ assert "transformer_blocks.0.ff.net.0.proj.weight" in transformer.state_dict()
+ assert not any("inner" in name for name, _ in transformer.named_modules())
+
+
+def test_chunking_handles_remainder_and_preserves_forward_backward() -> None:
+ torch.manual_seed(17)
+ direct = TransformerFake().double()
+ chunked = deepcopy(direct)
+ install_h3_feed_forward_chunking(chunked, max_tokens=4)
+ direct_ff = direct.transformer_blocks[0].ff
+ chunked_ff = chunked.transformer_blocks[0].ff
+ chunk_sizes: list[int] = []
+ handle = chunked_ff.net[0].register_forward_pre_hook(
+ lambda _module, inputs: chunk_sizes.append(inputs[0].shape[1])
+ )
+ direct_input = torch.randn(2, 9, 5, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+ output_gradient = torch.randn(2, 9, 5, dtype=torch.float64)
+
+ direct_output = direct_ff(direct_input)
+ chunked_output = chunked_ff(chunked_input)
+ direct_output.backward(output_gradient)
+ chunked_output.backward(output_gradient)
+ handle.remove()
+
+ assert chunk_sizes == [4, 4, 1]
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ for direct_parameter, chunked_parameter in zip(direct_ff.parameters(), chunked_ff.parameters()):
+ torch.testing.assert_close(chunked_parameter.grad, direct_parameter.grad)
+
+
+def test_short_feed_forward_uses_one_execution() -> None:
+ transformer = TransformerFake()
+ install_h3_feed_forward_chunking(transformer, max_tokens=4)
+ feed_forward = transformer.transformer_blocks[0].ff
+ chunk_sizes: list[int] = []
+ handle = feed_forward.net[0].register_forward_pre_hook(
+ lambda _module, inputs: chunk_sizes.append(inputs[0].shape[1])
+ )
+
+ feed_forward(torch.randn(2, 4, 5))
+ handle.remove()
+
+ assert chunk_sizes == [4]
+
+
+@pytest.mark.parametrize("max_tokens", [0, -1, True, 1.5])
+def test_chunking_rejects_invalid_max_tokens(max_tokens: object) -> None:
+ with pytest.raises(ValueError, match="max_tokens expected a positive int"):
+ install_h3_feed_forward_chunking(TransformerFake(), max_tokens=max_tokens)
+
+
+def test_chunking_rejects_conflicting_reinstallation() -> None:
+ transformer = TransformerFake()
+ install_h3_feed_forward_chunking(transformer, max_tokens=4)
+
+ with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
+ install_h3_feed_forward_chunking(transformer, max_tokens=8)
diff --git a/tests/models/minimax_h3/test_workflow_adapters.py b/tests/models/minimax_h3/test_workflow_adapters.py
index 8a3829018..e3ca52398 100644
--- a/tests/models/minimax_h3/test_workflow_adapters.py
+++ b/tests/models/minimax_h3/test_workflow_adapters.py
@@ -38,6 +38,36 @@ class UpstreamSchedulerFake:
"""Represent the lazy upstream scheduler replaced by Flow-Factory."""
+class SwiGLU(nn.Module):
+ """Match the upstream activation class name without adding parameters."""
+
+ def forward(self, value: torch.Tensor) -> torch.Tensor:
+ return value
+
+
+class FeedForwardFake(nn.Module):
+ """Expose the parameter-free upstream ``ff.net`` module structure."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ output = nn.Linear(1, 1, bias=False)
+ output.weight = None
+ self.net = nn.ModuleList([SwiGLU(), nn.Dropout(0.0), output])
+
+ def forward(self, value: torch.Tensor) -> torch.Tensor:
+ for module in self.net:
+ value = module(value)
+ return value
+
+
+class TransformerBlockFake(nn.Module):
+ """Expose one feed-forward child without adding trainable parameters."""
+
+ def __init__(self) -> None:
+ super().__init__()
+ self.ff = FeedForwardFake()
+
+
class TransformerFake(nn.Module):
"""Provide one parameter for BaseAdapter freeze and precision setup."""
@@ -46,6 +76,9 @@ class TransformerFake(nn.Module):
def __init__(self) -> None:
super().__init__()
self.weight = nn.Parameter(torch.ones(1))
+ self.token_refiner = nn.Module()
+ self.token_refiner.refiner_blocks = nn.ModuleList([TransformerBlockFake()])
+ self.transformer_blocks = nn.ModuleList([TransformerBlockFake()])
self.gradient_checkpointing_calls = 0
def forward(self, value: torch.Tensor) -> torch.Tensor:
@@ -222,8 +255,7 @@ def test_workflow_loader_preserves_hub_source() -> None:
assert isinstance(pipeline, WorkflowPipelineFake)
assert WorkflowPipelineFake.calls == [(model_name_or_path, "t2va")]
assert all(
- spec.pretrained_model_name_or_path == model_name_or_path
- and spec.revision == "main"
+ spec.pretrained_model_name_or_path == model_name_or_path and spec.revision == "main"
for spec in pipeline._component_specs.values()
)
@@ -304,6 +336,12 @@ def test_workflow_adapter_loads_pruned_runtime_and_exact_setup_components(
assert len(adapter.audio_scheduler.timesteps) == 4
assert not hasattr(adapter.pipeline, "unrelated")
assert transformer_name in adapter.component_runtime.materialized_component_names
+ transformer = adapter.get_component(transformer_name)
+ configured_blocks = [
+ *transformer.token_refiner.refiner_blocks,
+ *transformer.transformer_blocks,
+ ]
+ assert all(getattr(block.ff, "max_tokens", None) == 4096 for block in configured_blocks)
opposite = "transformer_ref" if transformer_name == "transformer" else "transformer"
assert opposite not in adapter.component_runtime.declared_component_names
From ed813541b256a5f6df5278286621b923ef33c733 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 23:12:34 +0800
Subject: [PATCH 50/76] [models] fix: tighten H3 feed-forward chunk bound
---
.agents/knowledge/topics/fix_patterns.md | 2 +-
src/flow_factory/models/minimax_h3/_chunking.py | 2 +-
tests/models/minimax_h3/test_workflow_adapters.py | 7 ++++++-
3 files changed, 8 insertions(+), 3 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 1028f3853..e21991060 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -539,7 +539,7 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Root Cause**: Diffusers' H3 blocks evaluated every feed-forward projection over the complete
dynamic sequence, and its generic chunk helper rejected non-divisible lengths such as 13,889.
- **Fix**: H3 runtime setup now reuses each existing `ff.net` parameter tree inside a remainder-safe
- 4,096-token executor for both token-refiner and main transformer blocks. Installation precedes
+ 2,048-token executor for both token-refiner and main transformer blocks. Installation precedes
Flow-Factory resume loading, LoRA, checkpointing, and distributed wrapping, so state-dict keys,
parameter identities, and execution policy remain consistent across rollout and training.
- **Lesson**: A token-local operation does not need to inherit the peak allocation of a packed
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 63dd64f7c..25ef591be 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -21,7 +21,7 @@
import torch
from torch import nn
-H3_MAX_FEED_FORWARD_TOKENS = 4096
+H3_MAX_FEED_FORWARD_TOKENS = 2048
class _ChunkedFeedForward(nn.Module):
diff --git a/tests/models/minimax_h3/test_workflow_adapters.py b/tests/models/minimax_h3/test_workflow_adapters.py
index e3ca52398..14fcdf318 100644
--- a/tests/models/minimax_h3/test_workflow_adapters.py
+++ b/tests/models/minimax_h3/test_workflow_adapters.py
@@ -23,6 +23,7 @@
from accelerate import DistributedType
from flow_factory.models.abc import BaseAdapter
+from flow_factory.models.minimax_h3._chunking import H3_MAX_FEED_FORWARD_TOKENS
from flow_factory.models.minimax_h3.adapters import (
MiniMaxH3FL2VAAdapter,
MiniMaxH3Ref2VAAdapter,
@@ -337,11 +338,15 @@ def test_workflow_adapter_loads_pruned_runtime_and_exact_setup_components(
assert not hasattr(adapter.pipeline, "unrelated")
assert transformer_name in adapter.component_runtime.materialized_component_names
transformer = adapter.get_component(transformer_name)
+ assert H3_MAX_FEED_FORWARD_TOKENS == 2048
configured_blocks = [
*transformer.token_refiner.refiner_blocks,
*transformer.transformer_blocks,
]
- assert all(getattr(block.ff, "max_tokens", None) == 4096 for block in configured_blocks)
+ assert all(
+ getattr(block.ff, "max_tokens", None) == H3_MAX_FEED_FORWARD_TOKENS
+ for block in configured_blocks
+ )
opposite = "transformer_ref" if transformer_name == "transformer" else "transformer"
assert opposite not in adapter.component_runtime.declared_component_names
From 1338bb976c3d776baf0b424064e6e1819f3fe321 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 23:16:24 +0800
Subject: [PATCH 51/76] [models] fix: finalize H3 feed-forward chunk bound
---
.agents/knowledge/topics/fix_patterns.md | 2 +-
src/flow_factory/models/minimax_h3/_chunking.py | 2 +-
tests/models/minimax_h3/test_workflow_adapters.py | 2 +-
3 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index e21991060..65fe09a21 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -539,7 +539,7 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Root Cause**: Diffusers' H3 blocks evaluated every feed-forward projection over the complete
dynamic sequence, and its generic chunk helper rejected non-divisible lengths such as 13,889.
- **Fix**: H3 runtime setup now reuses each existing `ff.net` parameter tree inside a remainder-safe
- 2,048-token executor for both token-refiner and main transformer blocks. Installation precedes
+ 1,024-token executor for both token-refiner and main transformer blocks. Installation precedes
Flow-Factory resume loading, LoRA, checkpointing, and distributed wrapping, so state-dict keys,
parameter identities, and execution policy remain consistent across rollout and training.
- **Lesson**: A token-local operation does not need to inherit the peak allocation of a packed
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 25ef591be..5a59339ac 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -21,7 +21,7 @@
import torch
from torch import nn
-H3_MAX_FEED_FORWARD_TOKENS = 2048
+H3_MAX_FEED_FORWARD_TOKENS = 1024
class _ChunkedFeedForward(nn.Module):
diff --git a/tests/models/minimax_h3/test_workflow_adapters.py b/tests/models/minimax_h3/test_workflow_adapters.py
index 14fcdf318..418bd939a 100644
--- a/tests/models/minimax_h3/test_workflow_adapters.py
+++ b/tests/models/minimax_h3/test_workflow_adapters.py
@@ -338,7 +338,7 @@ def test_workflow_adapter_loads_pruned_runtime_and_exact_setup_components(
assert not hasattr(adapter.pipeline, "unrelated")
assert transformer_name in adapter.component_runtime.materialized_component_names
transformer = adapter.get_component(transformer_name)
- assert H3_MAX_FEED_FORWARD_TOKENS == 2048
+ assert H3_MAX_FEED_FORWARD_TOKENS == 1024
configured_blocks = [
*transformer.token_refiner.refiner_blocks,
*transformer.transformer_blocks,
From 7100191754c8a382f50344d24bdc0d8c7096b279 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 23:24:09 +0800
Subject: [PATCH 52/76] [models] fix: bound H3 attention normalization memory
---
.agents/knowledge/topics/fix_patterns.md | 14 +++
.../models/minimax_h3/_chunking.py | 119 +++++++++++++++---
.../models/minimax_h3/workflow.py | 9 +-
tests/models/minimax_h3/test_chunking.py | 85 +++++++++++++
.../minimax_h3/test_workflow_adapters.py | 14 ++-
5 files changed, 222 insertions(+), 19 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 65fe09a21..8bec520ea 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -547,6 +547,20 @@ Based on the fix type, write the fix entry to the appropriate document:
parameter tree and accepting dynamic tail chunks.
- **Related Constraint**: N/A
+### Head-space normalization should not upcast an entire packed sequence
+- **Date**: 2026-08-30
+- **Symptom**: MiniMax H3 Ref2VA FSDP2 GRPO exhausted a 95 GiB device while
+ `attn.norm_q` requested one 380 MiB allocation before reaching the feed-forward layer.
+- **Root Cause**: PyTorch RMSNorm promoted the complete `[1, 13889, 56, 128]` BF16 query to a
+ 379.8 MiB FP32 temporary even though normalization reduces only the final head dimension.
+- **Fix**: H3 runtime setup now reuses each Q/K RMSNorm parameter inside a remainder-safe
+ 1,024-token executor on both repeated block stacks. Parameter identities and
+ `attn.norm_q.weight` / `attn.norm_k.weight` state-dict paths remain unchanged.
+- **Lesson**: Row-local normalization can preserve its exact reduction semantics while bounding
+ the number of independent rows promoted at once. Chunk the non-reduced sequence dimension,
+ never the normalized head dimension.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 5a59339ac..2c86965bd 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -12,7 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
-"""Bound peak memory for token-local MiniMax H3 feed-forward layers."""
+"""Bound peak memory for token-local MiniMax H3 operations."""
from __future__ import annotations
@@ -22,6 +22,7 @@
from torch import nn
H3_MAX_FEED_FORWARD_TOKENS = 1024
+H3_MAX_ATTENTION_NORM_TOKENS = 1024
class _ChunkedFeedForward(nn.Module):
@@ -57,6 +58,40 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
)
+class _ChunkedRMSNorm(nn.RMSNorm):
+ """Bound row-wise RMSNorm temporaries while preserving its parameter path."""
+
+ def __init__(self, norm: nn.RMSNorm, *, max_tokens: int) -> None:
+ nn.Module.__init__(self)
+ self.normalized_shape = norm.normalized_shape
+ self.eps = norm.eps
+ self.elementwise_affine = norm.elementwise_affine
+ self.register_parameter("weight", norm.weight)
+ self.max_tokens = _positive_int(max_tokens, "max_tokens")
+
+ def _forward_chunk(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return nn.functional.rms_norm(
+ hidden_states,
+ self.normalized_shape,
+ self.weight,
+ self.eps,
+ )
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ """Normalize independent token rows without materializing a full FP32 copy."""
+ if hidden_states.ndim < 3:
+ raise ValueError(
+ "MiniMax H3 RMSNorm chunking expected [batch, tokens, ..., hidden] "
+ f"input, received shape={tuple(hidden_states.shape)}"
+ )
+ if hidden_states.shape[1] <= self.max_tokens:
+ return self._forward_chunk(hidden_states)
+ return torch.cat(
+ [self._forward_chunk(chunk) for chunk in hidden_states.split(self.max_tokens, dim=1)],
+ dim=1,
+ )
+
+
def install_h3_feed_forward_chunking(
transformer: nn.Module,
*,
@@ -73,7 +108,7 @@ def install_h3_feed_forward_chunking(
Number of feed-forward layers configured across both stacks.
"""
max_tokens = _positive_int(max_tokens, "max_tokens")
- blocks = tuple(_h3_feed_forward_blocks(transformer))
+ blocks = tuple(_h3_repeated_blocks(transformer))
configured = 0
for name, block in blocks:
feed_forward = getattr(block, "ff", None)
@@ -111,19 +146,7 @@ def install_h3_feed_forward_chunking(
feed_forward.named_buffers(recurse=False)
):
raise TypeError(f"MiniMax H3 {name}.ff expected no direct parameters or buffers")
- hook_fields = (
- "_forward_pre_hooks",
- "_forward_hooks",
- "_backward_pre_hooks",
- "_backward_hooks",
- )
- active_hooks = tuple(field for field in hook_fields if getattr(feed_forward, field, None))
- if active_hooks or getattr(feed_forward, "_hf_hook", None) is not None:
- raise TypeError(
- f"MiniMax H3 {name}.ff must be configured before execution hooks; "
- f"received hooks={active_hooks}, hf_hook="
- f"{type(getattr(feed_forward, '_hf_hook', None)).__name__}"
- )
+ _reject_execution_hooks(feed_forward, f"{name}.ff")
replacement = _ChunkedFeedForward(net, max_tokens=max_tokens)
replacement.train(feed_forward.training)
block.ff = replacement
@@ -131,7 +154,55 @@ def install_h3_feed_forward_chunking(
return configured
-def _h3_feed_forward_blocks(transformer: nn.Module) -> Iterable[tuple[str, nn.Module]]:
+def install_h3_attention_norm_chunking(
+ transformer: nn.Module,
+ *,
+ max_tokens: int = H3_MAX_ATTENTION_NORM_TOKENS,
+) -> int:
+ """Install bounded Q/K head normalization on both H3 repeated stacks."""
+ max_tokens = _positive_int(max_tokens, "max_tokens")
+ configured = 0
+ for block_name, block in _h3_repeated_blocks(transformer):
+ attention = getattr(block, "attn", None)
+ if not isinstance(attention, nn.Module):
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn expected nn.Module, received "
+ f"{type(attention).__name__}"
+ )
+ for norm_name in ("norm_q", "norm_k"):
+ path = f"{block_name}.attn.{norm_name}"
+ norm = getattr(attention, norm_name, None)
+ if isinstance(norm, _ChunkedRMSNorm):
+ if norm.max_tokens != max_tokens:
+ raise ValueError(
+ f"MiniMax H3 {path} already uses max_tokens={norm.max_tokens}, "
+ f"received conflicting {max_tokens}"
+ )
+ configured += 1
+ continue
+ if not isinstance(norm, nn.RMSNorm):
+ raise TypeError(
+ f"MiniMax H3 {path} expected nn.RMSNorm, received " f"{type(norm).__name__}"
+ )
+ if tuple(norm.named_children()) or tuple(norm.named_buffers(recurse=False)):
+ raise TypeError(f"MiniMax H3 {path} expected no child modules or direct buffers")
+ expected_parameters = ("weight",) if norm.elementwise_affine else ()
+ if (
+ tuple(name for name, _ in norm.named_parameters(recurse=False))
+ != expected_parameters
+ ):
+ raise TypeError(
+ f"MiniMax H3 {path} expected direct parameters={expected_parameters}"
+ )
+ _reject_execution_hooks(norm, path)
+ replacement = _ChunkedRMSNorm(norm, max_tokens=max_tokens)
+ replacement.train(norm.training)
+ setattr(attention, norm_name, replacement)
+ configured += 1
+ return configured
+
+
+def _h3_repeated_blocks(transformer: nn.Module) -> Iterable[tuple[str, nn.Module]]:
token_refiner = getattr(transformer, "token_refiner", None)
stacks = (
("token_refiner.refiner_blocks", getattr(token_refiner, "refiner_blocks", None)),
@@ -147,6 +218,22 @@ def _h3_feed_forward_blocks(transformer: nn.Module) -> Iterable[tuple[str, nn.Mo
yield f"{stack_name}.{index}", block
+def _reject_execution_hooks(module: nn.Module, path: str) -> None:
+ hook_fields = (
+ "_forward_pre_hooks",
+ "_forward_hooks",
+ "_backward_pre_hooks",
+ "_backward_hooks",
+ )
+ active_hooks = tuple(field for field in hook_fields if getattr(module, field, None))
+ if active_hooks or getattr(module, "_hf_hook", None) is not None:
+ raise TypeError(
+ f"MiniMax H3 {path} must be configured before execution hooks; "
+ f"received hooks={active_hooks}, hf_hook="
+ f"{type(getattr(module, '_hf_hook', None)).__name__}"
+ )
+
+
def _positive_int(value: int, field: str) -> int:
if not isinstance(value, int) or isinstance(value, bool) or value < 1:
raise ValueError(f"MiniMax H3 {field} expected a positive int, received {value!r}")
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index 0b3376247..d0a9b0ea6 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -30,7 +30,10 @@
)
from ...scheduler import MiniMaxH3SDEScheduler, SchedulerGroup
from ..runtime import ModularPipelineRuntime
-from ._chunking import install_h3_feed_forward_chunking
+from ._chunking import (
+ install_h3_attention_norm_chunking,
+ install_h3_feed_forward_chunking,
+)
from ._common import (
build_structured_trajectories,
)
@@ -97,7 +100,9 @@ def build_h3_component_runtime(adapter: Any) -> ModularPipelineRuntime:
validate_h3_target_components(adapter)
runtime = ModularPipelineRuntime.from_adapter(adapter, adapter.load_pipeline())
runtime.materialize_components([adapter.transformer_component_name])
- install_h3_feed_forward_chunking(runtime.get_component(adapter.transformer_component_name))
+ transformer = runtime.get_component(adapter.transformer_component_name)
+ install_h3_feed_forward_chunking(transformer)
+ install_h3_attention_norm_chunking(transformer)
return runtime
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index 82f3fd66d..85a2eac82 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -17,9 +17,12 @@
import pytest
import torch
from torch import nn
+from torch.nn import functional as F
from flow_factory.models.minimax_h3._chunking import (
_ChunkedFeedForward,
+ _ChunkedRMSNorm,
+ install_h3_attention_norm_chunking,
install_h3_feed_forward_chunking,
)
@@ -55,6 +58,9 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
class BlockFake(nn.Module):
def __init__(self) -> None:
super().__init__()
+ self.attn = nn.Module()
+ self.attn.norm_q = nn.RMSNorm(5)
+ self.attn.norm_k = nn.RMSNorm(5)
self.ff = FeedForwardFake()
@@ -149,3 +155,82 @@ def test_chunking_rejects_conflicting_reinstallation() -> None:
with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
install_h3_feed_forward_chunking(transformer, max_tokens=8)
+
+
+def test_attention_norm_chunking_preserves_parameter_tree_and_is_idempotent() -> None:
+ transformer = TransformerFake()
+ state = deepcopy(transformer.state_dict())
+ keys_before = tuple(state)
+ parameter_ids_before = {
+ name: id(parameter) for name, parameter in transformer.named_parameters()
+ }
+
+ assert install_h3_attention_norm_chunking(transformer, max_tokens=4) == 10
+ assert install_h3_attention_norm_chunking(transformer, max_tokens=4) == 10
+
+ norms = [
+ norm
+ for block in (
+ *transformer.token_refiner.refiner_blocks,
+ *transformer.transformer_blocks,
+ )
+ for norm in (block.attn.norm_q, block.attn.norm_k)
+ ]
+ assert all(isinstance(norm, _ChunkedRMSNorm) for norm in norms)
+ assert tuple(transformer.state_dict()) == keys_before
+ assert {
+ name: id(parameter) for name, parameter in transformer.named_parameters()
+ } == parameter_ids_before
+ transformer.load_state_dict(state, strict=True)
+ assert "transformer_blocks.0.attn.norm_q.weight" in transformer.state_dict()
+ assert not any("inner" in name for name, _ in transformer.named_modules())
+
+
+def test_attention_norm_chunking_preserves_remainder_forward_backward(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ torch.manual_seed(23)
+ direct = TransformerFake().double()
+ chunked = deepcopy(direct)
+ install_h3_attention_norm_chunking(chunked, max_tokens=4)
+ direct_norm = direct.transformer_blocks[0].attn.norm_q
+ chunked_norm = chunked.transformer_blocks[0].attn.norm_q
+ direct_input = torch.randn(2, 9, 3, 5, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+ output_gradient = torch.randn(2, 9, 3, 5, dtype=torch.float64)
+
+ direct_output = direct_norm(direct_input)
+ direct_output.backward(output_gradient)
+ chunk_sizes: list[int] = []
+ direct_rms_norm = F.rms_norm
+
+ def recording_rms_norm(hidden_states, normalized_shape, weight=None, eps=None):
+ chunk_sizes.append(hidden_states.shape[1])
+ return direct_rms_norm(hidden_states, normalized_shape, weight, eps)
+
+ monkeypatch.setattr(F, "rms_norm", recording_rms_norm)
+ chunked_output = chunked_norm(chunked_input)
+ chunked_output.backward(output_gradient)
+
+ assert chunk_sizes == [4, 4, 1]
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ torch.testing.assert_close(chunked_norm.weight.grad, direct_norm.weight.grad)
+
+
+def test_attention_norm_chunking_preserves_bfloat16_dtype() -> None:
+ transformer = TransformerFake().bfloat16()
+ direct = transformer.transformer_blocks[0].attn.norm_q(torch.randn(1, 7, 3, 5).bfloat16())
+ install_h3_attention_norm_chunking(transformer, max_tokens=4)
+
+ chunked = transformer.transformer_blocks[0].attn.norm_q(torch.randn(1, 7, 3, 5).bfloat16())
+
+ assert chunked.dtype == direct.dtype == torch.bfloat16
+
+
+def test_attention_norm_chunking_rejects_conflicting_reinstallation() -> None:
+ transformer = TransformerFake()
+ install_h3_attention_norm_chunking(transformer, max_tokens=4)
+
+ with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
+ install_h3_attention_norm_chunking(transformer, max_tokens=8)
diff --git a/tests/models/minimax_h3/test_workflow_adapters.py b/tests/models/minimax_h3/test_workflow_adapters.py
index 418bd939a..941a8983c 100644
--- a/tests/models/minimax_h3/test_workflow_adapters.py
+++ b/tests/models/minimax_h3/test_workflow_adapters.py
@@ -23,7 +23,10 @@
from accelerate import DistributedType
from flow_factory.models.abc import BaseAdapter
-from flow_factory.models.minimax_h3._chunking import H3_MAX_FEED_FORWARD_TOKENS
+from flow_factory.models.minimax_h3._chunking import (
+ H3_MAX_ATTENTION_NORM_TOKENS,
+ H3_MAX_FEED_FORWARD_TOKENS,
+)
from flow_factory.models.minimax_h3.adapters import (
MiniMaxH3FL2VAAdapter,
MiniMaxH3Ref2VAAdapter,
@@ -66,6 +69,9 @@ class TransformerBlockFake(nn.Module):
def __init__(self) -> None:
super().__init__()
+ self.attn = nn.Module()
+ self.attn.norm_q = nn.RMSNorm(1, elementwise_affine=False)
+ self.attn.norm_k = nn.RMSNorm(1, elementwise_affine=False)
self.ff = FeedForwardFake()
@@ -339,6 +345,7 @@ def test_workflow_adapter_loads_pruned_runtime_and_exact_setup_components(
assert transformer_name in adapter.component_runtime.materialized_component_names
transformer = adapter.get_component(transformer_name)
assert H3_MAX_FEED_FORWARD_TOKENS == 1024
+ assert H3_MAX_ATTENTION_NORM_TOKENS == 1024
configured_blocks = [
*transformer.token_refiner.refiner_blocks,
*transformer.transformer_blocks,
@@ -347,6 +354,11 @@ def test_workflow_adapter_loads_pruned_runtime_and_exact_setup_components(
getattr(block.ff, "max_tokens", None) == H3_MAX_FEED_FORWARD_TOKENS
for block in configured_blocks
)
+ assert all(
+ getattr(block.attn.norm_q, "max_tokens", None) == H3_MAX_ATTENTION_NORM_TOKENS
+ and getattr(block.attn.norm_k, "max_tokens", None) == H3_MAX_ATTENTION_NORM_TOKENS
+ for block in configured_blocks
+ )
opposite = "transformer_ref" if transformer_name == "transformer" else "transformer"
assert opposite not in adapter.component_runtime.declared_component_names
From 9a4735a357de5b6ae37caf0ecfdee65f205cad17 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Sun, 30 Aug 2026 23:41:17 +0800
Subject: [PATCH 53/76] [models] fix: checkpoint H3 feed-forward chunks
---
.agents/knowledge/topics/fix_patterns.md | 16 ++++
.../models/minimax_h3/_chunking.py | 17 +++-
tests/models/minimax_h3/test_chunking.py | 94 ++++++++++++++++++-
3 files changed, 125 insertions(+), 2 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 8bec520ea..f45389a30 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -561,6 +561,22 @@ Based on the fix type, write the fix entry to the appropriate document:
never the normalized head dimension.
- **Related Constraint**: N/A
+### Outer activation checkpoint replay can retain every inner token chunk
+- **Date**: 2026-08-30
+- **Symptom**: MiniMax H3 Ref2VA ZeRO-2 offline DPO reached backward replay but exhausted
+ a 95 GiB device when the final feed-forward chunk requested a 28 MiB SiLU allocation.
+- **Root Cause**: The block-level non-reentrant checkpoint replay rebuilt one autograd graph
+ containing the intermediates from every sequential feed-forward chunk. Smaller chunks reduced
+ each allocation but did not bound their cumulative saved activations.
+- **Fix**: Long, grad-enabled H3 feed-forward chunks now use nested non-reentrant activation
+ checkpoints. Each chunk is replayed independently during backward; no-grad rollout and short
+ sequences retain their direct execution paths. Checkpoint replay preserves RNG state so later
+ parameter-efficient adapters cannot silently change stochastic-gradient semantics.
+- **Lesson**: Splitting a local operator bounds forward temporaries but not necessarily backward
+ replay state. When an outer checkpoint recomputes a sequence of chunks, checkpoint each chunk
+ as the lifetime boundary for its saved activations.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 2c86965bd..779f58f73 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -20,6 +20,7 @@
import torch
from torch import nn
+from torch.utils.checkpoint import checkpoint
H3_MAX_FEED_FORWARD_TOKENS = 1024
H3_MAX_ATTENTION_NORM_TOKENS = 1024
@@ -52,8 +53,22 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
)
if hidden_states.shape[1] <= self.max_tokens:
return self._forward_chunk(hidden_states)
+ chunks = hidden_states.split(self.max_tokens, dim=1)
+ if torch.is_grad_enabled():
+ return torch.cat(
+ [
+ checkpoint(
+ self._forward_chunk,
+ chunk,
+ use_reentrant=False,
+ preserve_rng_state=True,
+ )
+ for chunk in chunks
+ ],
+ dim=1,
+ )
return torch.cat(
- [self._forward_chunk(chunk) for chunk in hidden_states.split(self.max_tokens, dim=1)],
+ [self._forward_chunk(chunk) for chunk in chunks],
dim=1,
)
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index 85a2eac82..08ef21558 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -18,6 +18,7 @@
import torch
from torch import nn
from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
from flow_factory.models.minimax_h3._chunking import (
_ChunkedFeedForward,
@@ -121,7 +122,8 @@ def test_chunking_handles_remainder_and_preserves_forward_backward() -> None:
chunked_output.backward(output_gradient)
handle.remove()
- assert chunk_sizes == [4, 4, 1]
+ assert chunk_sizes[:3] == [4, 4, 1]
+ assert sorted(chunk_sizes[3:]) == [1, 4, 4]
torch.testing.assert_close(chunked_output, direct_output)
torch.testing.assert_close(chunked_input.grad, direct_input.grad)
for direct_parameter, chunked_parameter in zip(direct_ff.parameters(), chunked_ff.parameters()):
@@ -143,6 +145,96 @@ def test_short_feed_forward_uses_one_execution() -> None:
assert chunk_sizes == [4]
+def test_chunking_preserves_nested_checkpoint_backward() -> None:
+ torch.manual_seed(19)
+ direct = TransformerFake().double()
+ chunked = deepcopy(direct)
+ install_h3_feed_forward_chunking(chunked, max_tokens=4)
+ direct_ff = direct.transformer_blocks[0].ff
+ chunked_ff = chunked.transformer_blocks[0].ff
+ direct_input = torch.randn(2, 9, 5, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+ output_gradient = torch.randn(2, 9, 5, dtype=torch.float64)
+
+ direct_output = direct_ff(direct_input)
+ chunked_output = checkpoint(chunked_ff, chunked_input, use_reentrant=False)
+ direct_output.backward(output_gradient)
+ chunked_output.backward(output_gradient)
+
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ for direct_parameter, chunked_parameter in zip(direct_ff.parameters(), chunked_ff.parameters()):
+ torch.testing.assert_close(chunked_parameter.grad, direct_parameter.grad)
+
+
+def test_chunking_checkpoints_only_long_grad_enabled_chunks(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ transformer = TransformerFake()
+ install_h3_feed_forward_chunking(transformer, max_tokens=4)
+ feed_forward = transformer.transformer_blocks[0].ff
+ calls: list[tuple[int, bool, bool]] = []
+
+ def recording_checkpoint(function, chunk, *, use_reentrant, preserve_rng_state):
+ calls.append((chunk.shape[1], use_reentrant, preserve_rng_state))
+ return function(chunk)
+
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3._chunking.checkpoint",
+ recording_checkpoint,
+ )
+
+ feed_forward(torch.randn(2, 9, 5, requires_grad=True))
+ with torch.no_grad():
+ feed_forward(torch.randn(2, 9, 5))
+ feed_forward(torch.randn(2, 4, 5, requires_grad=True))
+
+ assert calls == [(4, False, True), (4, False, True), (1, False, True)]
+
+
+def test_chunking_supports_frozen_input_with_trainable_parameters() -> None:
+ direct = TransformerFake().double()
+ chunked = deepcopy(direct)
+ install_h3_feed_forward_chunking(chunked, max_tokens=4)
+ direct_ff = direct.transformer_blocks[0].ff
+ chunked_ff = chunked.transformer_blocks[0].ff
+ for parameter in direct_ff.parameters():
+ parameter.requires_grad_(False)
+ for parameter in chunked_ff.parameters():
+ parameter.requires_grad_(False)
+ direct_ff.net[2].weight.requires_grad_(True)
+ chunked_ff.net[2].weight.requires_grad_(True)
+ hidden_states = torch.randn(2, 9, 5, dtype=torch.float64)
+
+ direct_output = direct_ff(hidden_states)
+ chunked_output = chunked_ff(hidden_states)
+ direct_output.sum().backward()
+ chunked_output.sum().backward()
+
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_ff.net[2].weight.grad, direct_ff.net[2].weight.grad)
+
+
+def test_nested_chunk_checkpoint_preserves_cpu_autocast_dtype() -> None:
+ transformer = TransformerFake()
+ install_h3_feed_forward_chunking(transformer, max_tokens=4)
+ feed_forward = transformer.transformer_blocks[0].ff
+ hidden_states = torch.randn(2, 9, 5, requires_grad=True)
+ projected_dtypes: list[torch.dtype] = []
+ handle = feed_forward.net[2].register_forward_pre_hook(
+ lambda _module, inputs: projected_dtypes.append(inputs[0].dtype)
+ )
+
+ with torch.autocast("cpu", dtype=torch.bfloat16):
+ output = checkpoint(feed_forward, hidden_states, use_reentrant=False)
+ output.float().sum().backward()
+ handle.remove()
+
+ assert output.dtype == torch.bfloat16
+ assert len(projected_dtypes) >= 6
+ assert set(projected_dtypes) == {torch.bfloat16}
+
+
@pytest.mark.parametrize("max_tokens", [0, -1, True, 1.5])
def test_chunking_rejects_invalid_max_tokens(max_tokens: object) -> None:
with pytest.raises(ValueError, match="max_tokens expected a positive int"):
From 55aa8f31a4796330309c49a0b3ab484fafe59a7a Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 00:12:18 +0800
Subject: [PATCH 54/76] [loading,models] fix: bound H3 FSDP2 gather memory
---
.agents/knowledge/topics/fix_patterns.md | 16 ++++
src/flow_factory/loading/backend.py | 47 +++++++++++
src/flow_factory/models/abc.py | 3 +
.../models/minimax_h3/adapters.py | 6 ++
tests/loading/test_backend_runtime.py | 81 ++++++++++++++++++-
.../minimax_h3/test_workflow_adapters.py | 13 +++
6 files changed, 164 insertions(+), 2 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index f45389a30..a06d2e996 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -577,6 +577,22 @@ Based on the fix type, write the fix entry to the appropriate document:
as the lifetime boundary for its saved activations.
- **Related Constraint**: N/A
+### FSDP2 overlap can retain an all-gather buffer beyond a memory-tight block
+- **Date**: 2026-08-30
+- **Symptom**: MiniMax H3 Ref2VA FSDP2 GRPO exhausted a 95 GiB device when a block
+ pre-forward all-gather requested 1.21 GiB with only 1.03 GiB free.
+- **Root Cause**: Each H3 transformer block formed one 1.21 GiB gather unit, while FSDP2's default
+ implicit overlap also retained the current raw gather result through the next block's copy-in.
+ The long Ref2VA packed sequence left too little headroom for either lifetime.
+- **Fix**: Ref2VA extends the FSDP2 wrap policy with its attention, chunked feed-forward, and
+ 496 MiB AdaLN modulation modules, splitting each block into call-ordered gather units. It also
+ opts into default-stream unshard immediately after distributed preparation so each raw gather
+ result is released after copy-out, trading communication overlap for lower peak allocation.
+- **Lesson**: When model activations nearly fill a device, communication overlap is also a memory
+ policy. Apply the backend's explicit lifetime control at the prepared root instead of adding
+ allocator flushes or weakening model semantics.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/loading/backend.py b/src/flow_factory/loading/backend.py
index 49c2c45d6..2f68133ca 100644
--- a/src/flow_factory/loading/backend.py
+++ b/src/flow_factory/loading/backend.py
@@ -120,6 +120,53 @@ def load_scope(self, role: ComponentRole) -> Iterator[None]:
else:
os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = previous
+ def prepare(self, *objects: Any) -> Any:
+ """Prepare roots and apply adapter-requested FSDP2 communication policy."""
+ self._extend_fsdp2_wrap_policy(objects)
+ prepared = super().prepare(*objects)
+ plugin = self.accelerator.state.fsdp_plugin
+ if (getattr(plugin, "fsdp_version", 1) or 1) < 2 or not getattr(
+ self.adapter, "fsdp2_use_default_stream_unshard", False
+ ):
+ return prepared
+
+ prepared_objects = prepared if isinstance(prepared, (list, tuple)) else (prepared,)
+ configured = []
+ for original, candidate in zip(objects, prepared_objects):
+ if not isinstance(original, nn.Module):
+ continue
+ configure = getattr(candidate, "_set_unshard_async_op", None)
+ if not callable(configure):
+ raise TypeError(
+ "FSDP2 default-stream unshard requires prepared modules to expose "
+ f"_set_unshard_async_op(), received {type(candidate).__name__}"
+ )
+ configure(True)
+ configured.append(type(candidate).__name__)
+ if not configured:
+ raise TypeError("FSDP2 default-stream unshard requested without a prepared module")
+ logger.info("Enabled FSDP2 default-stream unshard for roots: %s", configured)
+ return prepared
+
+ def _extend_fsdp2_wrap_policy(self, objects: Sequence[Any]) -> None:
+ plugin = self.accelerator.state.fsdp_plugin
+ additional = tuple(getattr(self.adapter, "fsdp2_additional_wrap_module_names", ()))
+ if (getattr(plugin, "fsdp_version", 1) or 1) < 2 or not additional:
+ return
+ configured = getattr(plugin, "transformer_cls_names_to_wrap", None)
+ if configured is None:
+ configured = [
+ name
+ for candidate in objects
+ if isinstance(candidate, nn.Module)
+ for name in (getattr(candidate, "_no_split_modules", None) or ())
+ ]
+ plugin.transformer_cls_names_to_wrap = list(dict.fromkeys([*configured, *additional]))
+ logger.info(
+ "Extended FSDP2 transformer wrap modules: %s",
+ plugin.transformer_cls_names_to_wrap,
+ )
+
def bootstrap_targets(self) -> None:
if not self.cpu_ram_efficient_loading or self.accelerator.num_processes <= 1:
return
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index e2abfef2d..749ea103e 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -233,6 +233,9 @@ class BaseAdapter(ABC):
# ``cache_context``. The rollout cache accelerator rejects the default.
supports_diffusers_cache: ClassVar[bool] = False
supports_fsdp2_cpu_efficient_loading: ClassVar[bool] = False
+ # Opt in only when FSDP2 communication overlap exceeds the model's activation headroom.
+ fsdp2_use_default_stream_unshard: ClassVar[bool] = False
+ fsdp2_additional_wrap_module_names: ClassVar[Tuple[str, ...]] = ()
supports_ordered_references: ClassVar[bool] = False
preprocess_cache_fields: ClassVar[frozenset[str]] = frozenset()
preprocess_cache_version: ClassVar[str] = ""
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index a9c7e74bc..ac16a9c04 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -354,6 +354,12 @@ class MiniMaxH3Ref2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
workflow: ClassVar[str] = "ref2va"
transformer_component_name: ClassVar[str] = "transformer_ref"
+ fsdp2_use_default_stream_unshard: ClassVar[bool] = True
+ fsdp2_additional_wrap_module_names: ClassVar[Tuple[str, ...]] = (
+ "_ChunkedFeedForward",
+ "MiniMaxH3AdaLayerNormModulation",
+ "MiniMaxH3Attention",
+ )
supports_ordered_references: ClassVar[bool] = True
preprocessing_modules: ClassVar[List[str]] = [
"image_processor",
diff --git a/tests/loading/test_backend_runtime.py b/tests/loading/test_backend_runtime.py
index 6af076a83..14a8ed357 100644
--- a/tests/loading/test_backend_runtime.py
+++ b/tests/loading/test_backend_runtime.py
@@ -22,14 +22,18 @@ class _AcceleratorFake:
num_processes = 2
process_index = 0
- def __init__(self, *, efficient: bool) -> None:
+ def __init__(self, *, efficient: bool, fsdp_version: int = 2) -> None:
self.state = SimpleNamespace(
fsdp_plugin=SimpleNamespace(
- fsdp_version=2,
+ fsdp_version=fsdp_version,
cpu_ram_efficient_loading=efficient,
+ transformer_cls_names_to_wrap=None,
)
)
+ def prepare(self, *objects):
+ return list(objects) if len(objects) > 1 else objects[0]
+
def wait_for_everyone(self) -> None:
return None
@@ -105,6 +109,79 @@ def test_reward_scope_restores_target_loading_environment(
assert os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] == "true"
+class _FSDPModuleFake(torch.nn.Linear):
+ def __init__(self) -> None:
+ super().__init__(2, 2)
+ self.unshard_async_op_calls = []
+
+ def _set_unshard_async_op(self, enabled: bool) -> None:
+ self.unshard_async_op_calls.append(enabled)
+
+
+def test_prepare_enables_adapter_requested_fsdp2_default_stream_unshard() -> None:
+ module = _FSDPModuleFake()
+ optimizer = object()
+ runtime = FSDPBackendLoadRuntime(
+ _AcceleratorFake(efficient=True),
+ _plan(),
+ SimpleNamespace(fsdp2_use_default_stream_unshard=True),
+ )
+
+ prepared = runtime.prepare(module, optimizer)
+
+ assert prepared == [module, optimizer]
+ assert module.unshard_async_op_calls == [True]
+
+
+def test_prepare_extends_fsdp2_wrap_policy_before_distributed_preparation() -> None:
+ module = _FSDPModuleFake()
+ module._no_split_modules = ["TransformerBlock"]
+ accelerator = _AcceleratorFake(efficient=True)
+ runtime = FSDPBackendLoadRuntime(
+ accelerator,
+ _plan(),
+ SimpleNamespace(
+ fsdp2_additional_wrap_module_names=("ChunkedFeedForward",),
+ ),
+ )
+
+ assert runtime.prepare(module) is module
+ assert accelerator.state.fsdp_plugin.transformer_cls_names_to_wrap == [
+ "TransformerBlock",
+ "ChunkedFeedForward",
+ ]
+
+
+@pytest.mark.parametrize(
+ ("fsdp_version", "requested"),
+ [(1, True), (2, False)],
+)
+def test_prepare_preserves_default_unshard_without_fsdp2_adapter_opt_in(
+ fsdp_version: int,
+ requested: bool,
+) -> None:
+ module = _FSDPModuleFake()
+ runtime = FSDPBackendLoadRuntime(
+ _AcceleratorFake(efficient=True, fsdp_version=fsdp_version),
+ _plan(),
+ SimpleNamespace(fsdp2_use_default_stream_unshard=requested),
+ )
+
+ assert runtime.prepare(module) is module
+ assert module.unshard_async_op_calls == []
+
+
+def test_prepare_rejects_missing_fsdp2_default_stream_unshard_api() -> None:
+ runtime = FSDPBackendLoadRuntime(
+ _AcceleratorFake(efficient=True),
+ _plan(),
+ SimpleNamespace(fsdp2_use_default_stream_unshard=True),
+ )
+
+ with pytest.raises(TypeError, match="requires prepared modules.*_set_unshard_async_op"):
+ runtime.prepare(torch.nn.Linear(2, 2))
+
+
def test_physical_target_root_is_not_selected_as_auxiliary() -> None:
adapter = SimpleNamespace(
_resolve_component_names=lambda components: list(components),
diff --git a/tests/models/minimax_h3/test_workflow_adapters.py b/tests/models/minimax_h3/test_workflow_adapters.py
index 941a8983c..33bea489d 100644
--- a/tests/models/minimax_h3/test_workflow_adapters.py
+++ b/tests/models/minimax_h3/test_workflow_adapters.py
@@ -42,6 +42,19 @@ class UpstreamSchedulerFake:
"""Represent the lazy upstream scheduler replaced by Flow-Factory."""
+def test_only_ref2va_requests_fsdp2_default_stream_unshard() -> None:
+ assert MiniMaxH3Ref2VAAdapter.fsdp2_use_default_stream_unshard
+ assert MiniMaxH3Ref2VAAdapter.fsdp2_additional_wrap_module_names == (
+ "_ChunkedFeedForward",
+ "MiniMaxH3AdaLayerNormModulation",
+ "MiniMaxH3Attention",
+ )
+ assert not MiniMaxH3T2VAAdapter.fsdp2_use_default_stream_unshard
+ assert not MiniMaxH3T2VAAdapter.fsdp2_additional_wrap_module_names
+ assert not MiniMaxH3FL2VAAdapter.fsdp2_use_default_stream_unshard
+ assert not MiniMaxH3FL2VAAdapter.fsdp2_additional_wrap_module_names
+
+
class SwiGLU(nn.Module):
"""Match the upstream activation class name without adding parameters."""
From 12e574d6d6ade2289f54a27b8c23587e72f0adb8 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 00:24:40 +0800
Subject: [PATCH 55/76] [models] fix: bound H3 LoRA projection memory
---
.agents/knowledge/topics/fix_patterns.md | 16 ++
.../models/minimax_h3/_chunking.py | 130 ++++++++++++++++
.../models/minimax_h3/adapters.py | 48 ++++++
tests/models/minimax_h3/test_chunking.py | 144 ++++++++++++++++++
4 files changed, 338 insertions(+)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index a06d2e996..0861496c9 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -593,6 +593,22 @@ Based on the fix type, write the fix entry to the appropriate document:
allocator flushes or weakening model semantics.
- **Related Constraint**: N/A
+### PEFT projections should not materialize a full-sequence LoRA branch
+- **Date**: 2026-08-31
+- **Symptom**: MiniMax H3 Ref2VA FSDP2 GRPO passed every parameter all-gather but exhausted a
+ 95 GiB device when one adapted K projection requested a 190 MiB LoRA output with only
+ 114--134 MiB free.
+- **Root Cause**: PEFT evaluated the base projection and low-rank branch over the complete 13,889-token
+ packed sequence before adding them, so two full projection outputs overlapped at the memory peak.
+- **Fix**: Ref2VA FSDP2 changes each existing adapted Q/K/V/output PEFT Linear to a token-chunked
+ forward after LoRA injection. The complete PEFT contract runs per chunk and writes directly into
+ one preallocated final output, preserving module and parameter identities, hooks, adapter behavior,
+ and state-dict paths without a full-size concatenation copy.
+- **Lesson**: Bound parameter-efficient adaptation at the outer adapted-module boundary. Chunking only
+ the frozen base layer leaves the adapter branch unbounded, while concatenating chunk outputs can
+ recreate the same peak through an avoidable second full-size result.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 779f58f73..dc65a853e 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -19,11 +19,13 @@
from typing import Iterable
import torch
+from peft.tuners.lora.layer import Linear as LoraLinear
from torch import nn
from torch.utils.checkpoint import checkpoint
H3_MAX_FEED_FORWARD_TOKENS = 1024
H3_MAX_ATTENTION_NORM_TOKENS = 1024
+H3_MAX_LORA_PROJECTION_TOKENS = 1024
class _ChunkedFeedForward(nn.Module):
@@ -107,6 +109,42 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
)
+class _ChunkedLoraLinear(LoraLinear):
+ """Bound PEFT projection temporaries without replacing its module or state tree."""
+
+ flow_factory_max_tokens: int
+
+ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+ """Run the complete PEFT linear contract over sequence chunks."""
+ if kwargs.get("alora_offsets") is not None:
+ raise ValueError("MiniMax H3 LoRA projection chunking does not support aLoRA offsets")
+ max_tokens = self.flow_factory_max_tokens
+ if hidden_states.ndim < 3 or hidden_states.shape[1] <= max_tokens:
+ return super().forward(hidden_states, *args, **kwargs)
+
+ forward_chunk = super().forward
+ output = None
+ offset = 0
+ for chunk in hidden_states.split(max_tokens, dim=1):
+ chunk_output = forward_chunk(chunk, *args, **kwargs)
+ if chunk_output.shape[:-1] != chunk.shape[:-1]:
+ raise RuntimeError(
+ "MiniMax H3 LoRA projection must preserve input prefix dimensions, "
+ f"received input={tuple(chunk.shape)}, output={tuple(chunk_output.shape)}"
+ )
+ if output is None:
+ output = chunk_output.new_empty((*hidden_states.shape[:-1], chunk_output.shape[-1]))
+ output.narrow(1, offset, chunk_output.shape[1]).copy_(chunk_output)
+ offset += chunk_output.shape[1]
+
+ if output is None or offset != hidden_states.shape[1]:
+ raise RuntimeError(
+ "MiniMax H3 LoRA projection failed to assemble every input token, "
+ f"received input_tokens={hidden_states.shape[1]}, output_tokens={offset}"
+ )
+ return output
+
+
def install_h3_feed_forward_chunking(
transformer: nn.Module,
*,
@@ -217,6 +255,98 @@ def install_h3_attention_norm_chunking(
return configured
+def install_h3_lora_projection_chunking(
+ transformer: nn.Module,
+ *,
+ max_tokens: int = H3_MAX_LORA_PROJECTION_TOKENS,
+) -> int:
+ """Bound adapted H3 attention projections while preserving PEFT ownership.
+
+ The installer runs after PEFT injection and before distributed preparation. It
+ changes only the Python forward implementation on each existing PEFT Linear;
+ parameters, children, hooks, and state-dict paths remain owned by that same
+ module object.
+
+ Returns:
+ Number of adapted Q/K/V/output projections configured across both stacks.
+ """
+ max_tokens = _positive_int(max_tokens, "max_tokens")
+ configured = 0
+ for block_name, block in _h3_repeated_blocks(transformer):
+ attention = getattr(block, "attn", None)
+ if not isinstance(attention, nn.Module):
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn expected nn.Module, received "
+ f"{type(attention).__name__}"
+ )
+ if getattr(attention, "fused_projections", False):
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn must install LoRA chunking before "
+ "fusing its projections"
+ )
+ to_out = getattr(attention, "to_out", None)
+ if not isinstance(to_out, nn.ModuleList) or len(to_out) != 2:
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn.to_out expected two-entry ModuleList, "
+ f"received {type(to_out).__name__}"
+ )
+ projections = (
+ ("to_q", getattr(attention, "to_q", None)),
+ ("to_k", getattr(attention, "to_k", None)),
+ ("to_v", getattr(attention, "to_v", None)),
+ ("to_out.0", to_out[0]),
+ )
+ for projection_name, projection in projections:
+ path = f"{block_name}.attn.{projection_name}"
+ if isinstance(projection, _ChunkedLoraLinear):
+ if projection.flow_factory_max_tokens != max_tokens:
+ raise ValueError(
+ f"MiniMax H3 {path} already uses max_tokens="
+ f"{projection.flow_factory_max_tokens}, received conflicting "
+ f"{max_tokens}"
+ )
+ configured += 1
+ continue
+ if not isinstance(projection, LoraLinear):
+ continue
+ if type(projection) is not LoraLinear:
+ raise TypeError(
+ f"MiniMax H3 {path} expected the standard PEFT Linear before "
+ f"chunking, received {type(projection).__name__}"
+ )
+ if "forward" in projection.__dict__:
+ raise TypeError(
+ f"MiniMax H3 {path} must not shadow PEFT Linear.forward on the instance"
+ )
+ if getattr(projection, "_compiled_call_impl", None) is not None:
+ raise TypeError(
+ f"MiniMax H3 {path} must install LoRA chunking before torch.compile"
+ )
+ if type(getattr(projection, "base_layer", None)) is not nn.Linear:
+ raise TypeError(
+ f"MiniMax H3 {path} expected an exact nn.Linear base layer, received "
+ f"{type(getattr(projection, 'base_layer', None)).__name__}"
+ )
+ if any(getattr(projection, "use_dora", {}).values()):
+ raise TypeError(f"MiniMax H3 {path} LoRA chunking does not support DoRA")
+ if getattr(projection, "lora_variant", {}):
+ raise TypeError(f"MiniMax H3 {path} LoRA chunking supports only vanilla LoRA")
+ for adapter_name, dropout in projection.lora_dropout.items():
+ if isinstance(dropout, nn.Identity) or (
+ isinstance(dropout, nn.Dropout) and dropout.p == 0.0
+ ):
+ continue
+ raise TypeError(
+ f"MiniMax H3 {path} LoRA chunking requires zero dropout, received "
+ f"adapter={adapter_name!r}, dropout={dropout!r}"
+ )
+ _reject_execution_hooks(projection, path)
+ projection.__class__ = _ChunkedLoraLinear
+ projection.flow_factory_max_tokens = max_tokens
+ configured += 1
+ return configured
+
+
def _h3_repeated_blocks(transformer: nn.Module) -> Iterable[tuple[str, nn.Module]]:
token_refiner = getattr(transformer, "token_refiner", None)
stacks = (
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index ac16a9c04..f84c386b4 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -37,6 +37,7 @@
StackedSampleBatch,
)
from ...scheduler import MiniMaxH3SDEScheduler, SchedulerGroup
+from ...utils.logger_utils import setup_logger
from ..abc import BaseAdapter
from ..checkpointing import CheckpointUnit
from ..output_state import DecodedMediaBatch, EncodedOutputState, OutputStateCodec
@@ -46,6 +47,10 @@
audio_video_output_contract,
)
from ..runtime import ModularPipelineRuntime
+from ._chunking import (
+ H3_MAX_LORA_PROJECTION_TOKENS,
+ install_h3_lora_projection_chunking,
+)
from ._common import apply_forward_process_noise, draw_forward_process_noise
from ._condition import MiniMaxH3ConditionStatePreparer
from ._output import MiniMaxH3AVOutputCodec, validate_h3_encoded_output_geometry
@@ -65,6 +70,7 @@
_H3_PREPROCESS_CACHE_FIELDS = frozenset({"height", "width", "num_frames"})
_H3_PREPROCESS_CACHE_VERSION = "minimax-h3-v2"
+logger = setup_logger(__name__)
_H3_OPTIONAL_AUDIO_REFERENCE_FORMAT = MediaFormat(
type=MediaType.AUDIO,
fps=RateRequirement.NOT_APPLICABLE,
@@ -375,3 +381,45 @@ class MiniMaxH3Ref2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
"video_processor",
"audio_vae",
]
+
+ def apply_lora(
+ self,
+ target_modules: Union[str, List[str]],
+ components: Union[str, List[str]] = "transformer",
+ overwrite: bool = False,
+ ) -> Any:
+ """Apply PEFT and bound Ref2VA projection memory under FSDP2."""
+ component_names = (components,) if isinstance(components, str) else tuple(components)
+ result = super().apply_lora(
+ target_modules=target_modules,
+ components=components,
+ overwrite=overwrite,
+ )
+ if (
+ not result
+ or not self._is_fsdp2()
+ or self.transformer_component_name not in component_names
+ ):
+ return result
+
+ component = self.get_component(self.transformer_component_name)
+ get_base_model = getattr(component, "get_base_model", None)
+ if not callable(get_base_model):
+ raise TypeError(
+ "MiniMax H3 Ref2VA FSDP2 LoRA chunking expected a PEFT component "
+ f"with get_base_model(), received {type(component).__name__}"
+ )
+ transformer = get_base_model()
+ if not isinstance(transformer, torch.nn.Module):
+ raise TypeError(
+ "MiniMax H3 Ref2VA FSDP2 LoRA chunking expected an nn.Module base, "
+ f"received {type(transformer).__name__}"
+ )
+ configured = install_h3_lora_projection_chunking(transformer)
+ logger.info(
+ "Enabled token-chunked PEFT projections for MiniMax H3 Ref2VA FSDP2: "
+ "configured=%d, max_tokens=%d",
+ configured,
+ H3_MAX_LORA_PROJECTION_TOKENS,
+ )
+ return result
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index 08ef21558..98f6a942a 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -16,15 +16,18 @@
import pytest
import torch
+from peft import LoraConfig, get_peft_model
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
from flow_factory.models.minimax_h3._chunking import (
_ChunkedFeedForward,
+ _ChunkedLoraLinear,
_ChunkedRMSNorm,
install_h3_attention_norm_chunking,
install_h3_feed_forward_chunking,
+ install_h3_lora_projection_chunking,
)
@@ -60,8 +63,12 @@ class BlockFake(nn.Module):
def __init__(self) -> None:
super().__init__()
self.attn = nn.Module()
+ self.attn.to_q = nn.Linear(5, 5, bias=False)
+ self.attn.to_k = nn.Linear(5, 5, bias=False)
+ self.attn.to_v = nn.Linear(5, 5, bias=False)
self.attn.norm_q = nn.RMSNorm(5)
self.attn.norm_k = nn.RMSNorm(5)
+ self.attn.to_out = nn.ModuleList([nn.Linear(5, 5, bias=False), nn.Dropout(0.0)])
self.ff = FeedForwardFake()
@@ -80,6 +87,23 @@ def _feed_forwards(transformer: TransformerFake) -> list[nn.Module]:
]
+def _lora_transformer(*, lora_dropout: float = 0.0) -> nn.Module:
+ model = get_peft_model(
+ TransformerFake(),
+ LoraConfig(
+ r=2,
+ lora_alpha=4,
+ target_modules=["to_q", "to_k", "to_v", "to_out.0"],
+ init_lora_weights="gaussian",
+ lora_dropout=lora_dropout,
+ ),
+ )
+ for module in model.modules():
+ if hasattr(module, "lora_B") and "default" in module.lora_B:
+ nn.init.normal_(module.lora_B["default"].weight)
+ return model
+
+
def test_chunking_preserves_parameter_tree_and_is_idempotent() -> None:
transformer = TransformerFake()
state = deepcopy(transformer.state_dict())
@@ -326,3 +350,123 @@ def test_attention_norm_chunking_rejects_conflicting_reinstallation() -> None:
with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
install_h3_attention_norm_chunking(transformer, max_tokens=8)
+
+
+def test_lora_projection_chunking_preserves_peft_tree_and_is_idempotent() -> None:
+ model = _lora_transformer()
+ transformer = model.get_base_model()
+ state = deepcopy(model.state_dict())
+ keys_before = tuple(state)
+ parameter_ids_before = {name: id(parameter) for name, parameter in model.named_parameters()}
+
+ assert install_h3_lora_projection_chunking(transformer, max_tokens=4) == 20
+ assert install_h3_lora_projection_chunking(transformer, max_tokens=4) == 20
+
+ projections = [
+ projection
+ for block in (
+ *transformer.token_refiner.refiner_blocks,
+ *transformer.transformer_blocks,
+ )
+ for projection in (
+ block.attn.to_q,
+ block.attn.to_k,
+ block.attn.to_v,
+ block.attn.to_out[0],
+ )
+ ]
+ assert all(isinstance(projection, _ChunkedLoraLinear) for projection in projections)
+ assert tuple(model.state_dict()) == keys_before
+ assert {
+ name: id(parameter) for name, parameter in model.named_parameters()
+ } == parameter_ids_before
+ model.load_state_dict(state, strict=True)
+
+
+def test_lora_projection_chunking_preserves_remainder_forward_backward() -> None:
+ torch.manual_seed(29)
+ direct_model = _lora_transformer().double()
+ chunked_model = deepcopy(direct_model)
+ chunked_transformer = chunked_model.get_base_model()
+ install_h3_lora_projection_chunking(chunked_transformer, max_tokens=4)
+ direct = direct_model.get_base_model().transformer_blocks[0].attn.to_k
+ chunked = chunked_transformer.transformer_blocks[0].attn.to_k
+ chunk_sizes: list[int] = []
+ handle = chunked.base_layer.register_forward_pre_hook(
+ lambda _module, inputs: chunk_sizes.append(inputs[0].shape[1])
+ )
+ direct_input = torch.randn(2, 9, 5, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+ output_gradient = torch.randn(2, 9, 5, dtype=torch.float64)
+
+ direct_output = direct(direct_input)
+ chunked_output = chunked(chunked_input)
+ direct_output.backward(output_gradient)
+ chunked_output.backward(output_gradient)
+ handle.remove()
+
+ assert chunk_sizes == [4, 4, 1]
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ for direct_parameter, chunked_parameter in zip(direct.parameters(), chunked.parameters()):
+ torch.testing.assert_close(chunked_parameter.grad, direct_parameter.grad)
+
+
+def test_lora_projection_chunking_preserves_mixed_batch_adapter_names() -> None:
+ direct_model = _lora_transformer()
+ chunked_model = deepcopy(direct_model)
+ install_h3_lora_projection_chunking(chunked_model.get_base_model(), max_tokens=4)
+ direct = direct_model.get_base_model().transformer_blocks[0].attn.to_v
+ chunked = chunked_model.get_base_model().transformer_blocks[0].attn.to_v
+ hidden_states = torch.randn(2, 9, 5)
+ adapter_names = ["default", "__base__"]
+
+ direct_output = direct(hidden_states, adapter_names=adapter_names)
+ chunked_output = chunked(hidden_states, adapter_names=adapter_names)
+
+ torch.testing.assert_close(chunked_output, direct_output)
+
+
+def test_lora_projection_chunking_rejects_conflicting_reinstallation() -> None:
+ transformer = _lora_transformer().get_base_model()
+ install_h3_lora_projection_chunking(transformer, max_tokens=4)
+
+ with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
+ install_h3_lora_projection_chunking(transformer, max_tokens=8)
+
+
+def test_lora_projection_chunking_rejects_nonzero_dropout() -> None:
+ transformer = _lora_transformer(lora_dropout=0.1).get_base_model()
+
+ with pytest.raises(TypeError, match="requires zero dropout"):
+ install_h3_lora_projection_chunking(transformer, max_tokens=4)
+
+
+def test_lora_projection_chunking_rejects_existing_execution_hooks() -> None:
+ transformer = _lora_transformer().get_base_model()
+ transformer.transformer_blocks[0].attn.to_q.register_forward_pre_hook(
+ lambda _module, _inputs: None
+ )
+
+ with pytest.raises(TypeError, match="must be configured before execution hooks"):
+ install_h3_lora_projection_chunking(transformer, max_tokens=4)
+
+
+def test_lora_projection_chunking_preserves_non_reentrant_checkpoint_backward() -> None:
+ direct_model = _lora_transformer().double()
+ chunked_model = deepcopy(direct_model)
+ install_h3_lora_projection_chunking(chunked_model.get_base_model(), max_tokens=4)
+ direct = direct_model.get_base_model().transformer_blocks[0].attn.to_q
+ chunked = chunked_model.get_base_model().transformer_blocks[0].attn.to_q
+ direct_input = torch.randn(2, 9, 5, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+
+ direct_output = direct(direct_input)
+ chunked_output = checkpoint(chunked, chunked_input, use_reentrant=False)
+ direct_output.sum().backward()
+ chunked_output.sum().backward()
+
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ for direct_parameter, chunked_parameter in zip(direct.parameters(), chunked.parameters()):
+ torch.testing.assert_close(chunked_parameter.grad, direct_parameter.grad)
From 182262a0d0ff88de8bdddbc0386955e9b5fae741 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 00:29:29 +0800
Subject: [PATCH 56/76] [models] fix: avoid H3 chunk assembly copies
---
.agents/knowledge/topics/fix_patterns.md | 14 +++
.../models/minimax_h3/_chunking.py | 90 +++++++++++--------
tests/models/minimax_h3/test_chunking.py | 19 ++++
3 files changed, 87 insertions(+), 36 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 0861496c9..c5ae0a635 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -609,6 +609,20 @@ Based on the fix type, write the fix entry to the appropriate document:
recreate the same peak through an avoidable second full-size result.
- **Related Constraint**: N/A
+### Token chunk aggregation must not duplicate the complete packed output
+- **Date**: 2026-08-31
+- **Symptom**: After H3 Ref2VA FSDP2 crossed the adapted projection peak, training exhausted a
+ 95 GiB device when Q RMSNorm's chunk aggregation requested a 190 MiB output with only
+ 154--174 MiB free.
+- **Root Cause**: The chunk executors retained a list whose outputs already totaled the complete
+ packed tensor, then `torch.cat` allocated a second complete tensor to assemble that list.
+- **Fix**: H3 feed-forward, Q/K RMSNorm, and PEFT projection chunk executors now allocate their final
+ output once and copy each non-overlapping token slice into it. CopySlices preserves input and
+ parameter gradients, including nested non-reentrant checkpoint replay, without a concatenation copy.
+- **Lesson**: Bounding each operator invocation is insufficient if aggregation recreates a full-size
+ peak. Treat chunk assembly as part of the memory contract and retain exactly one final output.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index dc65a853e..f169cef0c 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -16,7 +16,7 @@
from __future__ import annotations
-from typing import Iterable
+from typing import Callable, Iterable
import torch
from peft.tuners.lora.layer import Linear as LoraLinear
@@ -55,23 +55,23 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
)
if hidden_states.shape[1] <= self.max_tokens:
return self._forward_chunk(hidden_states)
- chunks = hidden_states.split(self.max_tokens, dim=1)
if torch.is_grad_enabled():
- return torch.cat(
- [
- checkpoint(
- self._forward_chunk,
- chunk,
- use_reentrant=False,
- preserve_rng_state=True,
- )
- for chunk in chunks
- ],
- dim=1,
+ return _assemble_sequence_chunks(
+ hidden_states,
+ self.max_tokens,
+ lambda chunk: checkpoint(
+ self._forward_chunk,
+ chunk,
+ use_reentrant=False,
+ preserve_rng_state=True,
+ ),
+ operation="feed-forward",
)
- return torch.cat(
- [self._forward_chunk(chunk) for chunk in chunks],
- dim=1,
+ return _assemble_sequence_chunks(
+ hidden_states,
+ self.max_tokens,
+ self._forward_chunk,
+ operation="feed-forward",
)
@@ -103,9 +103,11 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
)
if hidden_states.shape[1] <= self.max_tokens:
return self._forward_chunk(hidden_states)
- return torch.cat(
- [self._forward_chunk(chunk) for chunk in hidden_states.split(self.max_tokens, dim=1)],
- dim=1,
+ return _assemble_sequence_chunks(
+ hidden_states,
+ self.max_tokens,
+ self._forward_chunk,
+ operation="RMSNorm",
)
@@ -123,26 +125,42 @@ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
return super().forward(hidden_states, *args, **kwargs)
forward_chunk = super().forward
- output = None
- offset = 0
- for chunk in hidden_states.split(max_tokens, dim=1):
- chunk_output = forward_chunk(chunk, *args, **kwargs)
- if chunk_output.shape[:-1] != chunk.shape[:-1]:
- raise RuntimeError(
- "MiniMax H3 LoRA projection must preserve input prefix dimensions, "
- f"received input={tuple(chunk.shape)}, output={tuple(chunk_output.shape)}"
- )
- if output is None:
- output = chunk_output.new_empty((*hidden_states.shape[:-1], chunk_output.shape[-1]))
- output.narrow(1, offset, chunk_output.shape[1]).copy_(chunk_output)
- offset += chunk_output.shape[1]
+ return _assemble_sequence_chunks(
+ hidden_states,
+ max_tokens,
+ lambda chunk: forward_chunk(chunk, *args, **kwargs),
+ operation="LoRA projection",
+ )
+
- if output is None or offset != hidden_states.shape[1]:
+def _assemble_sequence_chunks(
+ hidden_states: torch.Tensor,
+ max_tokens: int,
+ forward_chunk: Callable[[torch.Tensor], torch.Tensor],
+ *,
+ operation: str,
+) -> torch.Tensor:
+ """Write token-local chunk results into one final allocation."""
+ output = None
+ offset = 0
+ for chunk in hidden_states.split(max_tokens, dim=1):
+ chunk_output = forward_chunk(chunk)
+ if chunk_output.shape[:-1] != chunk.shape[:-1]:
raise RuntimeError(
- "MiniMax H3 LoRA projection failed to assemble every input token, "
- f"received input_tokens={hidden_states.shape[1]}, output_tokens={offset}"
+ f"MiniMax H3 {operation} must preserve input prefix dimensions, "
+ f"received input={tuple(chunk.shape)}, output={tuple(chunk_output.shape)}"
)
- return output
+ if output is None:
+ output = chunk_output.new_empty((*hidden_states.shape[:-1], chunk_output.shape[-1]))
+ output.narrow(1, offset, chunk_output.shape[1]).copy_(chunk_output)
+ offset += chunk_output.shape[1]
+
+ if output is None or offset != hidden_states.shape[1]:
+ raise RuntimeError(
+ f"MiniMax H3 {operation} failed to assemble every input token, "
+ f"received input_tokens={hidden_states.shape[1]}, output_tokens={offset}"
+ )
+ return output
def install_h3_feed_forward_chunking(
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index 98f6a942a..740ef1766 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -352,6 +352,25 @@ def test_attention_norm_chunking_rejects_conflicting_reinstallation() -> None:
install_h3_attention_norm_chunking(transformer, max_tokens=8)
+def test_chunked_token_local_operations_avoid_full_output_concatenation(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ transformer = TransformerFake()
+ install_h3_feed_forward_chunking(transformer, max_tokens=4)
+ install_h3_attention_norm_chunking(transformer, max_tokens=4)
+
+ monkeypatch.setattr(
+ torch,
+ "cat",
+ lambda *args, **kwargs: pytest.fail("chunk aggregation must not call torch.cat"),
+ )
+
+ hidden_states = torch.randn(2, 9, 5, requires_grad=True)
+ feed_forward = transformer.transformer_blocks[0].ff(hidden_states)
+ normalized = transformer.transformer_blocks[0].attn.norm_q(hidden_states)
+ (feed_forward + normalized).sum().backward()
+
+
def test_lora_projection_chunking_preserves_peft_tree_and_is_idempotent() -> None:
model = _lora_transformer()
transformer = model.get_base_model()
From 46a3c163b87bb5bfc4f6b1355e178e5608edfced Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 00:44:13 +0800
Subject: [PATCH 57/76] [models] fix: bound H3 rotary embedding memory
---
.agents/knowledge/topics/fix_patterns.md | 14 ++
.../models/minimax_h3/_chunking.py | 169 ++++++++++++++++++
.../models/minimax_h3/adapters.py | 16 +-
.../models/minimax_h3/dependency.py | 16 +-
tests/models/minimax_h3/test_chunking.py | 127 +++++++++++++
tests/models/minimax_h3/test_modular_core.py | 23 +++
6 files changed, 362 insertions(+), 3 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index c5ae0a635..6ebcce643 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -623,6 +623,20 @@ Based on the fix type, write the fix entry to the appropriate document:
peak. Treat chunk assembly as part of the memory contract and retain exactly one final output.
- **Related Constraint**: N/A
+### Rotary embedding should rotate packed rows in bounded slices
+- **Date**: 2026-08-31
+- **Symptom**: MiniMax H3 Ref2VA FSDP2 GRPO crossed projection, normalization, and aggregation
+ peaks but exhausted a 95 GiB device when rotary embedding requested a 144 MiB elementwise
+ product with only 74--134 MiB free.
+- **Root Cause**: Diffusers built rotate-half, cosine-product, sine-product, sum, and concatenation
+ intermediates over all 13,889 query or key rows simultaneously.
+- **Fix**: Ref2VA FSDP2 installs an instance-local H3 attention processor before distributed
+ preparation. It preserves projection, backend, and output behavior while applying Q/K rotary
+ embedding in aligned 1,024-token slices directly into one final output allocation.
+- **Lesson**: Positional rotation is row-local even when the following attention is global. Bound
+ its elementwise intermediates independently and keep cosine/sine slices aligned with token rows.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index f169cef0c..3a6ad366a 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -23,9 +23,12 @@
from torch import nn
from torch.utils.checkpoint import checkpoint
+from .dependency import require_minimax_h3_support
+
H3_MAX_FEED_FORWARD_TOKENS = 1024
H3_MAX_ATTENTION_NORM_TOKENS = 1024
H3_MAX_LORA_PROJECTION_TOKENS = 1024
+H3_MAX_ROTARY_TOKENS = 1024
class _ChunkedFeedForward(nn.Module):
@@ -133,6 +136,126 @@ def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
)
+class _ChunkedH3AttnProcessor:
+ """Preserve the upstream H3 attention contract with bounded rotary work."""
+
+ _attention_backend = None
+ _parallel_config = None
+ flow_factory_max_tokens: int
+ flow_factory_dispatch_attention_fn: Callable[..., torch.Tensor]
+
+ def __call__(
+ self,
+ attn: nn.Module,
+ hidden_states: torch.Tensor,
+ rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ if attn.fused_projections:
+ query, key, value = attn.to_qkv(hidden_states).chunk(3, dim=-1)
+ else:
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ query = query.unflatten(-1, (attn.heads, -1))
+ key = key.unflatten(-1, (attn.heads, -1))
+ value = value.unflatten(-1, (attn.heads, -1))
+
+ query = attn.norm_q(query)
+ key = attn.norm_k(key)
+
+ if rotary_emb is not None:
+ query = _apply_h3_rotary_chunks(
+ query,
+ *rotary_emb,
+ max_tokens=self.flow_factory_max_tokens,
+ )
+ key = _apply_h3_rotary_chunks(
+ key,
+ *rotary_emb,
+ max_tokens=self.flow_factory_max_tokens,
+ )
+
+ hidden_states = self.flow_factory_dispatch_attention_fn(
+ query,
+ key,
+ value,
+ attn_mask=attention_mask,
+ dropout_p=0.0,
+ is_causal=False,
+ backend=self._attention_backend,
+ parallel_config=self._parallel_config,
+ )
+ hidden_states = hidden_states.flatten(2, 3).type_as(query)
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+ return hidden_states
+
+
+def _apply_h3_rotary_chunks(
+ hidden_states: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ *,
+ max_tokens: int,
+) -> torch.Tensor:
+ """Apply row-local H3 rotary embedding without full-sequence temporaries."""
+ max_tokens = _positive_int(max_tokens, "max_tokens")
+ if hidden_states.ndim != 4:
+ raise ValueError(
+ "MiniMax H3 rotary embedding expected [batch, tokens, heads, head_dim], "
+ f"received shape={tuple(hidden_states.shape)}"
+ )
+ if cos.ndim != 2 or sin.shape != cos.shape:
+ raise ValueError(
+ "MiniMax H3 rotary embedding expected matching [tokens, rotary_dim] cos/sin, "
+ f"received cos={tuple(cos.shape)}, sin={tuple(sin.shape)}"
+ )
+ if cos.shape[0] != hidden_states.shape[1]:
+ raise ValueError(
+ "MiniMax H3 rotary embedding sequence mismatch, received "
+ f"hidden_tokens={hidden_states.shape[1]}, rotary_tokens={cos.shape[0]}"
+ )
+ rotary_dim = cos.shape[-1]
+ if rotary_dim < 2 or rotary_dim % 2 or rotary_dim > hidden_states.shape[-1]:
+ raise ValueError(
+ "MiniMax H3 rotary_dim must be positive, even, and no larger than head_dim, "
+ f"received rotary_dim={rotary_dim}, head_dim={hidden_states.shape[-1]}"
+ )
+
+ cos = cos.to(hidden_states.dtype)
+ sin = sin.to(hidden_states.dtype)
+ if hidden_states.shape[1] <= max_tokens:
+ return _apply_h3_rotary_chunk(hidden_states, cos, sin)
+
+ cos_chunks = iter(cos.split(max_tokens, dim=0))
+ sin_chunks = iter(sin.split(max_tokens, dim=0))
+ return _assemble_sequence_chunks(
+ hidden_states,
+ max_tokens,
+ lambda chunk: _apply_h3_rotary_chunk(chunk, next(cos_chunks), next(sin_chunks)),
+ operation="rotary embedding",
+ )
+
+
+def _apply_h3_rotary_chunk(
+ hidden_states: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+) -> torch.Tensor:
+ """Apply the upstream rotate-half convention to one aligned token chunk."""
+ rotary_dim = cos.shape[-1]
+ hidden_states_rotary = hidden_states[..., :rotary_dim]
+ hidden_states_pass = hidden_states[..., rotary_dim:]
+ cos = cos[None, :, None, :]
+ sin = sin[None, :, None, :]
+ first, second = hidden_states_rotary.chunk(2, dim=-1)
+ hidden_states_rotated = torch.cat((-second, first), dim=-1)
+ hidden_states_rotary = hidden_states_rotary * cos + hidden_states_rotated * sin
+ return torch.cat((hidden_states_rotary, hidden_states_pass), dim=-1).contiguous()
+
+
def _assemble_sequence_chunks(
hidden_states: torch.Tensor,
max_tokens: int,
@@ -365,6 +488,52 @@ def install_h3_lora_projection_chunking(
return configured
+def install_h3_rotary_chunking(
+ transformer: nn.Module,
+ *,
+ max_tokens: int = H3_MAX_ROTARY_TOKENS,
+) -> int:
+ """Install the bounded processor on every standard H3 attention instance."""
+ max_tokens = _positive_int(max_tokens, "max_tokens")
+ symbols = require_minimax_h3_support()
+ processor_class = symbols.MiniMaxH3AttnProcessor
+ configured = 0
+ for block_name, block in _h3_repeated_blocks(transformer):
+ attention = getattr(block, "attn", None)
+ if not isinstance(attention, nn.Module):
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn expected nn.Module, received "
+ f"{type(attention).__name__}"
+ )
+ processor = getattr(attention, "processor", None)
+ if isinstance(processor, _ChunkedH3AttnProcessor):
+ if processor.flow_factory_max_tokens != max_tokens:
+ raise ValueError(
+ f"MiniMax H3 {block_name}.attn.processor already uses max_tokens="
+ f"{processor.flow_factory_max_tokens}, received conflicting {max_tokens}"
+ )
+ configured += 1
+ continue
+ if type(processor) is not processor_class:
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn expected the standard attention processor, "
+ f"received {type(processor).__name__}"
+ )
+ if "__call__" in processor.__dict__:
+ raise TypeError(
+ f"MiniMax H3 {block_name}.attn processor must not shadow __call__ on the instance"
+ )
+ attention_backend = getattr(processor, "_attention_backend", None)
+ parallel_config = getattr(processor, "_parallel_config", None)
+ processor.__class__ = _ChunkedH3AttnProcessor
+ processor.flow_factory_max_tokens = max_tokens
+ processor.flow_factory_dispatch_attention_fn = symbols.dispatch_attention_fn
+ processor._attention_backend = attention_backend
+ processor._parallel_config = parallel_config
+ configured += 1
+ return configured
+
+
def _h3_repeated_blocks(transformer: nn.Module) -> Iterable[tuple[str, nn.Module]]:
token_refiner = getattr(transformer, "token_refiner", None)
stacks = (
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index f84c386b4..411077c05 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -49,7 +49,9 @@
from ..runtime import ModularPipelineRuntime
from ._chunking import (
H3_MAX_LORA_PROJECTION_TOKENS,
+ H3_MAX_ROTARY_TOKENS,
install_h3_lora_projection_chunking,
+ install_h3_rotary_chunking,
)
from ._common import apply_forward_process_noise, draw_forward_process_noise
from ._condition import MiniMaxH3ConditionStatePreparer
@@ -138,7 +140,19 @@ def load_pipeline(self) -> Any:
def build_component_runtime(self) -> ModularPipelineRuntime:
"""Build the workflow-pruned modular runtime."""
- return build_h3_component_runtime(self)
+ runtime = build_h3_component_runtime(self)
+ if self.workflow != "ref2va" or not self._is_fsdp2():
+ return runtime
+ configured = install_h3_rotary_chunking(
+ runtime.get_component(self.transformer_component_name)
+ )
+ logger.info(
+ "Enabled token-chunked rotary embedding for MiniMax H3 Ref2VA FSDP2: "
+ "configured=%d, max_tokens=%d",
+ configured,
+ H3_MAX_ROTARY_TOKENS,
+ )
+ return runtime
def load_scheduler(self) -> MiniMaxH3SDEScheduler:
"""Build the canonical shift-12 video scheduler."""
diff --git a/src/flow_factory/models/minimax_h3/dependency.py b/src/flow_factory/models/minimax_h3/dependency.py
index 2a0f93258..936685ee3 100644
--- a/src/flow_factory/models/minimax_h3/dependency.py
+++ b/src/flow_factory/models/minimax_h3/dependency.py
@@ -15,7 +15,7 @@
import inspect
from dataclasses import dataclass
-from typing import Any, Tuple, Type
+from typing import Any, Callable, Tuple, Type
import torch
@@ -60,6 +60,8 @@ class MiniMaxH3Symbols:
ModularPipeline: Type[Any]
MiniMaxH3ModularPipeline: Type[Any]
MiniMaxH3Blocks: Type[Any]
+ MiniMaxH3AttnProcessor: Type[Any]
+ dispatch_attention_fn: Callable[..., torch.Tensor]
PipelineState: Type[Any]
ResizeStep: Type[Any]
RefSetupStep: Type[Any]
@@ -85,7 +87,8 @@ class MiniMaxH3Symbols:
try:
- from diffusers import ModularPipeline
+ from diffusers.models.attention_dispatch import dispatch_attention_fn
+ from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3AttnProcessor
from diffusers.modular_pipelines.minimax_h3.before_denoise import (
MiniMaxH3FL2VAPrepareLatentsStep,
MiniMaxH3NoKeyframeAnchorsStep,
@@ -121,10 +124,14 @@ class MiniMaxH3Symbols:
)
from diffusers.modular_pipelines.modular_pipeline import PipelineState
+ from diffusers import ModularPipeline
+
_SYMBOLS = MiniMaxH3Symbols(
ModularPipeline=ModularPipeline,
MiniMaxH3ModularPipeline=MiniMaxH3ModularPipeline,
MiniMaxH3Blocks=MiniMaxH3Blocks,
+ MiniMaxH3AttnProcessor=MiniMaxH3AttnProcessor,
+ dispatch_attention_fn=dispatch_attention_fn,
PipelineState=PipelineState,
ResizeStep=MiniMaxH3ResizeStep,
RefSetupStep=MiniMaxH3Ref2VASetupStep,
@@ -172,6 +179,11 @@ def require_minimax_h3_support() -> MiniMaxH3Symbols:
def _probe_symbol_bundle(symbols: MiniMaxH3Symbols) -> None:
+ processor = symbols.MiniMaxH3AttnProcessor()
+ if not callable(processor):
+ raise TypeError("MiniMaxH3AttnProcessor instance must be callable")
+ if not callable(symbols.dispatch_attention_fn):
+ raise TypeError("dispatch_attention_fn must be callable")
state_values = {"probe": object()}
try:
state = symbols.PipelineState(values=state_values)
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index 740ef1766..447e8d0ac 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -16,18 +16,26 @@
import pytest
import torch
+from diffusers.models.transformers.transformer_minimax_h3 import (
+ MiniMaxH3AttnProcessor,
+ _apply_rotary_emb,
+)
from peft import LoraConfig, get_peft_model
from torch import nn
from torch.nn import functional as F
from torch.utils.checkpoint import checkpoint
from flow_factory.models.minimax_h3._chunking import (
+ _apply_h3_rotary_chunk,
+ _apply_h3_rotary_chunks,
_ChunkedFeedForward,
+ _ChunkedH3AttnProcessor,
_ChunkedLoraLinear,
_ChunkedRMSNorm,
install_h3_attention_norm_chunking,
install_h3_feed_forward_chunking,
install_h3_lora_projection_chunking,
+ install_h3_rotary_chunking,
)
@@ -63,12 +71,15 @@ class BlockFake(nn.Module):
def __init__(self) -> None:
super().__init__()
self.attn = nn.Module()
+ self.attn.heads = 1
+ self.attn.fused_projections = False
self.attn.to_q = nn.Linear(5, 5, bias=False)
self.attn.to_k = nn.Linear(5, 5, bias=False)
self.attn.to_v = nn.Linear(5, 5, bias=False)
self.attn.norm_q = nn.RMSNorm(5)
self.attn.norm_k = nn.RMSNorm(5)
self.attn.to_out = nn.ModuleList([nn.Linear(5, 5, bias=False), nn.Dropout(0.0)])
+ self.attn.processor = MiniMaxH3AttnProcessor()
self.ff = FeedForwardFake()
@@ -489,3 +500,119 @@ def test_lora_projection_chunking_preserves_non_reentrant_checkpoint_backward()
torch.testing.assert_close(chunked_input.grad, direct_input.grad)
for direct_parameter, chunked_parameter in zip(direct.parameters(), chunked.parameters()):
torch.testing.assert_close(chunked_parameter.grad, direct_parameter.grad)
+
+
+def test_rotary_chunking_preserves_remainder_forward_backward(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ torch.manual_seed(31)
+ direct_input = torch.randn(2, 9, 3, 8, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+ direct_cos = torch.randn(9, 6, dtype=torch.float64, requires_grad=True)
+ chunked_cos = direct_cos.detach().clone().requires_grad_(True)
+ direct_sin = torch.randn(9, 6, dtype=torch.float64, requires_grad=True)
+ chunked_sin = direct_sin.detach().clone().requires_grad_(True)
+ output_gradient = torch.randn(2, 9, 3, 8, dtype=torch.float64)
+ chunk_sizes: list[int] = []
+
+ def recording_rotary_chunk(hidden_states, chunk_cos, chunk_sin):
+ chunk_sizes.append(hidden_states.shape[1])
+ return _apply_h3_rotary_chunk(hidden_states, chunk_cos, chunk_sin)
+
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3._chunking._apply_h3_rotary_chunk",
+ recording_rotary_chunk,
+ )
+
+ direct_output = _apply_rotary_emb(direct_input, direct_cos, direct_sin)
+ chunked_output = _apply_h3_rotary_chunks(
+ chunked_input,
+ chunked_cos,
+ chunked_sin,
+ max_tokens=4,
+ )
+ direct_output.backward(output_gradient)
+ chunked_output.backward(output_gradient)
+
+ assert chunk_sizes == [4, 4, 1]
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ torch.testing.assert_close(chunked_cos.grad, direct_cos.grad)
+ torch.testing.assert_close(chunked_sin.grad, direct_sin.grad)
+
+
+def test_rotary_chunking_preserves_processor_and_parameter_tree() -> None:
+ transformer = TransformerFake()
+ attention_backend = object()
+ parallel_config = object()
+ first_processor = transformer.token_refiner.refiner_blocks[0].attn.processor
+ first_processor._attention_backend = attention_backend
+ first_processor._parallel_config = parallel_config
+ state = deepcopy(transformer.state_dict())
+ keys_before = tuple(state)
+ parameter_ids_before = {
+ name: id(parameter) for name, parameter in transformer.named_parameters()
+ }
+ processor_ids_before = [
+ id(block.attn.processor)
+ for block in (
+ *transformer.token_refiner.refiner_blocks,
+ *transformer.transformer_blocks,
+ )
+ ]
+
+ assert install_h3_rotary_chunking(transformer, max_tokens=4) == 5
+ assert install_h3_rotary_chunking(transformer, max_tokens=4) == 5
+
+ processors = [
+ block.attn.processor
+ for block in (
+ *transformer.token_refiner.refiner_blocks,
+ *transformer.transformer_blocks,
+ )
+ ]
+ assert all(isinstance(processor, _ChunkedH3AttnProcessor) for processor in processors)
+ assert [id(processor) for processor in processors] == processor_ids_before
+ assert processors[0]._attention_backend is attention_backend
+ assert processors[0]._parallel_config is parallel_config
+ assert tuple(transformer.state_dict()) == keys_before
+ assert {
+ name: id(parameter) for name, parameter in transformer.named_parameters()
+ } == parameter_ids_before
+ transformer.load_state_dict(state, strict=True)
+
+
+def test_rotary_chunking_processor_preserves_attention_forward_backward() -> None:
+ torch.manual_seed(37)
+ direct_transformer = TransformerFake().double()
+ chunked_transformer = deepcopy(direct_transformer)
+ install_h3_rotary_chunking(chunked_transformer, max_tokens=4)
+ direct = direct_transformer.transformer_blocks[0].attn
+ chunked = chunked_transformer.transformer_blocks[0].attn
+ direct_input = torch.randn(2, 9, 5, dtype=torch.float64, requires_grad=True)
+ chunked_input = direct_input.detach().clone().requires_grad_(True)
+ cos = torch.randn(9, 4, dtype=torch.float64)
+ sin = torch.randn(9, 4, dtype=torch.float64)
+ output_gradient = torch.randn(2, 9, 5, dtype=torch.float64)
+
+ direct_output = direct.processor(direct, direct_input, (cos, sin))
+ chunked_output = checkpoint(
+ lambda value: chunked.processor(chunked, value, (cos, sin)),
+ chunked_input,
+ use_reentrant=False,
+ )
+ direct_output.backward(output_gradient)
+ chunked_output.backward(output_gradient)
+
+ torch.testing.assert_close(chunked_output, direct_output)
+ torch.testing.assert_close(chunked_input.grad, direct_input.grad)
+ for direct_parameter, chunked_parameter in zip(direct.parameters(), chunked.parameters()):
+ torch.testing.assert_close(chunked_parameter.grad, direct_parameter.grad)
+
+
+def test_rotary_chunking_rejects_conflicting_reinstallation() -> None:
+ transformer = TransformerFake()
+ install_h3_rotary_chunking(transformer, max_tokens=4)
+
+ with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
+ install_h3_rotary_chunking(transformer, max_tokens=8)
diff --git a/tests/models/minimax_h3/test_modular_core.py b/tests/models/minimax_h3/test_modular_core.py
index 160b02397..5146fa77c 100644
--- a/tests/models/minimax_h3/test_modular_core.py
+++ b/tests/models/minimax_h3/test_modular_core.py
@@ -202,6 +202,15 @@ class FakeMiniMaxH3Blocks:
}
+class FakeMiniMaxH3AttnProcessor:
+ def __call__(self, *args, **kwargs):
+ return None
+
+
+def fake_dispatch_attention_fn(*args, **kwargs):
+ return None
+
+
class FakeReference:
pass
@@ -210,6 +219,8 @@ class FakeReference:
class FakeSymbols:
MiniMaxH3ModularPipeline: type = FakeModularPipeline
MiniMaxH3Blocks: type = FakeMiniMaxH3Blocks
+ MiniMaxH3AttnProcessor: type = FakeMiniMaxH3AttnProcessor
+ dispatch_attention_fn: object = fake_dispatch_attention_fn
PipelineState: type = FakePipelineState
ResizeStep: type = ResizeStep
RefSetupStep: type = RefSetupStep
@@ -1052,6 +1063,7 @@ def __call__(self):
[
"PipelineState",
"MiniMaxH3Blocks",
+ "MiniMaxH3AttnProcessor",
"PrepareLatentsStep",
"Ref2VATextEncoderStep",
"SetTimestepsStep",
@@ -1075,6 +1087,17 @@ def test_dependency_probe_rejects_incompatible_api_with_actionable_requirement(
assert broken_field in message
+def test_dependency_probe_rejects_non_callable_attention_dispatch(monkeypatch):
+ from flow_factory.models.minimax_h3 import dependency
+
+ bundle = dataclasses.replace(FakeSymbols(), dispatch_attention_fn=object())
+ monkeypatch.setattr(dependency, "_SYMBOLS", bundle)
+ monkeypatch.setattr(dependency, "_IMPORT_ERROR", None)
+
+ with pytest.raises(ImportError, match="dispatch_attention_fn must be callable"):
+ dependency.require_minimax_h3_support()
+
+
def test_pyproject_requires_released_h3_diffusers():
text = Path("pyproject.toml").read_text()
requirement = "diffusers>=0.40.0"
From 71526e06301a8197b28c41eb8947e1b72756fc19 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 00:48:59 +0800
Subject: [PATCH 58/76] [models] fix: release H3 attention inputs promptly
---
.agents/knowledge/topics/fix_patterns.md | 14 +++++++++++
.../models/minimax_h3/_chunking.py | 4 +++-
tests/models/minimax_h3/test_chunking.py | 23 +++++++++++++++++++
3 files changed, 40 insertions(+), 1 deletion(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 6ebcce643..c87a7b68d 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -637,6 +637,20 @@ Based on the fix type, write the fix entry to the appropriate document:
its elementwise intermediates independently and keep cosine/sine slices aligned with token rows.
- **Related Constraint**: N/A
+### Checkpointed attention should release QKV before its output projection
+- **Date**: 2026-08-31
+- **Symptom**: After bounded rotary embedding passed, H3 Ref2VA FSDP2 GRPO missed a 144 MiB
+ output-projection allocation by roughly 10 MiB while the attention result was already available.
+- **Root Cause**: The attention processor kept strong Python references to complete Q/K/V tensors
+ through the output projection. FSDP2 activation checkpointing had discarded their saved-tensor
+ storage requirements, but the local variables still extended their lifetime.
+- **Fix**: The bounded H3 processor captures the output dtype, releases Q/K/V immediately after
+ attention dispatch, and only then flattens and projects the attention result.
+- **Lesson**: Under activation checkpointing, autograd may no longer own an intermediate while a
+ Python local still does. End large tensor lifetimes at their last semantic use before allocating
+ the next full-size result.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 3a6ad366a..e55920816 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -177,6 +177,7 @@ def __call__(
max_tokens=self.flow_factory_max_tokens,
)
+ output_dtype = query.dtype
hidden_states = self.flow_factory_dispatch_attention_fn(
query,
key,
@@ -187,7 +188,8 @@ def __call__(
backend=self._attention_backend,
parallel_config=self._parallel_config,
)
- hidden_states = hidden_states.flatten(2, 3).type_as(query)
+ del query, key, value
+ hidden_states = hidden_states.flatten(2, 3).to(dtype=output_dtype)
hidden_states = attn.to_out[0](hidden_states)
hidden_states = attn.to_out[1](hidden_states)
return hidden_states
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index 447e8d0ac..f17f5c3db 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -13,6 +13,7 @@
# limitations under the License.
from copy import deepcopy
+from weakref import ref
import pytest
import torch
@@ -616,3 +617,25 @@ def test_rotary_chunking_rejects_conflicting_reinstallation() -> None:
with pytest.raises(ValueError, match="already uses max_tokens=4.*conflicting 8"):
install_h3_rotary_chunking(transformer, max_tokens=8)
+
+
+def test_chunked_attention_releases_qkv_before_output_projection() -> None:
+ transformer = TransformerFake()
+ install_h3_rotary_chunking(transformer, max_tokens=4)
+ attention = transformer.transformer_blocks[0].attn
+ qkv_refs = []
+
+ def recording_dispatch(query, key, value, **kwargs):
+ del kwargs
+ qkv_refs.extend((ref(query), ref(key), ref(value)))
+ return torch.zeros_like(query)
+
+ attention.processor.flow_factory_dispatch_attention_fn = recording_dispatch
+
+ def require_released_qkv(_module, _inputs):
+ assert len(qkv_refs) == 3
+ assert all(tensor_ref() is None for tensor_ref in qkv_refs)
+
+ handle = attention.to_out[0].register_forward_pre_hook(require_released_qkv)
+ attention.processor(attention, torch.randn(2, 9, 5))
+ handle.remove()
From 84229f6ea495f0fe8f0772436fa29a634b481d36 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 01:19:09 +0800
Subject: [PATCH 59/76] [loading,models] fix: bound H3 FSDP2 checkpoint memory
---
.agents/knowledge/topics/fix_patterns.md | 21 ++++
src/flow_factory/loading/backend.py | 103 +++++++++++++++---
src/flow_factory/models/abc.py | 4 +
.../models/minimax_h3/_chunking.py | 68 ++++++++++++
.../models/minimax_h3/adapters.py | 23 ++++
tests/loading/test_backend_runtime.py | 55 ++++++++++
tests/models/minimax_h3/test_chunking.py | 75 +++++++++++++
.../minimax_h3/test_workflow_adapters.py | 6 +
8 files changed, 340 insertions(+), 15 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index c87a7b68d..ef927458a 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -651,6 +651,27 @@ Based on the fix type, write the fix entry to the appropriate document:
the next full-size result.
- **Related Constraint**: N/A
+### FSDP2 checkpoint and backward-prefetch policies must have independent memory boundaries
+- **Date**: 2026-08-31
+- **Symptom**: MiniMax H3 Ref2VA FSDP2 GRPO reached the training forward only after several
+ operator-level bounds, then exhausted each 95 GiB device during forward activation retention or
+ the first backward all-gather. The final backward failure requested 286 MiB while the 442 MiB
+ feed-forward unit was still resident.
+- **Root Cause**: Accelerate reused the FSDP2 transformer wrap policy for activation checkpointing
+ and checkpointed every direct child of each H3 block, retaining several complete packed-sequence
+ inputs per block. After replacing those boundaries, PyTorch's implicit backward prefetch still
+ overlapped the current feed-forward unit with the next attention unit's all-gather.
+- **Fix**: Ref2VA now installs one non-reentrant checkpoint inside every materialized H3 block after
+ the block's FSDP mixed-precision input cast, with its saved BF16 inputs held in pinned CPU memory.
+ Backend preparation disables Accelerate's duplicate checkpoint owner and opts every prepared
+ FSDP2 unit out of implicit next-unit backward prefetch, so the current unit reshards before its
+ successor gathers. All variant instances are configured only after materialization. A two-rank
+ GRPO sentinel completed two forward/backward/optimizer cycles with stable gradients.
+- **Lesson**: FSDP wrap granularity, activation recomputation, saved-input placement, and collective
+ prefetch are separate memory policies. Keep each boundary explicit; a policy that is correct for
+ parameter sharding may multiply activation lifetimes or overlap adjacent full-parameter units.
+- **Related Constraint**: #9, #20
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/loading/backend.py b/src/flow_factory/loading/backend.py
index 2f68133ca..56d3e3a71 100644
--- a/src/flow_factory/loading/backend.py
+++ b/src/flow_factory/loading/backend.py
@@ -122,30 +122,103 @@ def load_scope(self, role: ComponentRole) -> Iterator[None]:
def prepare(self, *objects: Any) -> Any:
"""Prepare roots and apply adapter-requested FSDP2 communication policy."""
- self._extend_fsdp2_wrap_policy(objects)
- prepared = super().prepare(*objects)
plugin = self.accelerator.state.fsdp_plugin
- if (getattr(plugin, "fsdp_version", 1) or 1) < 2 or not getattr(
- self.adapter, "fsdp2_use_default_stream_unshard", False
- ):
+ use_in_forward_checkpointing = (
+ (getattr(plugin, "fsdp_version", 1) or 1) >= 2
+ and bool(getattr(plugin, "activation_checkpointing", False))
+ and bool(
+ getattr(
+ self.adapter,
+ "fsdp2_use_in_forward_activation_checkpointing",
+ False,
+ )
+ )
+ )
+ if use_in_forward_checkpointing:
+ modules = [candidate for candidate in objects if isinstance(candidate, nn.Module)]
+ configure = getattr(
+ self.adapter,
+ "configure_fsdp2_in_forward_activation_checkpointing",
+ None,
+ )
+ if len(modules) != 1 or not callable(configure):
+ raise TypeError(
+ "adapter-owned FSDP2 activation checkpointing requires exactly one "
+ "model root and a callable configuration hook"
+ )
+ configured = configure(modules[0])
+ if not isinstance(configured, int) or isinstance(configured, bool) or configured < 1:
+ raise TypeError(
+ "adapter-owned FSDP2 activation checkpointing expected a positive "
+ f"configured block count, received {configured!r}"
+ )
+ plugin.activation_checkpointing = False
+ logger.info(
+ "Delegated FSDP2 activation checkpointing to adapter-owned in-forward "
+ "block boundaries: configured=%d",
+ configured,
+ )
+
+ try:
+ self._extend_fsdp2_wrap_policy(objects)
+ prepared = super().prepare(*objects)
+ finally:
+ if use_in_forward_checkpointing:
+ plugin.activation_checkpointing = True
+
+ if (getattr(plugin, "fsdp_version", 1) or 1) < 2:
+ return prepared
+
+ use_default_stream = bool(getattr(self.adapter, "fsdp2_use_default_stream_unshard", False))
+ disable_backward_prefetch = bool(
+ getattr(self.adapter, "fsdp2_disable_backward_prefetch", False)
+ )
+ if not use_default_stream and not disable_backward_prefetch:
return prepared
prepared_objects = prepared if isinstance(prepared, (list, tuple)) else (prepared,)
- configured = []
+ prepared_modules = []
for original, candidate in zip(objects, prepared_objects):
if not isinstance(original, nn.Module):
continue
- configure = getattr(candidate, "_set_unshard_async_op", None)
- if not callable(configure):
+ prepared_modules.append(candidate)
+
+ if not prepared_modules:
+ raise TypeError("FSDP2 memory policy requested without a prepared module")
+
+ if use_default_stream:
+ configured = []
+ for candidate in prepared_modules:
+ configure = getattr(candidate, "_set_unshard_async_op", None)
+ if not callable(configure):
+ raise TypeError(
+ "FSDP2 default-stream unshard requires prepared modules to expose "
+ f"_set_unshard_async_op(), received {type(candidate).__name__}"
+ )
+ configure(True)
+ configured.append(type(candidate).__name__)
+ logger.info("Enabled FSDP2 default-stream unshard for roots: %s", configured)
+
+ if disable_backward_prefetch:
+ configured = 0
+ for root in prepared_modules:
+ for module in root.modules():
+ configure = getattr(module, "set_modules_to_backward_prefetch", None)
+ if not callable(configure):
+ continue
+ # FSDP2 treats an empty list as "use default next-unit prefetch".
+ # Prefetching the current, already-unsharded unit is a public-API
+ # no-op that selects explicit mode and disables that default.
+ configure([module])
+ configured += 1
+ if configured < 1:
raise TypeError(
- "FSDP2 default-stream unshard requires prepared modules to expose "
- f"_set_unshard_async_op(), received {type(candidate).__name__}"
+ "FSDP2 backward-prefetch opt-out requires prepared FSDPModule roots"
)
- configure(True)
- configured.append(type(candidate).__name__)
- if not configured:
- raise TypeError("FSDP2 default-stream unshard requested without a prepared module")
- logger.info("Enabled FSDP2 default-stream unshard for roots: %s", configured)
+ logger.info(
+ "Disabled default FSDP2 next-unit backward prefetch: configured=%d",
+ configured,
+ )
return prepared
def _extend_fsdp2_wrap_policy(self, objects: Sequence[Any]) -> None:
diff --git a/src/flow_factory/models/abc.py b/src/flow_factory/models/abc.py
index 749ea103e..1b81b5536 100644
--- a/src/flow_factory/models/abc.py
+++ b/src/flow_factory/models/abc.py
@@ -236,6 +236,10 @@ class BaseAdapter(ABC):
# Opt in only when FSDP2 communication overlap exceeds the model's activation headroom.
fsdp2_use_default_stream_unshard: ClassVar[bool] = False
fsdp2_additional_wrap_module_names: ClassVar[Tuple[str, ...]] = ()
+ # The adapter may place one checkpoint inside each FSDP-wrapped block forward.
+ fsdp2_use_in_forward_activation_checkpointing: ClassVar[bool] = False
+ # Opt in when backward all-gather overlap exceeds the model's peak headroom.
+ fsdp2_disable_backward_prefetch: ClassVar[bool] = False
supports_ordered_references: ClassVar[bool] = False
preprocess_cache_fields: ClassVar[frozenset[str]] = frozenset()
preprocess_cache_version: ClassVar[str] = ""
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index e55920816..3f346bd5c 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -16,11 +16,13 @@
from __future__ import annotations
+from types import MethodType
from typing import Callable, Iterable
import torch
from peft.tuners.lora.layer import Linear as LoraLinear
from torch import nn
+from torch.autograd.graph import save_on_cpu
from torch.utils.checkpoint import checkpoint
from .dependency import require_minimax_h3_support
@@ -31,6 +33,29 @@
H3_MAX_ROTARY_TOKENS = 1024
+def _checkpoint_h3_block_forward(self: nn.Module, *args, **kwargs):
+ """Checkpoint one block after its FSDP cast and offload saved inputs to CPU."""
+ original_forward = self.__dict__["flow_factory_uncheckpointed_forward"]
+ if not torch.is_grad_enabled():
+ return original_forward(self, *args, **kwargs)
+ input_device_type = next(
+ (argument.device.type for argument in args if isinstance(argument, torch.Tensor)),
+ "cuda",
+ )
+ with save_on_cpu(
+ pin_memory=input_device_type == "cuda" and torch.cuda.is_available(),
+ device_type=input_device_type,
+ ):
+ return checkpoint(
+ original_forward,
+ self,
+ *args,
+ use_reentrant=False,
+ preserve_rng_state=False,
+ **kwargs,
+ )
+
+
class _ChunkedFeedForward(nn.Module):
"""Run one existing feed-forward network over remainder-safe token chunks.
@@ -350,6 +375,49 @@ def install_h3_feed_forward_chunking(
return configured
+def install_h3_in_forward_block_checkpointing(transformer: nn.Module) -> int:
+ """Place one checkpoint inside every repeated block's FSDP execution boundary.
+
+ Diffusers model-level checkpointing surrounds the block call and therefore saves
+ inputs before FSDP2 casts them for mixed-precision compute. Accelerate's generic
+ FSDP2 policy checkpoints each direct child separately. This instance-local
+ forward delegates the complete original block body to one non-reentrant
+ checkpoint after the block's FSDP pre-forward hook has already run. Its saved
+ BF16 block inputs live in pinned CPU memory until their backward replay.
+ """
+ configured = 0
+ for block_name, block in _h3_repeated_blocks(transformer):
+ installed = block.__dict__.get("flow_factory_uncheckpointed_forward")
+ shadowed_forward = block.__dict__.get("forward")
+ if installed is not None:
+ if not (
+ callable(installed)
+ and isinstance(shadowed_forward, MethodType)
+ and shadowed_forward.__func__ is _checkpoint_h3_block_forward
+ ):
+ raise TypeError(
+ f"MiniMax H3 {block_name} has an inconsistent in-forward "
+ "checkpoint installation"
+ )
+ configured += 1
+ continue
+ if shadowed_forward is not None:
+ raise TypeError(
+ f"MiniMax H3 {block_name} must not shadow forward before in-forward "
+ "checkpoint installation"
+ )
+ original_forward = type(block).forward
+ if not callable(original_forward):
+ raise TypeError(
+ f"MiniMax H3 {block_name} expected a callable class forward, received "
+ f"{original_forward!r}"
+ )
+ block.flow_factory_uncheckpointed_forward = original_forward
+ block.forward = MethodType(_checkpoint_h3_block_forward, block)
+ configured += 1
+ return configured
+
+
def install_h3_attention_norm_chunking(
transformer: nn.Module,
*,
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index 411077c05..db5a6a633 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -50,6 +50,7 @@
from ._chunking import (
H3_MAX_LORA_PROJECTION_TOKENS,
H3_MAX_ROTARY_TOKENS,
+ install_h3_in_forward_block_checkpointing,
install_h3_lora_projection_chunking,
install_h3_rotary_chunking,
)
@@ -154,6 +155,26 @@ def build_component_runtime(self) -> ModularPipelineRuntime:
)
return runtime
+ def configure_fsdp2_in_forward_activation_checkpointing(
+ self,
+ model_root: torch.nn.Module,
+ ) -> int:
+ """Install one inner checkpoint on every materialized H3 block variant."""
+ members = getattr(model_root, "members", None)
+ if not isinstance(members, torch.nn.ModuleDict):
+ raise TypeError(
+ "MiniMax H3 FSDP2 checkpointing expected a ModelBundle ModuleDict, "
+ f"received {type(members).__name__}"
+ )
+ configured = 0
+ seen = set()
+ for component in members.values():
+ if id(component) in seen:
+ continue
+ seen.add(id(component))
+ configured += install_h3_in_forward_block_checkpointing(component)
+ return configured
+
def load_scheduler(self) -> MiniMaxH3SDEScheduler:
"""Build the canonical shift-12 video scheduler."""
return build_h3_scheduler(self.config.scheduler_args, shift=12.0)
@@ -375,6 +396,8 @@ class MiniMaxH3Ref2VAAdapter(_MiniMaxH3WorkflowAdapter, BaseAdapter):
workflow: ClassVar[str] = "ref2va"
transformer_component_name: ClassVar[str] = "transformer_ref"
fsdp2_use_default_stream_unshard: ClassVar[bool] = True
+ fsdp2_use_in_forward_activation_checkpointing: ClassVar[bool] = True
+ fsdp2_disable_backward_prefetch: ClassVar[bool] = True
fsdp2_additional_wrap_module_names: ClassVar[Tuple[str, ...]] = (
"_ChunkedFeedForward",
"MiniMaxH3AdaLayerNormModulation",
diff --git a/tests/loading/test_backend_runtime.py b/tests/loading/test_backend_runtime.py
index 14a8ed357..586836d01 100644
--- a/tests/loading/test_backend_runtime.py
+++ b/tests/loading/test_backend_runtime.py
@@ -113,10 +113,14 @@ class _FSDPModuleFake(torch.nn.Linear):
def __init__(self) -> None:
super().__init__(2, 2)
self.unshard_async_op_calls = []
+ self.backward_prefetch_calls = []
def _set_unshard_async_op(self, enabled: bool) -> None:
self.unshard_async_op_calls.append(enabled)
+ def set_modules_to_backward_prefetch(self, modules) -> None:
+ self.backward_prefetch_calls.append(modules)
+
def test_prepare_enables_adapter_requested_fsdp2_default_stream_unshard() -> None:
module = _FSDPModuleFake()
@@ -133,6 +137,20 @@ def test_prepare_enables_adapter_requested_fsdp2_default_stream_unshard() -> Non
assert module.unshard_async_op_calls == [True]
+def test_prepare_disables_default_fsdp2_backward_prefetch_recursively() -> None:
+ module = _FSDPModuleFake()
+ module.child = _FSDPModuleFake()
+ runtime = FSDPBackendLoadRuntime(
+ _AcceleratorFake(efficient=True),
+ _plan(),
+ SimpleNamespace(fsdp2_disable_backward_prefetch=True),
+ )
+
+ assert runtime.prepare(module) is module
+ assert module.backward_prefetch_calls == [[module]]
+ assert module.child.backward_prefetch_calls == [[module.child]]
+
+
def test_prepare_extends_fsdp2_wrap_policy_before_distributed_preparation() -> None:
module = _FSDPModuleFake()
module._no_split_modules = ["TransformerBlock"]
@@ -152,6 +170,43 @@ def test_prepare_extends_fsdp2_wrap_policy_before_distributed_preparation() -> N
]
+def test_prepare_delegates_fsdp2_activation_checkpointing_during_prepare() -> None:
+ module = _FSDPModuleFake()
+ module._no_split_modules = ["TransformerBlock"]
+ accelerator = _AcceleratorFake(efficient=True)
+ plugin = accelerator.state.fsdp_plugin
+ plugin.activation_checkpointing = True
+ prepare_observations = []
+ accelerator.prepare = (
+ lambda *objects: prepare_observations.append(
+ (
+ plugin.activation_checkpointing,
+ tuple(plugin.transformer_cls_names_to_wrap),
+ )
+ )
+ or objects[0]
+ )
+ configured_roots = []
+ runtime = FSDPBackendLoadRuntime(
+ accelerator,
+ _plan(),
+ SimpleNamespace(
+ fsdp2_additional_wrap_module_names=("ChunkedFeedForward",),
+ fsdp2_use_in_forward_activation_checkpointing=True,
+ configure_fsdp2_in_forward_activation_checkpointing=lambda root: (
+ configured_roots.append(root) or 2
+ ),
+ ),
+ )
+
+ assert runtime.prepare(module) is module
+ assert configured_roots == [module]
+ assert prepare_observations == [
+ (False, ("TransformerBlock", "ChunkedFeedForward")),
+ ]
+ assert plugin.activation_checkpointing is True
+
+
@pytest.mark.parametrize(
("fsdp_version", "requested"),
[(1, True), (2, False)],
diff --git a/tests/models/minimax_h3/test_chunking.py b/tests/models/minimax_h3/test_chunking.py
index f17f5c3db..a0342228c 100644
--- a/tests/models/minimax_h3/test_chunking.py
+++ b/tests/models/minimax_h3/test_chunking.py
@@ -35,6 +35,7 @@
_ChunkedRMSNorm,
install_h3_attention_norm_chunking,
install_h3_feed_forward_chunking,
+ install_h3_in_forward_block_checkpointing,
install_h3_lora_projection_chunking,
install_h3_rotary_chunking,
)
@@ -83,6 +84,9 @@ def __init__(self) -> None:
self.attn.processor = MiniMaxH3AttnProcessor()
self.ff = FeedForwardFake()
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return self.ff(hidden_states)
+
class TransformerFake(nn.Module):
def __init__(self) -> None:
@@ -137,6 +141,77 @@ def test_chunking_preserves_parameter_tree_and_is_idempotent() -> None:
assert not any("inner" in name for name, _ in transformer.named_modules())
+def test_in_forward_checkpointing_preserves_state_and_backward() -> None:
+ torch.manual_seed(13)
+ direct = TransformerFake().double()
+ checkpointed = deepcopy(direct)
+ state_keys = tuple(checkpointed.state_dict())
+ parameter_ids = {name: id(parameter) for name, parameter in checkpointed.named_parameters()}
+
+ assert install_h3_in_forward_block_checkpointing(checkpointed) == 5
+ assert install_h3_in_forward_block_checkpointing(checkpointed) == 5
+ assert tuple(checkpointed.state_dict()) == state_keys
+ assert {
+ name: id(parameter) for name, parameter in checkpointed.named_parameters()
+ } == parameter_ids
+
+ direct_input = torch.randn(2, 7, 5, dtype=torch.float64, requires_grad=True)
+ checkpointed_input = direct_input.detach().clone().requires_grad_(True)
+ direct_block = direct.transformer_blocks[0]
+ checkpointed_block = checkpointed.transformer_blocks[0]
+ direct_output = direct_block(direct_input)
+ checkpointed_output = checkpointed_block(checkpointed_input)
+ output_gradient = torch.randn_like(direct_output)
+ direct_output.backward(output_gradient)
+ checkpointed_output.backward(output_gradient)
+
+ torch.testing.assert_close(checkpointed_output, direct_output)
+ torch.testing.assert_close(checkpointed_input.grad, direct_input.grad)
+ for direct_parameter, checkpointed_parameter in zip(
+ direct_block.parameters(), checkpointed_block.parameters()
+ ):
+ torch.testing.assert_close(checkpointed_parameter.grad, direct_parameter.grad)
+
+
+def test_in_forward_checkpoint_starts_after_module_pre_forward_hook(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ transformer = TransformerFake().double()
+ install_h3_in_forward_block_checkpointing(transformer)
+ block = transformer.transformer_blocks[0]
+ checkpoint_dtypes = []
+ save_on_cpu_calls = []
+
+ class RecordingSaveOnCPU:
+ def __init__(self, *, pin_memory: bool, device_type: str) -> None:
+ save_on_cpu_calls.append((pin_memory, device_type))
+
+ def __enter__(self) -> None:
+ return None
+
+ def __exit__(self, *exc_info) -> None:
+ return None
+
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3._chunking.save_on_cpu",
+ RecordingSaveOnCPU,
+ )
+ monkeypatch.setattr(
+ "flow_factory.models.minimax_h3._chunking.checkpoint",
+ lambda function, owner, hidden_states, **kwargs: checkpoint_dtypes.append(
+ hidden_states.dtype
+ )
+ or function(owner, hidden_states),
+ )
+ handle = block.register_forward_pre_hook(lambda _module, inputs: (inputs[0].to(torch.float64),))
+
+ block(torch.randn(2, 7, 5, dtype=torch.float32, requires_grad=True))
+ handle.remove()
+
+ assert checkpoint_dtypes == [torch.float64]
+ assert save_on_cpu_calls == [(False, "cpu")]
+
+
def test_chunking_handles_remainder_and_preserves_forward_backward() -> None:
torch.manual_seed(17)
direct = TransformerFake().double()
diff --git a/tests/models/minimax_h3/test_workflow_adapters.py b/tests/models/minimax_h3/test_workflow_adapters.py
index 33bea489d..da9dac2cc 100644
--- a/tests/models/minimax_h3/test_workflow_adapters.py
+++ b/tests/models/minimax_h3/test_workflow_adapters.py
@@ -44,14 +44,20 @@ class UpstreamSchedulerFake:
def test_only_ref2va_requests_fsdp2_default_stream_unshard() -> None:
assert MiniMaxH3Ref2VAAdapter.fsdp2_use_default_stream_unshard
+ assert MiniMaxH3Ref2VAAdapter.fsdp2_use_in_forward_activation_checkpointing
+ assert MiniMaxH3Ref2VAAdapter.fsdp2_disable_backward_prefetch
assert MiniMaxH3Ref2VAAdapter.fsdp2_additional_wrap_module_names == (
"_ChunkedFeedForward",
"MiniMaxH3AdaLayerNormModulation",
"MiniMaxH3Attention",
)
assert not MiniMaxH3T2VAAdapter.fsdp2_use_default_stream_unshard
+ assert not MiniMaxH3T2VAAdapter.fsdp2_use_in_forward_activation_checkpointing
+ assert not MiniMaxH3T2VAAdapter.fsdp2_disable_backward_prefetch
assert not MiniMaxH3T2VAAdapter.fsdp2_additional_wrap_module_names
assert not MiniMaxH3FL2VAAdapter.fsdp2_use_default_stream_unshard
+ assert not MiniMaxH3FL2VAAdapter.fsdp2_use_in_forward_activation_checkpointing
+ assert not MiniMaxH3FL2VAAdapter.fsdp2_disable_backward_prefetch
assert not MiniMaxH3FL2VAAdapter.fsdp2_additional_wrap_module_names
From b6c8ffd6444bc099890ea7f7090f3e54a14bd4e9 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 01:43:44 +0800
Subject: [PATCH 60/76] [loading] fix: unshard repeated FSDP2 checkpoint
replays
---
.agents/knowledge/topics/fix_patterns.md | 19 ++++++++++
src/flow_factory/loading/backend.py | 46 +++++++++++++++++++++++-
tests/loading/test_backend_runtime.py | 35 ++++++++++++++++++
3 files changed, 99 insertions(+), 1 deletion(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index ef927458a..2c6c1a173 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -672,6 +672,25 @@ Based on the fix type, write the fix entry to the appropriate document:
parameter sharding may multiply activation lifetimes or overlap adjacent full-parameter units.
- **Related Constraint**: #9, #20
+### Nested FSDP2 checkpoint replay must verify that parameters remain unsharded
+- **Date**: 2026-08-31
+- **Symptom**: MiniMax H3 Ref2VA FSDP2 offline DPO completed both policy-arm forwards but
+ failed during the second checkpoint replay with `got mixed torch.Tensor and DTensor`.
+- **Root Cause**: PyTorch FSDP2 releases a nested unit after the first arm's post-backward while
+ its state remains `PRE_BACKWARD`; the second arm's checkpoint replay therefore takes the
+ pre-forward early return and uses sharded DTensor parameters with ordinary tensor inputs. This
+ is the upstream PyTorch issue #153354, fixed only after PyTorch 2.10 by commit 6579652.
+- **Fix**: Adapter-owned FSDP2 activation checkpointing now appends a post-prepare pre-forward
+ hook to every prepared FSDP unit. The hook calls the public synchronous `unshard()` API after
+ FSDP's own pre-forward hook, which is a no-op for normal forwards and restores parameters only
+ when an earlier checkpoint graph has already resharded them. A two-rank nested-FSDP reproducer
+ now completes two forward graphs and one combined backward without changing DPO semantics.
+- **Lesson**: A distributed training state is not proof of parameter residency. Multiple
+ checkpointed forwards may interleave one unit's post-backward with another graph's replay, so a
+ compatibility backport must repair residency at the FSDP lifecycle boundary rather than split
+ a coupled objective or expose the workaround in algorithm code.
+- **Related Constraint**: #9, #20
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/loading/backend.py b/src/flow_factory/loading/backend.py
index 56d3e3a71..388ade55d 100644
--- a/src/flow_factory/loading/backend.py
+++ b/src/flow_factory/loading/backend.py
@@ -32,6 +32,15 @@
logger = setup_logger(__name__)
+def _unshard_fsdp2_checkpoint_replay(module: nn.Module, _inputs: Any) -> None:
+ """Backport the nested-FSDP multi-forward replay fix through public APIs."""
+ unshard = getattr(module, "unshard", None)
+ if not callable(unshard):
+ module_type = type(module).__name__
+ raise TypeError(f"FSDP2 checkpoint replay hook requires unshard(), received {module_type}")
+ unshard()
+
+
def configure_backend_loading(accelerator: Accelerator, adapter_class: type) -> None:
"""Apply adapter capabilities before any pretrained component is loaded."""
if accelerator.distributed_type != DistributedType.FSDP:
@@ -173,7 +182,11 @@ def prepare(self, *objects: Any) -> Any:
disable_backward_prefetch = bool(
getattr(self.adapter, "fsdp2_disable_backward_prefetch", False)
)
- if not use_default_stream and not disable_backward_prefetch:
+ if (
+ not use_in_forward_checkpointing
+ and not use_default_stream
+ and not disable_backward_prefetch
+ ):
return prepared
prepared_objects = prepared if isinstance(prepared, (list, tuple)) else (prepared,)
@@ -186,6 +199,37 @@ def prepare(self, *objects: Any) -> Any:
if not prepared_modules:
raise TypeError("FSDP2 memory policy requested without a prepared module")
+ if use_in_forward_checkpointing:
+ configured = 0
+ newly_registered = 0
+ for root in prepared_modules:
+ for module in root.modules():
+ if not callable(getattr(module, "unshard", None)):
+ continue
+ configured += 1
+ if any(
+ hook is _unshard_fsdp2_checkpoint_replay
+ for hook in module._forward_pre_hooks.values()
+ ):
+ continue
+ # PyTorch <=2.10 returns early from its FSDP pre-forward hook
+ # during activation-checkpoint replay. With two forward graphs,
+ # a preceding post-backward may already have resharded this unit.
+ # Appending this hook after FSDP's hook mirrors the upstream fix
+ # and makes the public call a no-op when parameters remain full.
+ module.register_forward_pre_hook(_unshard_fsdp2_checkpoint_replay)
+ newly_registered += 1
+ if configured < 1:
+ raise TypeError(
+ "adapter-owned FSDP2 checkpointing requires prepared FSDPModule roots"
+ )
+ logger.info(
+ "Enabled nested FSDP2 multi-forward checkpoint replay unshard: "
+ "configured=%d, newly_registered=%d",
+ configured,
+ newly_registered,
+ )
+
if use_default_stream:
configured = []
for candidate in prepared_modules:
diff --git a/tests/loading/test_backend_runtime.py b/tests/loading/test_backend_runtime.py
index 586836d01..c367b45ef 100644
--- a/tests/loading/test_backend_runtime.py
+++ b/tests/loading/test_backend_runtime.py
@@ -114,6 +114,7 @@ def __init__(self) -> None:
super().__init__(2, 2)
self.unshard_async_op_calls = []
self.backward_prefetch_calls = []
+ self.replay_hook_events = []
def _set_unshard_async_op(self, enabled: bool) -> None:
self.unshard_async_op_calls.append(enabled)
@@ -121,6 +122,9 @@ def _set_unshard_async_op(self, enabled: bool) -> None:
def set_modules_to_backward_prefetch(self, modules) -> None:
self.backward_prefetch_calls.append(modules)
+ def unshard(self) -> None:
+ self.replay_hook_events.append("unshard")
+
def test_prepare_enables_adapter_requested_fsdp2_default_stream_unshard() -> None:
module = _FSDPModuleFake()
@@ -172,6 +176,13 @@ def test_prepare_extends_fsdp2_wrap_policy_before_distributed_preparation() -> N
def test_prepare_delegates_fsdp2_activation_checkpointing_during_prepare() -> None:
module = _FSDPModuleFake()
+ module.child = _FSDPModuleFake()
+ module.register_forward_pre_hook(
+ lambda _module, _inputs: module.replay_hook_events.append("fsdp-pre-forward")
+ )
+ module.child.register_forward_pre_hook(
+ lambda _module, _inputs: module.child.replay_hook_events.append("fsdp-pre-forward")
+ )
module._no_split_modules = ["TransformerBlock"]
accelerator = _AcceleratorFake(efficient=True)
plugin = accelerator.state.fsdp_plugin
@@ -205,6 +216,30 @@ def test_prepare_delegates_fsdp2_activation_checkpointing_during_prepare() -> No
(False, ("TransformerBlock", "ChunkedFeedForward")),
]
assert plugin.activation_checkpointing is True
+ module(torch.ones(1, 2))
+ assert module.replay_hook_events == ["fsdp-pre-forward", "unshard"]
+ module.child(torch.ones(1, 2))
+ assert module.child.replay_hook_events == ["fsdp-pre-forward", "unshard"]
+
+
+def test_prepare_replay_unshard_hook_is_idempotent() -> None:
+ module = _FSDPModuleFake()
+ accelerator = _AcceleratorFake(efficient=True)
+ accelerator.state.fsdp_plugin.activation_checkpointing = True
+ runtime = FSDPBackendLoadRuntime(
+ accelerator,
+ _plan(),
+ SimpleNamespace(
+ fsdp2_use_in_forward_activation_checkpointing=True,
+ configure_fsdp2_in_forward_activation_checkpointing=lambda _root: 1,
+ ),
+ )
+
+ assert runtime.prepare(module) is module
+ assert runtime.prepare(module) is module
+ module(torch.ones(1, 2))
+
+ assert module.replay_hook_events == ["unshard"]
@pytest.mark.parametrize(
From c84bfae051c5cc49197d238f659112ca6217fe46 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 03:14:29 +0800
Subject: [PATCH 61/76] [trainer] fix: reject unsupported optimizers before
model load
---
.agents/knowledge/topics/fix_patterns.md | 14 +++
src/flow_factory/trainers/loader.py | 4 +-
.../trainers/multirole/__init__.py | 2 +
.../trainers/multirole/backend.py | 65 +++++++----
.../test_distributed_plan_validation.py | 108 ++++++++++++------
tests/trainers/test_multirole_loader.py | 1 +
6 files changed, 133 insertions(+), 61 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 2c6c1a173..37b7202ff 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -691,6 +691,20 @@ Based on the fix type, write the fix entry to the appropriate document:
a coupled objective or expose the workaround in algorithm code.
- **Related Constraint**: #9, #20
+### Optimizer/backend compatibility must fail before model loading
+- **Date**: 2026-08-31
+- **Symptom**: A Muon run configured with DeepSpeed ZeRO-2 failed with the intended compatibility
+ error only after SD3.5 weights, LoRA adapters, and training data had been loaded.
+- **Root Cause**: Optimizer/backend validation lived only in `_init_optimizer`, whose lifecycle
+ position is necessarily after model adapter construction and data preprocessing.
+- **Fix**: The compatibility contract is now one shared backend-plan validator called by the
+ trainer loader before model construction and defensively called again before optimizer
+ construction. The two lifecycle gates therefore cannot drift to different rules.
+- **Lesson**: Validate compatibility from configuration as soon as the runtime backend is known.
+ Keep a second check at the resource-construction boundary, but delegate both checks to one
+ implementation so early rejection does not create a parallel source of truth.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/src/flow_factory/trainers/loader.py b/src/flow_factory/trainers/loader.py
index 4554f342a..2554c8c7a 100644
--- a/src/flow_factory/trainers/loader.py
+++ b/src/flow_factory/trainers/loader.py
@@ -31,7 +31,8 @@
from ..models.registry import get_model_adapter_class
from ..utils.env_utils import reconcile_config
from ..utils.logger_utils import setup_logger
-from .abc import BaseTrainer, validate_supported_distributed_plan
+from .abc import BaseTrainer
+from .multirole import validate_optimizer_backend_plan, validate_supported_distributed_plan
from .registry import get_trainer_class, list_registered_trainers
logger = setup_logger(__name__)
@@ -137,6 +138,7 @@ def load_trainer(config: Arguments) -> BaseTrainer:
# constructing an adapter under ZeRO-3 can shard parameters immediately, so
# rejecting it in BaseTrainer.__init__ is too late.
validate_supported_distributed_plan(accelerator)
+ validate_optimizer_backend_plan(accelerator, tuple(config.optimizer_args))
set_seed(config.training_args.seed, device_specific=True)
# Reconcile config with runtime distributed state (before any consumer reads it)
diff --git a/src/flow_factory/trainers/multirole/__init__.py b/src/flow_factory/trainers/multirole/__init__.py
index abeaf5475..1649ff235 100644
--- a/src/flow_factory/trainers/multirole/__init__.py
+++ b/src/flow_factory/trainers/multirole/__init__.py
@@ -3,6 +3,7 @@
from .backend import (
MultiRoleBackendValidationMixin,
configure_deepspeed_micro_batch_size,
+ validate_optimizer_backend_plan,
validate_supported_distributed_plan,
)
from .checkpointing import MULTIROLE_RUNTIME_CHILD_NAME, MultiRoleCheckpointingMixin
@@ -12,5 +13,6 @@
"MultiRoleCheckpointingMixin",
"MULTIROLE_RUNTIME_CHILD_NAME",
"configure_deepspeed_micro_batch_size",
+ "validate_optimizer_backend_plan",
"validate_supported_distributed_plan",
]
diff --git a/src/flow_factory/trainers/multirole/backend.py b/src/flow_factory/trainers/multirole/backend.py
index 323d7a49f..4a04813c6 100644
--- a/src/flow_factory/trainers/multirole/backend.py
+++ b/src/flow_factory/trainers/multirole/backend.py
@@ -25,6 +25,47 @@ def validate_supported_distributed_plan(accelerator: Accelerator) -> None:
)
+def validate_optimizer_backend_plan(
+ accelerator: Accelerator,
+ optimizer_args: Sequence[OptimizerArguments],
+) -> None:
+ """Reject optimizer/backend pairings before pretrained weights are loaded.
+
+ Args:
+ accelerator: Runtime backend whose optimizer support is being validated.
+ optimizer_args: Parsed optimizer configurations for every trainable role.
+
+ Returns:
+ None.
+
+ Raises:
+ ValueError: If Muon is paired with DeepSpeed or FSDP1.
+ """
+ if not uses_muon(optimizer_args):
+ return
+ if accelerator.distributed_type == DistributedType.DEEPSPEED:
+ raise ValueError(
+ "Muon with DeepSpeed is not verified in this framework: Muon rejects "
+ "non-matrix parameters, so it runs inside a CompositeOptimizer, and "
+ "DeepSpeed rebuilds its own optimizer wrapper around the object it "
+ "receives. Use DDP or FSDP2 with Muon, or select the adamw optimizer."
+ )
+ if accelerator.distributed_type != DistributedType.FSDP:
+ return
+ fsdp_plugin = getattr(accelerator.state, "fsdp_plugin", None)
+ fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1
+ if fsdp_version >= 2:
+ return
+ raise ValueError(
+ "Muon with FSDP1 does not work: FSDP1 flattens each wrapped unit into a "
+ "1D FlatParameter, so Muon is constructed over matrices and then receives "
+ "a 1D gradient, failing with 'Param gradient must be a 2D matrix' at the "
+ "first optimizer step. Set `fsdp_version: 2` in the accelerate config "
+ "(config/accelerate_configs/fsdp2.yaml), use DDP, or select the adamw "
+ "optimizer."
+ )
+
+
def configure_deepspeed_micro_batch_size(
accelerator: Accelerator, per_device_batch_size: int
) -> None:
@@ -61,29 +102,7 @@ def _validate_optimizer_backend(
optimizer_args: Sequence[OptimizerArguments],
) -> None:
"""Reject optimizer and distributed-backend pairings that are not verified."""
- if not uses_muon(optimizer_args):
- return
- if self.accelerator.distributed_type == DistributedType.DEEPSPEED:
- raise ValueError(
- "Muon with DeepSpeed is not verified in this framework: Muon rejects "
- "non-matrix parameters, so it runs inside a CompositeOptimizer, and "
- "DeepSpeed rebuilds its own optimizer wrapper around the object it "
- "receives. Use DDP or FSDP2 with Muon, or select the adamw optimizer."
- )
- if self.accelerator.distributed_type != DistributedType.FSDP:
- return
- fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None)
- fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1
- if fsdp_version >= 2:
- return
- raise ValueError(
- "Muon with FSDP1 does not work: FSDP1 flattens each wrapped unit into a "
- "1D FlatParameter, so Muon is constructed over matrices and then receives "
- "a 1D gradient, failing with 'Param gradient must be a 2D matrix' at the "
- "first optimizer step. Set `fsdp_version: 2` in the accelerate config "
- "(config/accelerate_configs/fsdp2.yaml), use DDP, or select the adamw "
- "optimizer."
- )
+ validate_optimizer_backend_plan(self.accelerator, optimizer_args)
def _validate_trainable_parameters_survived_prepare(self) -> None:
"""Reject a prepared root that no rank can train."""
diff --git a/tests/trainers/test_distributed_plan_validation.py b/tests/trainers/test_distributed_plan_validation.py
index 0766c89f3..0ca9d46a7 100644
--- a/tests/trainers/test_distributed_plan_validation.py
+++ b/tests/trainers/test_distributed_plan_validation.py
@@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
+import inspect
from pathlib import Path
from types import SimpleNamespace
@@ -19,7 +20,13 @@
import torch
from accelerate.utils import DistributedType
+from flow_factory.hparams.optimizer_args import (
+ AdamWOptimizerArguments,
+ MuonOptimizerArguments,
+)
+from flow_factory.trainers import loader
from flow_factory.trainers.abc import (
+ BaseTrainer,
configure_deepspeed_micro_batch_size,
validate_supported_distributed_plan,
)
@@ -43,8 +50,6 @@ def test_zero_three_is_rejected_before_any_weights_load() -> None:
def test_loader_rejects_zero_three_before_loading_model(monkeypatch: pytest.MonkeyPatch) -> None:
"""The trainer factory must reject ZeRO-3 before constructing an adapter."""
- from flow_factory.trainers import loader
-
accelerator = _accelerator(DistributedType.DEEPSPEED, zero_stage=3)
model_load_attempted = False
@@ -59,6 +64,7 @@ def unexpected_model_load(**kwargs: object) -> None:
config = SimpleNamespace(
mixed_precision="bf16",
+ optimizer_args=(),
model_args=SimpleNamespace(model_type="test"),
log_args=SimpleNamespace(save_dir="/tmp", run_name="zero3-rejection-test"),
training_args=SimpleNamespace(
@@ -81,6 +87,69 @@ def unexpected_model_load(**kwargs: object) -> None:
assert model_load_attempted is False
+def _fsdp_accelerator(fsdp_version: int) -> SimpleNamespace:
+ """Build an Accelerator stub reporting the requested FSDP major version."""
+ accelerator = _accelerator(DistributedType.FSDP)
+ accelerator.state.fsdp_plugin = SimpleNamespace(fsdp_version=fsdp_version)
+ return accelerator
+
+
+@pytest.mark.parametrize(
+ ("accelerator", "error_pattern", "run_name"),
+ [
+ (
+ _accelerator(DistributedType.DEEPSPEED, zero_stage=2),
+ "Muon with DeepSpeed is not verified",
+ "muon-deepspeed-rejection-test",
+ ),
+ (
+ _fsdp_accelerator(fsdp_version=1),
+ "Muon with FSDP1 does not work",
+ "muon-fsdp1-rejection-test",
+ ),
+ ],
+)
+def test_loader_rejects_unsupported_muon_backend_before_loading_model(
+ monkeypatch: pytest.MonkeyPatch,
+ accelerator: SimpleNamespace,
+ error_pattern: str,
+ run_name: str,
+) -> None:
+ """The trainer factory must reject unsupported Muon plans before model loading."""
+ model_load_attempted = False
+
+ class Adapter:
+ ddp_find_unused_parameters = False
+
+ def unexpected_model_load(**kwargs: object) -> None:
+ del kwargs
+ nonlocal model_load_attempted
+ model_load_attempted = True
+ raise AssertionError("load_model must not run for an unsupported Muon backend")
+
+ config = SimpleNamespace(
+ mixed_precision="bf16",
+ optimizer_args=(MuonOptimizerArguments(name="base"),),
+ model_args=SimpleNamespace(model_type="test"),
+ log_args=SimpleNamespace(save_dir="/tmp", run_name=run_name),
+ training_args=SimpleNamespace(
+ gradient_accumulation_steps=1,
+ max_grad_norm=1.0,
+ seed=42,
+ trainer_type="grpo",
+ required_trainable_roles=None,
+ ),
+ )
+ monkeypatch.setattr(loader, "get_model_adapter_class", lambda model_type: Adapter)
+ monkeypatch.setattr(loader, "Accelerator", lambda **kwargs: accelerator)
+ monkeypatch.setattr(loader, "load_model", unexpected_model_load)
+
+ with pytest.raises(ValueError, match=error_pattern):
+ loader.load_trainer(config)
+
+ assert model_load_attempted is False
+
+
@pytest.mark.parametrize("zero_stage", [1, 2])
def test_supported_deepspeed_stages_pass(zero_stage: int) -> None:
"""ZeRO-1 and ZeRO-2 are the supported DeepSpeed configurations."""
@@ -114,12 +183,6 @@ def test_deepspeed_micro_batch_size_is_set_for_custom_train_loader() -> None:
def test_muon_with_deepspeed_is_rejected_as_unverified() -> None:
"""Muon runs inside a composite; DeepSpeed rebuilds its own optimizer wrapper."""
- from flow_factory.hparams.optimizer_args import (
- AdamWOptimizerArguments,
- MuonOptimizerArguments,
- )
- from flow_factory.trainers.abc import BaseTrainer
-
trainer = SimpleNamespace(accelerator=_accelerator(DistributedType.DEEPSPEED, zero_stage=2))
with pytest.raises(ValueError, match="Muon with DeepSpeed is not verified"):
@@ -131,21 +194,8 @@ def test_muon_with_deepspeed_is_rejected_as_unverified() -> None:
BaseTrainer._validate_optimizer_backend(fsdp2_trainer, (MuonOptimizerArguments(name="base"),))
-def _fsdp_accelerator(fsdp_version: int) -> SimpleNamespace:
- """Accelerator reporting an FSDP plan of the requested major version."""
- accelerator = _accelerator(DistributedType.FSDP)
- accelerator.state.fsdp_plugin = SimpleNamespace(fsdp_version=fsdp_version)
- return accelerator
-
-
def test_muon_with_fsdp1_is_rejected_before_a_rollout_is_paid_for() -> None:
"""FSDP1 flattens to 1D, so Muon would only fail after the first full rollout."""
- from flow_factory.hparams.optimizer_args import (
- AdamWOptimizerArguments,
- MuonOptimizerArguments,
- )
- from flow_factory.trainers.abc import BaseTrainer
-
trainer = SimpleNamespace(accelerator=_fsdp_accelerator(fsdp_version=1))
with pytest.raises(ValueError, match="Muon with FSDP1 does not work"):
@@ -170,10 +220,6 @@ def test_deepspeed_gradient_clipping_is_wired_from_the_configured_norm() -> None
plugin, so leaving it unset ships an unresolved "auto" and max_grad_norm never
takes effect on that backend.
"""
- import inspect
-
- from flow_factory.trainers import loader
-
source = inspect.getsource(loader.load_trainer)
assert "ACCELERATE_GRADIENT_CLIPPING" in source
assert source.index("ACCELERATE_GRADIENT_CLIPPING") < source.index("accelerator = Accelerator(")
@@ -185,8 +231,6 @@ def _prepared_trainer(distributed_type: DistributedType, local: int, others: int
``reduce`` stands in for the collective: the guard asks whether ANY rank holds
trainable elements, so the stub adds what the peers would report.
"""
- from flow_factory.trainers.abc import BaseTrainer
-
accelerator = _accelerator(distributed_type)
accelerator.device = torch.device("cpu")
accelerator.num_processes = 2
@@ -220,8 +264,6 @@ def test_a_rank_holding_no_shard_of_the_adapter_is_accepted() -> None:
def test_tdm_r1_fsdp1_disables_incompatible_activation_checkpointing() -> None:
- from flow_factory.trainers.abc import BaseTrainer
-
plugin = SimpleNamespace(fsdp_version=1, activation_checkpointing=True)
disabled = []
trainer = SimpleNamespace(
@@ -250,8 +292,6 @@ def test_tdm_r1_fsdp1_disables_incompatible_activation_checkpointing() -> None:
def test_fsdp2_disables_model_checkpointing_and_keeps_backend_checkpointing(
checkpoint_policy: object,
) -> None:
- from flow_factory.trainers.abc import BaseTrainer
-
plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
disabled = []
trainer = SimpleNamespace(
@@ -274,8 +314,6 @@ def test_fsdp2_disables_model_checkpointing_and_keeps_backend_checkpointing(
def test_fsdp1_keeps_model_checkpointing_and_disables_nested_backend_checkpointing() -> None:
- from flow_factory.trainers.abc import BaseTrainer
-
plugin = SimpleNamespace(fsdp_version=1, activation_checkpointing=True)
trainer = SimpleNamespace(
accelerator=SimpleNamespace(
@@ -300,8 +338,6 @@ def test_fsdp1_keeps_model_checkpointing_and_disables_nested_backend_checkpointi
def test_fsdp2_rejects_selective_model_and_backend_checkpointing() -> None:
- from flow_factory.trainers.abc import BaseTrainer
-
plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
trainer = SimpleNamespace(
accelerator=SimpleNamespace(
@@ -327,8 +363,6 @@ def test_fsdp2_rejects_selective_model_and_backend_checkpointing() -> None:
def test_fsdp2_keeps_backend_checkpointing_when_model_policy_is_disabled() -> None:
- from flow_factory.trainers.abc import BaseTrainer
-
plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
trainer = SimpleNamespace(
accelerator=SimpleNamespace(
diff --git a/tests/trainers/test_multirole_loader.py b/tests/trainers/test_multirole_loader.py
index 9645726ee..693b9a1fd 100644
--- a/tests/trainers/test_multirole_loader.py
+++ b/tests/trainers/test_multirole_loader.py
@@ -52,6 +52,7 @@ def _config(
) -> SimpleNamespace:
return SimpleNamespace(
mixed_precision="no",
+ optimizer_args=(),
model_args=SimpleNamespace(model_type="tiny"),
training_args=SimpleNamespace(
trainer_type=trainer_type,
From 7501d1585fea30e58df07cb18edd202ac4a3cf2b Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 03:23:45 +0800
Subject: [PATCH 62/76] [samples] refactor: enforce reconstruction at gather
boundary
---
.agents/knowledge/topics/fix_patterns.md | 5 ++++-
src/flow_factory/rewards/reward_processor.py | 5 -----
src/flow_factory/utils/dist.py | 7 +++++--
tests/trainers/test_collective_packing.py | 18 +++++++++++++++++-
4 files changed, 26 insertions(+), 9 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 37b7202ff..1f946ecf2 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -58,7 +58,10 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Date**: 2026-08-30
- **Symptom**: Two-rank H3 Ref2VA GRPO failed before reward execution because `MiniMaxH3Ref2VASample` was reconstructed with `reference_manifest=None`.
- **Root Cause**: The distributed group-reward path gathered only reward-consumed fields, although `gather_samples` reconstructs the concrete sample class and that class can require additional state.
-- **Fix**: `BaseSample` now declares an empty `reconstruction_required_fields` contract, `OrderedReferenceConditionSample` adds `reference_manifest`, and `RewardProcessor` unions that contract into its distributed gather fields without forwarding it to the reward call.
+- **Fix**: `BaseSample` now declares an empty `reconstruction_required_fields` contract,
+ `OrderedReferenceConditionSample` adds `reference_manifest`, and `gather_samples` automatically
+ unions that contract at the concrete reconstruction boundary. Reward callers therefore transport
+ constructor state without forwarding it to the reward call or duplicating gather policy.
- **Lesson**: Communication payload requirements and reward-call requirements are distinct contracts; partial gathers must preserve constructor invariants even for fields that downstream computation does not consume.
- **Related Constraint**: N/A
diff --git a/src/flow_factory/rewards/reward_processor.py b/src/flow_factory/rewards/reward_processor.py
index 8632ff06c..ccd000867 100644
--- a/src/flow_factory/rewards/reward_processor.py
+++ b/src/flow_factory/rewards/reward_processor.py
@@ -474,11 +474,6 @@ def _compute_groupwise_distributed(
for model in models.values():
required_fields.update(model.required_fields)
- # ``gather_samples`` reconstructs the concrete sample class from the
- # transported fields. Preserve fields required by that class's constructor
- # invariants even when the reward itself does not consume them.
- required_fields.update(type(samples[0]).reconstruction_required_fields)
-
# Always include the typed source bookkeeping — the gate needs
# `source` (and ideally `source_id`) on the gathered side. Now
# that they're real dataclass fields on `BaseSample`, gathering
diff --git a/src/flow_factory/utils/dist.py b/src/flow_factory/utils/dist.py
index 330731502..33af7e60d 100644
--- a/src/flow_factory/utils/dist.py
+++ b/src/flow_factory/utils/dist.py
@@ -505,6 +505,7 @@ def gather_samples(
samples: Local samples on this rank.
field_names: Fields to gather. When ``'extra_kwargs'`` is included,
each key inside the dict is gathered independently and reassembled.
+ Concrete sample reconstruction fields are added automatically.
device: Target device for tensor fields in the returned samples.
Returns:
@@ -517,8 +518,10 @@ def gather_samples(
device = torch.device(device)
# Separate extra_kwargs from regular fields
- has_extra_kwargs = "extra_kwargs" in field_names
- regular_fields = sorted(f for f in field_names if f != "extra_kwargs")
+ reconstruction_fields = sample_cls.reconstruction_required_fields
+ gathered_fields = set(field_names) | set(reconstruction_fields)
+ has_extra_kwargs = "extra_kwargs" in gathered_fields
+ regular_fields = sorted(f for f in gathered_fields if f != "extra_kwargs")
extra_keys: List[str] = []
if has_extra_kwargs:
extra_keys = sorted({k for s in samples for k in s.extra_kwargs})
diff --git a/tests/trainers/test_collective_packing.py b/tests/trainers/test_collective_packing.py
index 665f0a401..75bd90bc8 100644
--- a/tests/trainers/test_collective_packing.py
+++ b/tests/trainers/test_collective_packing.py
@@ -6,7 +6,7 @@
import flow_factory.utils.dist as dist_utils
from flow_factory.advantage import AdvantageProcessor
-from flow_factory.samples import BaseSample
+from flow_factory.samples import BaseSample, MiniMaxH3Ref2VASample
from flow_factory.utils.dist import gather_aligned_floating_tensors, gather_samples
@@ -115,6 +115,22 @@ def test_gather_samples_packs_same_dtype_fields_and_preserves_other_fields():
torch.testing.assert_close(gathered[2].prompt_ids, samples[0].prompt_ids)
+def test_gather_samples_preserves_concrete_reconstruction_fields() -> None:
+ accelerator = GatherRecorder()
+ manifest = '[{"kind":"image","path":"condition.png"}]'
+ sample = MiniMaxH3Ref2VASample(
+ prompt="A reference-conditioned prompt",
+ reference_manifest=manifest,
+ )
+
+ gathered = gather_samples(accelerator, [sample], ["prompt"])
+
+ assert len(gathered) == 1
+ assert isinstance(gathered[0], MiniMaxH3Ref2VASample)
+ assert gathered[0].prompt == sample.prompt
+ assert gathered[0].reference_manifest == manifest
+
+
def test_gather_samples_keeps_large_cpu_fields_on_separate_paths(monkeypatch):
monkeypatch.setattr(dist_utils, "_CPU_PACKED_GATHER_MAX_BYTES", 1)
accelerator = GatherRecorder()
From 83ebee787725cd873d5bb29a82dbf384a67fbe6e Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 03:31:25 +0800
Subject: [PATCH 63/76] [trainer] refactor: resolve checkpoint owner before
model load
---
.agents/knowledge/topics/fix_patterns.md | 13 +-
guidance/new_model.md | 18 +--
src/flow_factory/trainers/abc.py | 57 +-------
src/flow_factory/trainers/loader.py | 7 +-
.../trainers/multirole/__init__.py | 2 +
.../trainers/multirole/backend.py | 79 +++++++++++
.../test_distributed_plan_validation.py | 124 +++++++++++++++++-
7 files changed, 228 insertions(+), 72 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 1f946ecf2..18119fcbb 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -405,12 +405,13 @@ Based on the fix type, write the fix entry to the appropriate document:
- **Root Cause**: Model-level checkpointing captured FP32 block inputs before FSDP2's forward-input
cast, while backward replay re-entered a block in `PRE_BACKWARD` state where PyTorch deliberately
skips that cast.
-- **Fix**: When full model checkpointing and FSDP2 activation checkpointing are both requested, the
- trainer now disables model-level boundaries and keeps Accelerate's backend checkpoint wrappers,
- which replay inside the fully-sharded mixed-precision boundary. Selective policies fail closed
- because backend checkpointing cannot preserve their exact selection, while FSDP1 retains its
- existing owner. Wan FSDP2 GRPO, TDM, SFT, and offline DPO plus SD3.5 and Bagel regressions verify
- the shared path.
+- **Fix**: Before loading the model, the trainer resolves one checkpoint owner. Any FSDP2 full
+ model policy is normalized to Accelerate's backend checkpoint wrappers, even when backend
+ checkpointing was initially disabled, because those wrappers replay inside the fully-sharded
+ mixed-precision boundary. Every selective FSDP2 model policy fails closed because its boundary
+ remains outside the input cast. FSDP1 retains its existing owner, and direct trainer construction
+ defensively reuses the same resolver after model realization. Wan FSDP2 GRPO, TDM, SFT, and
+ offline DPO plus SD3.5 and Bagel regressions verify the shared path.
- **Lesson**: Checkpoint placement is part of distributed precision semantics. A recompute boundary
outside a sharded module may not replay its forward hooks, so backend-aligned checkpoint wrappers
must own FSDP2 full checkpointing instead of nesting model-level boundaries around sharded blocks.
diff --git a/guidance/new_model.md b/guidance/new_model.md
index 8c86bd80c..b3b948edc 100644
--- a/guidance/new_model.md
+++ b/guidance/new_model.md
@@ -69,15 +69,15 @@ Diffusers model's `_repeated_blocks` declaration. Adapters with multiple forward
stacks should override `_gradient_checkpointing_units()` and return their blocks
in execution order.
-Checkpointing has one owner. When FSDP2 full model checkpointing and backend
-activation checkpointing are both enabled, the model policy yields ownership to
-the backend so recomputation stays inside the sharded mixed-precision boundary.
-FSDP1 keeps model-level ownership. FSDP2 rejects a selective train-level policy
-combined with backend activation checkpointing, because the backend cannot
-preserve the requested `fraction`, `every_n`, or `layers` boundary; disable
-backend activation checkpointing when using those policies. Transformers-style
-components support full checkpointing through `gradient_checkpointing_enable()`,
-but must expose the Diffusers callback API to support selective modes.
+Checkpointing has one owner. Before model loading, an FSDP2 full model policy is
+normalized to backend activation checkpointing, even when backend checkpointing
+was not explicitly enabled, so recomputation stays inside the sharded
+mixed-precision boundary. FSDP1 keeps model-level ownership. FSDP2 rejects every
+selective train-level policy because model-level `fraction`, `every_n`, or
+`layers` boundaries sit outside the FSDP2 input-cast boundary and cannot replay
+it safely. Transformers-style components support full checkpointing through
+`gradient_checkpointing_enable()`, but must expose the Diffusers callback API to
+support selective modes on compatible backends.
## Step-by-Step Implementation
diff --git a/src/flow_factory/trainers/abc.py b/src/flow_factory/trainers/abc.py
index 0ecb77551..2ceee1438 100644
--- a/src/flow_factory/trainers/abc.py
+++ b/src/flow_factory/trainers/abc.py
@@ -101,6 +101,7 @@
MULTIROLE_RUNTIME_CHILD_NAME,
MultiRoleBackendValidationMixin,
MultiRoleCheckpointingMixin,
+ configure_checkpointing_backend_plan,
configure_deepspeed_micro_batch_size,
validate_supported_distributed_plan,
)
@@ -1163,59 +1164,13 @@ def _validate_paradigm_dynamics(self) -> None:
)
def _apply_backend_checkpointing_constraints(self) -> None:
- """Select one checkpoint owner and reject the unsafe TDM-R1/FSDP1 case."""
- if self.accelerator.distributed_type != DistributedType.FSDP:
- return
- fsdp_plugin = getattr(self.accelerator.state, "fsdp_plugin", None)
- model_checkpointing = bool(
- getattr(
- self.training_args,
- "gradient_checkpointing_enabled",
- getattr(self.training_args, "enable_gradient_checkpointing", False),
- )
+ """Apply the shared owner plan to an adapter that may already be realized."""
+ disable_realized_model_checkpointing = configure_checkpointing_backend_plan(
+ self.accelerator,
+ self.training_args,
)
- fsdp_checkpointing = bool(getattr(fsdp_plugin, "activation_checkpointing", False))
- fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) or 1
- if self.training_args.trainer_type == "tdm-r1" and fsdp_version < 2:
- if not model_checkpointing and not fsdp_checkpointing:
- return
+ if disable_realized_model_checkpointing:
self.adapter.disable_gradient_checkpointing()
- self.training_args.enable_gradient_checkpointing = False
- if fsdp_plugin is not None:
- fsdp_plugin.activation_checkpointing = False
- logger.warning(
- "Disabled model and FSDP activation checkpointing for TDM-R1 on FSDP1: "
- "the surrogate objective runs reference/snapshot forwards between its live "
- "forward and backward, so FSDP1 recomputation saves a different graph. "
- "FSDP2 does not require this fallback."
- )
- return
-
- if model_checkpointing and fsdp_checkpointing:
- if fsdp_version >= 2:
- checkpoint_policy = self.training_args.enable_gradient_checkpointing
- full_checkpointing = checkpoint_policy is True or (
- getattr(checkpoint_policy, "mode", None) == "full"
- )
- if not full_checkpointing:
- raise ValueError(
- "FSDP2 activation checkpointing cannot preserve selective model "
- "checkpointing boundaries. Disable fsdp_activation_checkpointing or "
- "use train.enable_gradient_checkpointing=true/mode=full."
- )
- self.adapter.disable_gradient_checkpointing()
- self.training_args.enable_gradient_checkpointing = False
- logger.info(
- "Disabled model gradient checkpointing because FSDP2 activation "
- "checkpointing is enabled; checkpoint recomputation must stay inside "
- "the FSDP2 mixed-precision boundary."
- )
- else:
- fsdp_plugin.activation_checkpointing = False
- logger.info(
- "Disabled FSDP activation checkpointing because train-level model "
- "checkpointing is enabled; nested checkpoint boundaries duplicate recompute."
- )
def _initialization(self):
self._validate_paradigm_dynamics()
diff --git a/src/flow_factory/trainers/loader.py b/src/flow_factory/trainers/loader.py
index 2554c8c7a..673e5790e 100644
--- a/src/flow_factory/trainers/loader.py
+++ b/src/flow_factory/trainers/loader.py
@@ -32,7 +32,11 @@
from ..utils.env_utils import reconcile_config
from ..utils.logger_utils import setup_logger
from .abc import BaseTrainer
-from .multirole import validate_optimizer_backend_plan, validate_supported_distributed_plan
+from .multirole import (
+ configure_checkpointing_backend_plan,
+ validate_optimizer_backend_plan,
+ validate_supported_distributed_plan,
+)
from .registry import get_trainer_class, list_registered_trainers
logger = setup_logger(__name__)
@@ -139,6 +143,7 @@ def load_trainer(config: Arguments) -> BaseTrainer:
# rejecting it in BaseTrainer.__init__ is too late.
validate_supported_distributed_plan(accelerator)
validate_optimizer_backend_plan(accelerator, tuple(config.optimizer_args))
+ configure_checkpointing_backend_plan(accelerator, config.training_args)
set_seed(config.training_args.seed, device_specific=True)
# Reconcile config with runtime distributed state (before any consumer reads it)
diff --git a/src/flow_factory/trainers/multirole/__init__.py b/src/flow_factory/trainers/multirole/__init__.py
index 1649ff235..4cedfa522 100644
--- a/src/flow_factory/trainers/multirole/__init__.py
+++ b/src/flow_factory/trainers/multirole/__init__.py
@@ -2,6 +2,7 @@
from .backend import (
MultiRoleBackendValidationMixin,
+ configure_checkpointing_backend_plan,
configure_deepspeed_micro_batch_size,
validate_optimizer_backend_plan,
validate_supported_distributed_plan,
@@ -12,6 +13,7 @@
"MultiRoleBackendValidationMixin",
"MultiRoleCheckpointingMixin",
"MULTIROLE_RUNTIME_CHILD_NAME",
+ "configure_checkpointing_backend_plan",
"configure_deepspeed_micro_batch_size",
"validate_optimizer_backend_plan",
"validate_supported_distributed_plan",
diff --git a/src/flow_factory/trainers/multirole/backend.py b/src/flow_factory/trainers/multirole/backend.py
index 4a04813c6..8b87603d0 100644
--- a/src/flow_factory/trainers/multirole/backend.py
+++ b/src/flow_factory/trainers/multirole/backend.py
@@ -7,7 +7,11 @@
from accelerate.utils import DistributedType
from ...hparams.optimizer_args import OptimizerArguments
+from ...hparams.training_args import TrainingArguments
from ...optimizer import uses_muon
+from ...utils.logger_utils import setup_logger
+
+logger = setup_logger(__name__)
def validate_supported_distributed_plan(accelerator: Accelerator) -> None:
@@ -66,6 +70,81 @@ def validate_optimizer_backend_plan(
)
+def configure_checkpointing_backend_plan(
+ accelerator: Accelerator,
+ training_args: TrainingArguments,
+) -> bool:
+ """Select a checkpoint owner before model loading and distributed preparation.
+
+ Args:
+ accelerator: Runtime backend whose checkpointing policy is being configured.
+ training_args: Parsed algorithm arguments containing the model checkpoint policy.
+
+ Returns:
+ Whether a previously realized adapter must disable model checkpointing.
+
+ Raises:
+ ValueError: If FSDP2 is paired with a selective model checkpoint policy.
+ """
+ if accelerator.distributed_type != DistributedType.FSDP:
+ return False
+ fsdp_plugin = getattr(accelerator.state, "fsdp_plugin", None)
+ model_checkpointing = bool(
+ getattr(
+ training_args,
+ "gradient_checkpointing_enabled",
+ getattr(training_args, "enable_gradient_checkpointing", False),
+ )
+ )
+ fsdp_checkpointing = bool(getattr(fsdp_plugin, "activation_checkpointing", False))
+ fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) or 1
+
+ if training_args.trainer_type == "tdm-r1" and fsdp_version < 2:
+ if not model_checkpointing and not fsdp_checkpointing:
+ return False
+ training_args.enable_gradient_checkpointing = False
+ if fsdp_plugin is not None:
+ fsdp_plugin.activation_checkpointing = False
+ logger.warning(
+ "Disabled model and FSDP activation checkpointing for TDM-R1 on FSDP1: "
+ "the surrogate objective runs reference/snapshot forwards between its live "
+ "forward and backward, so FSDP1 recomputation saves a different graph. "
+ "FSDP2 does not require this fallback."
+ )
+ return model_checkpointing
+
+ if fsdp_version >= 2 and model_checkpointing:
+ checkpoint_policy = training_args.enable_gradient_checkpointing
+ full_checkpointing = checkpoint_policy is True or (
+ getattr(checkpoint_policy, "mode", None) == "full"
+ )
+ if not full_checkpointing:
+ raise ValueError(
+ "FSDP2 activation checkpointing cannot preserve selective model "
+ "checkpointing boundaries. Disable model checkpointing or use "
+ "train.enable_gradient_checkpointing=true/mode=full."
+ )
+ training_args.enable_gradient_checkpointing = False
+ if fsdp_plugin is None:
+ raise RuntimeError(
+ "FSDP2 full activation checkpointing requires an FSDP plugin, received None"
+ )
+ fsdp_plugin.activation_checkpointing = True
+ logger.info(
+ "Selected FSDP2 backend activation checkpointing and disabled train-level "
+ "model checkpointing so recomputation stays inside the mixed-precision boundary."
+ )
+ return True
+
+ if fsdp_version < 2 and model_checkpointing and fsdp_checkpointing:
+ fsdp_plugin.activation_checkpointing = False
+ logger.info(
+ "Disabled FSDP activation checkpointing because train-level model "
+ "checkpointing is enabled; nested checkpoint boundaries duplicate recompute."
+ )
+ return False
+
+
def configure_deepspeed_micro_batch_size(
accelerator: Accelerator, per_device_batch_size: int
) -> None:
diff --git a/tests/trainers/test_distributed_plan_validation.py b/tests/trainers/test_distributed_plan_validation.py
index 0ca9d46a7..29671761d 100644
--- a/tests/trainers/test_distributed_plan_validation.py
+++ b/tests/trainers/test_distributed_plan_validation.py
@@ -90,7 +90,11 @@ def unexpected_model_load(**kwargs: object) -> None:
def _fsdp_accelerator(fsdp_version: int) -> SimpleNamespace:
"""Build an Accelerator stub reporting the requested FSDP major version."""
accelerator = _accelerator(DistributedType.FSDP)
- accelerator.state.fsdp_plugin = SimpleNamespace(fsdp_version=fsdp_version)
+ accelerator.state.fsdp_plugin = SimpleNamespace(
+ fsdp_version=fsdp_version,
+ activation_checkpointing=False,
+ cpu_ram_efficient_loading=False,
+ )
return accelerator
@@ -150,6 +154,105 @@ def unexpected_model_load(**kwargs: object) -> None:
assert model_load_attempted is False
+@pytest.mark.parametrize("backend_checkpointing", [False, True])
+def test_loader_rejects_selective_fsdp2_checkpointing_before_loading_model(
+ monkeypatch: pytest.MonkeyPatch,
+ backend_checkpointing: bool,
+) -> None:
+ """The trainer factory must reject selective FSDP2 checkpointing before loading."""
+ accelerator = _fsdp_accelerator(fsdp_version=2)
+ accelerator.state.fsdp_plugin.activation_checkpointing = backend_checkpointing
+ model_load_attempted = False
+
+ class Adapter:
+ ddp_find_unused_parameters = False
+
+ def unexpected_model_load(**kwargs: object) -> None:
+ del kwargs
+ nonlocal model_load_attempted
+ model_load_attempted = True
+ raise AssertionError("load_model must not run for selective FSDP2 checkpointing")
+
+ config = SimpleNamespace(
+ mixed_precision="bf16",
+ optimizer_args=(),
+ model_args=SimpleNamespace(model_type="test"),
+ log_args=SimpleNamespace(save_dir="/tmp", run_name="checkpoint-plan-rejection-test"),
+ training_args=SimpleNamespace(
+ gradient_accumulation_steps=1,
+ max_grad_norm=1.0,
+ seed=42,
+ trainer_type="grpo",
+ required_trainable_roles=None,
+ enable_gradient_checkpointing=SimpleNamespace(mode="every_n"),
+ ),
+ )
+ monkeypatch.setattr(loader, "get_model_adapter_class", lambda model_type: Adapter)
+ monkeypatch.setattr(loader, "Accelerator", lambda **kwargs: accelerator)
+ monkeypatch.setattr(loader, "load_model", unexpected_model_load)
+
+ with pytest.raises(ValueError, match="cannot preserve selective"):
+ loader.load_trainer(config)
+
+ assert model_load_attempted is False
+
+
+@pytest.mark.parametrize("backend_checkpointing", [False, True])
+def test_loader_selects_fsdp2_backend_checkpointing_before_loading_model(
+ monkeypatch: pytest.MonkeyPatch,
+ backend_checkpointing: bool,
+) -> None:
+ """A full model policy must become the safe backend owner before loading."""
+ accelerator = _fsdp_accelerator(fsdp_version=2)
+ accelerator.state.fsdp_plugin.activation_checkpointing = backend_checkpointing
+ observed_plan = []
+ adapter = object()
+
+ class Adapter:
+ ddp_find_unused_parameters = False
+
+ class FakeTrainer:
+ def __init__(self, **kwargs: object) -> None:
+ self.kwargs = kwargs
+
+ config = SimpleNamespace(
+ mixed_precision="bf16",
+ optimizer_args=(),
+ model_args=SimpleNamespace(model_type="test"),
+ log_args=SimpleNamespace(save_dir="/tmp", run_name="checkpoint-plan-owner-test"),
+ training_args=SimpleNamespace(
+ gradient_accumulation_steps=1,
+ max_grad_norm=1.0,
+ seed=42,
+ trainer_type="grpo",
+ required_trainable_roles=None,
+ enable_gradient_checkpointing=True,
+ ),
+ )
+
+ def load_model_with_resolved_plan(**kwargs: object) -> object:
+ del kwargs
+ observed_plan.append(
+ (
+ config.training_args.enable_gradient_checkpointing,
+ accelerator.state.fsdp_plugin.activation_checkpointing,
+ )
+ )
+ return adapter
+
+ monkeypatch.setattr(loader, "get_trainer_class", lambda trainer_type: FakeTrainer)
+ monkeypatch.setattr(loader, "get_model_adapter_class", lambda model_type: Adapter)
+ monkeypatch.setattr(loader, "Accelerator", lambda **kwargs: accelerator)
+ monkeypatch.setattr(loader, "set_seed", lambda *args, **kwargs: None)
+ monkeypatch.setattr(loader, "reconcile_config", lambda *args, **kwargs: None)
+ monkeypatch.setattr(loader, "load_model", load_model_with_resolved_plan)
+
+ trainer = loader.load_trainer(config)
+
+ assert observed_plan == [(False, True)]
+ assert trainer.kwargs["adapter"] is adapter
+
+
@pytest.mark.parametrize("zero_stage", [1, 2])
def test_supported_deepspeed_stages_pass(zero_stage: int) -> None:
"""ZeRO-1 and ZeRO-2 are the supported DeepSpeed configurations."""
@@ -289,10 +392,15 @@ def test_tdm_r1_fsdp1_disables_incompatible_activation_checkpointing() -> None:
"checkpoint_policy",
[True, SimpleNamespace(mode="full")],
)
+@pytest.mark.parametrize("backend_checkpointing", [False, True])
def test_fsdp2_disables_model_checkpointing_and_keeps_backend_checkpointing(
checkpoint_policy: object,
+ backend_checkpointing: bool,
) -> None:
- plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
+ plugin = SimpleNamespace(
+ fsdp_version=2,
+ activation_checkpointing=backend_checkpointing,
+ )
disabled = []
trainer = SimpleNamespace(
accelerator=SimpleNamespace(
@@ -337,8 +445,14 @@ def test_fsdp1_keeps_model_checkpointing_and_disables_nested_backend_checkpointi
assert plugin.activation_checkpointing is False
-def test_fsdp2_rejects_selective_model_and_backend_checkpointing() -> None:
- plugin = SimpleNamespace(fsdp_version=2, activation_checkpointing=True)
+@pytest.mark.parametrize("backend_checkpointing", [False, True])
+def test_fsdp2_rejects_selective_model_checkpointing(
+ backend_checkpointing: bool,
+) -> None:
+ plugin = SimpleNamespace(
+ fsdp_version=2,
+ activation_checkpointing=backend_checkpointing,
+ )
trainer = SimpleNamespace(
accelerator=SimpleNamespace(
distributed_type=DistributedType.FSDP,
@@ -359,7 +473,7 @@ def test_fsdp2_rejects_selective_model_and_backend_checkpointing() -> None:
BaseTrainer._apply_backend_checkpointing_constraints(trainer)
assert trainer.training_args.enable_gradient_checkpointing.mode == "every_n"
- assert plugin.activation_checkpointing is True
+ assert plugin.activation_checkpointing is backend_checkpointing
def test_fsdp2_keeps_backend_checkpointing_when_model_policy_is_disabled() -> None:
From 5349f6a937c432c097a3c8c3b0ad3d9d64e73140 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 03:42:52 +0800
Subject: [PATCH 64/76] [optimizer] fix: reject unavailable Muon before model
load
---
.../knowledge/topics/component_variants.md | 7 +--
.agents/knowledge/topics/fix_patterns.md | 18 ++++++++
guidance/workflow.md | 6 ++-
src/flow_factory/optimizer/__init__.py | 8 +++-
src/flow_factory/optimizer/loader.py | 15 +++++++
.../trainers/multirole/backend.py | 31 +++++++------
tests/hparams/test_optimizer_args.py | 20 +++++++++
.../test_distributed_plan_validation.py | 45 ++++++++++++++++++-
8 files changed, 127 insertions(+), 23 deletions(-)
diff --git a/.agents/knowledge/topics/component_variants.md b/.agents/knowledge/topics/component_variants.md
index 37e98a298..a6c1d5735 100644
--- a/.agents/knowledge/topics/component_variants.md
+++ b/.agents/knowledge/topics/component_variants.md
@@ -153,9 +153,10 @@ its children's groups as one list. An all-AdamW run still gets a plain
`torch.optim.AdamW`, unchanged.
Muon therefore gives one variant **two** parameter groups, which is why
-`OptimizationRole.optimizer_group_ids` is a tuple. Muon combined with DeepSpeed is
-rejected at startup as unverified: DDP and FSDP only read `param_groups` and call
-`step`, but DeepSpeed rebuilds its own optimizer wrapper around the object it receives.
+`OptimizationRole.optimizer_group_ids` is a tuple. It requires a PyTorch build that
+exposes `torch.optim.Muon`. DeepSpeed is rejected as unverified because it rebuilds
+its own optimizer wrapper, while FSDP1 is rejected because its flat parameters erase
+the required matrix rank; the supported distributed plans are DDP and FSDP2.
## Optimizer and backend contract
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 18119fcbb..6e1bda21b 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -709,6 +709,24 @@ Based on the fix type, write the fix entry to the appropriate document:
implementation so early rejection does not create a parallel source of truth.
- **Related Constraint**: N/A
+### Optional optimizer APIs must be validated before model loading
+- **Date**: 2026-08-31
+- **Symptom**: A Muon configuration on the supported PyTorch 2.6 baseline could pass backend
+ validation, load pretrained weights, and then fail when optimizer construction accessed the
+ unavailable `torch.optim.Muon` attribute.
+- **Root Cause**: Backend validation assumed that parsing a Muon optimizer configuration implied
+ its optional PyTorch implementation existed. Flow-Factory's minimum PyTorch version predates
+ that API, so configuration support and runtime capability are independent contracts.
+- **Fix**: One optimizer capability validator now checks the concrete `torch.optim.Muon` API.
+ The shared pre-load execution-plan validator calls it after rejecting intrinsically
+ incompatible backends, so supported DDP/FSDP2 plans fail before model loading with an
+ actionable upgrade message. Direct optimizer construction and the defensive pre-optimizer
+ plan check reuse the same rule.
+- **Lesson**: Optional APIs gated by dependency versions belong in the same early execution-plan
+ validation as backend compatibility. Detect the capability itself instead of trusting a version
+ string, while keeping a late defensive call at the construction boundary.
+- **Related Constraint**: N/A
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/guidance/workflow.md b/guidance/workflow.md
index 491eb3d14..a52e24c3d 100644
--- a/guidance/workflow.md
+++ b/guidance/workflow.md
@@ -668,8 +668,10 @@ so a Muon variant is driven by two algorithms at once: Muon for its matrices and
AdamW for its biases, normalization scales and embeddings, which the `fallback_`
fields configure. `optimizer/loader.py` wraps that pair in a `CompositeOptimizer` so
the framework still prepares exactly one root. An all-AdamW run gets a plain
-`torch.optim.AdamW`, unchanged. Muon combined with DeepSpeed is refused at startup as
-unverified; use DDP or FSDP.
+`torch.optim.AdamW`, unchanged. Muon requires a PyTorch build that exposes
+`torch.optim.Muon` (2.10 or newer in supported environments). Muon combined with
+DeepSpeed is refused at startup as unverified, and FSDP1 flattens matrices into
+incompatible parameters; use DDP or FSDP2.
### Key Points
diff --git a/src/flow_factory/optimizer/__init__.py b/src/flow_factory/optimizer/__init__.py
index a0b70131e..f207f28a4 100644
--- a/src/flow_factory/optimizer/__init__.py
+++ b/src/flow_factory/optimizer/__init__.py
@@ -15,11 +15,17 @@
"""Build the single optimizer root from per-variant configurations."""
from .composite import CompositeOptimizer
-from .loader import build_optimizer, split_muon_parameters, uses_muon
+from .loader import (
+ build_optimizer,
+ split_muon_parameters,
+ uses_muon,
+ validate_muon_available,
+)
__all__ = [
"CompositeOptimizer",
"build_optimizer",
"split_muon_parameters",
"uses_muon",
+ "validate_muon_available",
]
diff --git a/src/flow_factory/optimizer/loader.py b/src/flow_factory/optimizer/loader.py
index aa56045ae..eb3232890 100644
--- a/src/flow_factory/optimizer/loader.py
+++ b/src/flow_factory/optimizer/loader.py
@@ -30,6 +30,20 @@
MUON_MINIMUM_DIMENSIONS = 2
+def validate_muon_available() -> None:
+ """Require the optional PyTorch Muon implementation.
+
+ Raises:
+ ValueError: If this PyTorch build does not expose ``torch.optim.Muon``.
+ """
+ if not hasattr(torch.optim, "Muon"):
+ raise ValueError(
+ f"torch.optim.Muon is unavailable in PyTorch {torch.__version__}. Install a "
+ "PyTorch build that provides Muon (2.10 or newer in supported environments), "
+ "or select the adamw optimizer."
+ )
+
+
def split_muon_parameters(
parameters: Sequence[torch.nn.Parameter],
) -> Tuple[List[torch.nn.Parameter], List[torch.nn.Parameter]]:
@@ -136,6 +150,7 @@ def build_optimizer(
if not muon_groups:
return torch.optim.AdamW(adamw_groups)
+ validate_muon_available()
children: List[torch.optim.Optimizer] = [torch.optim.Muon(muon_groups)]
if adamw_groups:
children.append(torch.optim.AdamW(adamw_groups))
diff --git a/src/flow_factory/trainers/multirole/backend.py b/src/flow_factory/trainers/multirole/backend.py
index 8b87603d0..e877d9956 100644
--- a/src/flow_factory/trainers/multirole/backend.py
+++ b/src/flow_factory/trainers/multirole/backend.py
@@ -8,7 +8,7 @@
from ...hparams.optimizer_args import OptimizerArguments
from ...hparams.training_args import TrainingArguments
-from ...optimizer import uses_muon
+from ...optimizer import uses_muon, validate_muon_available
from ...utils.logger_utils import setup_logger
logger = setup_logger(__name__)
@@ -43,7 +43,7 @@ def validate_optimizer_backend_plan(
None.
Raises:
- ValueError: If Muon is paired with DeepSpeed or FSDP1.
+ ValueError: If Muon is unavailable or paired with DeepSpeed or FSDP1.
"""
if not uses_muon(optimizer_args):
return
@@ -54,20 +54,19 @@ def validate_optimizer_backend_plan(
"DeepSpeed rebuilds its own optimizer wrapper around the object it "
"receives. Use DDP or FSDP2 with Muon, or select the adamw optimizer."
)
- if accelerator.distributed_type != DistributedType.FSDP:
- return
- fsdp_plugin = getattr(accelerator.state, "fsdp_plugin", None)
- fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1
- if fsdp_version >= 2:
- return
- raise ValueError(
- "Muon with FSDP1 does not work: FSDP1 flattens each wrapped unit into a "
- "1D FlatParameter, so Muon is constructed over matrices and then receives "
- "a 1D gradient, failing with 'Param gradient must be a 2D matrix' at the "
- "first optimizer step. Set `fsdp_version: 2` in the accelerate config "
- "(config/accelerate_configs/fsdp2.yaml), use DDP, or select the adamw "
- "optimizer."
- )
+ if accelerator.distributed_type == DistributedType.FSDP:
+ fsdp_plugin = getattr(accelerator.state, "fsdp_plugin", None)
+ fsdp_version = getattr(fsdp_plugin, "fsdp_version", 1) if fsdp_plugin else 1
+ if fsdp_version < 2:
+ raise ValueError(
+ "Muon with FSDP1 does not work: FSDP1 flattens each wrapped unit into a "
+ "1D FlatParameter, so Muon is constructed over matrices and then receives "
+ "a 1D gradient, failing with 'Param gradient must be a 2D matrix' at the "
+ "first optimizer step. Set `fsdp_version: 2` in the accelerate config "
+ "(config/accelerate_configs/fsdp2.yaml), use DDP, or select the adamw "
+ "optimizer."
+ )
+ validate_muon_available()
def configure_checkpointing_backend_plan(
diff --git a/tests/hparams/test_optimizer_args.py b/tests/hparams/test_optimizer_args.py
index 3e61c4941..32e775382 100644
--- a/tests/hparams/test_optimizer_args.py
+++ b/tests/hparams/test_optimizer_args.py
@@ -22,6 +22,11 @@
)
from flow_factory.optimizer import CompositeOptimizer, build_optimizer, split_muon_parameters
+requires_muon_api = pytest.mark.skipif(
+ not hasattr(torch.optim, "Muon"),
+ reason="functional Muon tests require torch.optim.Muon",
+)
+
def test_the_optimizer_key_selects_the_arguments_subclass() -> None:
"""AdamW and Muon hyperparameters live on separate classes, chosen by one key."""
@@ -66,6 +71,7 @@ def test_all_adamw_configurations_build_one_plain_adamw() -> None:
assert [group["role_name"] for group in optimizer.param_groups] == ["base"]
+@requires_muon_api
def test_muon_splits_a_variant_into_its_matrices_and_the_adamw_remainder() -> None:
"""torch.optim.Muon rejects non-matrix parameters, so the rest needs AdamW."""
matrix = torch.nn.Parameter(torch.randn(4, 3))
@@ -84,6 +90,7 @@ def test_muon_splits_a_variant_into_its_matrices_and_the_adamw_remainder() -> No
assert [group["role_name"] for group in optimizer.param_groups] == ["base", "base"]
+@requires_muon_api
def test_a_composite_steps_every_child_and_round_trips_its_state() -> None:
"""The framework prepares one optimizer, so the composite must behave like one."""
matrix = torch.nn.Parameter(torch.randn(4, 3))
@@ -111,6 +118,7 @@ def test_a_composite_steps_every_child_and_round_trips_its_state() -> None:
assert matrix.grad is None
+@requires_muon_api
def test_mixed_optimizers_across_variants_share_one_root() -> None:
"""One variant on Muon and another on AdamW still yields a single optimizer."""
generator = torch.nn.Parameter(torch.randn(4, 3))
@@ -137,6 +145,18 @@ def test_muon_requires_at_least_one_matrix_parameter() -> None:
)
+def test_muon_requires_the_optional_pytorch_api(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Direct optimizer construction must share the trainer's capability error."""
+ monkeypatch.delattr(torch.optim, "Muon", raising=False)
+
+ with pytest.raises(ValueError, match="torch.optim.Muon is unavailable"):
+ build_optimizer(
+ (MuonOptimizerArguments(name="base"),),
+ {"base": [torch.nn.Parameter(torch.randn(4, 3))]},
+ )
+
+
+@requires_muon_api
def test_a_composite_refuses_to_have_its_groups_reassigned() -> None:
"""The group list is a view; reassigning it would detach the children silently."""
optimizer = build_optimizer(
diff --git a/tests/trainers/test_distributed_plan_validation.py b/tests/trainers/test_distributed_plan_validation.py
index 29671761d..5e0e26fe0 100644
--- a/tests/trainers/test_distributed_plan_validation.py
+++ b/tests/trainers/test_distributed_plan_validation.py
@@ -154,6 +154,46 @@ def unexpected_model_load(**kwargs: object) -> None:
assert model_load_attempted is False
+def test_loader_rejects_unavailable_muon_before_loading_model(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A supported backend still requires the optimizer API before model loading."""
+ accelerator = _accelerator(DistributedType.MULTI_GPU)
+ model_load_attempted = False
+
+ class Adapter:
+ ddp_find_unused_parameters = False
+
+ def unexpected_model_load(**kwargs: object) -> None:
+ del kwargs
+ nonlocal model_load_attempted
+ model_load_attempted = True
+ raise AssertionError("load_model must not run without torch.optim.Muon")
+
+ config = SimpleNamespace(
+ mixed_precision="bf16",
+ optimizer_args=(MuonOptimizerArguments(name="base"),),
+ model_args=SimpleNamespace(model_type="test"),
+ log_args=SimpleNamespace(save_dir="/tmp", run_name="muon-api-rejection-test"),
+ training_args=SimpleNamespace(
+ gradient_accumulation_steps=1,
+ max_grad_norm=1.0,
+ seed=42,
+ trainer_type="grpo",
+ required_trainable_roles=None,
+ ),
+ )
+ monkeypatch.delattr(torch.optim, "Muon", raising=False)
+ monkeypatch.setattr(loader, "get_model_adapter_class", lambda model_type: Adapter)
+ monkeypatch.setattr(loader, "Accelerator", lambda **kwargs: accelerator)
+ monkeypatch.setattr(loader, "load_model", unexpected_model_load)
+
+ with pytest.raises(ValueError, match="torch.optim.Muon is unavailable"):
+ loader.load_trainer(config)
+
+ assert model_load_attempted is False
+
+
@pytest.mark.parametrize("backend_checkpointing", [False, True])
def test_loader_rejects_selective_fsdp2_checkpointing_before_loading_model(
monkeypatch: pytest.MonkeyPatch,
@@ -284,7 +324,9 @@ def test_deepspeed_micro_batch_size_is_set_for_custom_train_loader() -> None:
)
-def test_muon_with_deepspeed_is_rejected_as_unverified() -> None:
+def test_muon_with_deepspeed_is_rejected_as_unverified(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
"""Muon runs inside a composite; DeepSpeed rebuilds its own optimizer wrapper."""
trainer = SimpleNamespace(accelerator=_accelerator(DistributedType.DEEPSPEED, zero_stage=2))
@@ -294,6 +336,7 @@ def test_muon_with_deepspeed_is_rejected_as_unverified() -> None:
# AdamW is unaffected, and Muon is fine on the backends that preserve parameter rank.
BaseTrainer._validate_optimizer_backend(trainer, (AdamWOptimizerArguments(name="base"),))
fsdp2_trainer = SimpleNamespace(accelerator=_fsdp_accelerator(fsdp_version=2))
+ monkeypatch.setattr(torch.optim, "Muon", object(), raising=False)
BaseTrainer._validate_optimizer_backend(fsdp2_trainer, (MuonOptimizerArguments(name="base"),))
From 1846911e01bf46e745815de5e2ac1a7cd3091a84 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:27:55 +0800
Subject: [PATCH 65/76] [opd] fix: keep monitoring rewards eval-only
---
.agents/knowledge/topics/fix_patterns.md | 13 +++++++++
.../opd/lora/sd3_5/DiffusionOPD_aligned.yaml | 28 +------------------
.../opd/lora/sd3_5/geneval_pickscore_ocr.yaml | 27 +-----------------
.../sd3_5/geneval_pickscore_ocr_x0_norm.yaml | 27 +-----------------
guidance/algorithms.md | 4 +--
.../trainers/test_distillation_evaluation.py | 12 ++++++--
6 files changed, 28 insertions(+), 83 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 6e1bda21b..b5dd4c067 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -727,6 +727,19 @@ Based on the fix type, write the fix entry to the appropriate document:
string, while keeping a late defensive call at the construction boundary.
- **Related Constraint**: N/A
+### No-feedback algorithms must keep monitoring rewards eval-only
+- **Date**: 2026-08-31
+- **Symptom**: Every shipped DiffusionOPD example failed argument loading even though its reward
+ comments described monitoring rather than a training signal.
+- **Root Cause**: The examples duplicated evaluation rewards under top-level `rewards`, but the
+ algorithm declares a no-feedback execution contract that rejects every training reward.
+- **Fix**: The DiffusionOPD examples now keep monitoring models only under `eval_rewards`; the
+ algorithm guide states that evaluation scores never enter the distillation loss, and the shared
+ distillation contract tests cover DiffusionOPD alongside DMD2 and TDM.
+- **Lesson**: Monitoring intent does not change configuration semantics. A no-feedback algorithm
+ must express quality metrics through the evaluation-only reward surface.
+- **Related Constraint**: #7
+
## Cross-refs
- `constraints.md` (archival target for constraint violations)
diff --git a/examples/opd/lora/sd3_5/DiffusionOPD_aligned.yaml b/examples/opd/lora/sd3_5/DiffusionOPD_aligned.yaml
index beb1ccd91..de986bdd8 100644
--- a/examples/opd/lora/sd3_5/DiffusionOPD_aligned.yaml
+++ b/examples/opd/lora/sd3_5/DiffusionOPD_aligned.yaml
@@ -151,33 +151,7 @@ eval:
eval_freq: 30 # Run evaluation every N epochs (0 to disable) — upstream mopd eval_freq
seed: 42 # Eval random seed
-# Reward Model Configuration (monitoring only — NOT used by the distillation loss)
-# Routing via `applicable_datasets`: geneval -> geneval, ocr -> ocr,
-# pick_score -> all three datasets (covers the `pickscore` dataset and provides
-# a general-quality signal on the others).
-rewards:
- - name: "geneval" # Reward identifier (appears in log keys)
- reward_model: "GenEval" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: float32 # Precision for reward model inference
- applicable_datasets: ["geneval"] # Only fires on geneval source batches
-
- - name: "ocr" # Reward identifier (appears in log keys)
- reward_model: "OCR" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: bfloat16 # Precision for reward model inference
- applicable_datasets: ["ocr"] # Only fires on ocr source batches
-
- - name: "pick_score" # General image-quality reward (covers the pickscore dataset)
- reward_model: "PickScore" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: bfloat16 # Precision for reward model inference
- # applicable_datasets omitted -> applies to ALL training datasets (pickscore, ocr, geneval)
-
-# Eval Reward Models (same routing as training rewards)
+# Evaluation-only Reward Models (used by evaluate(); never by the distillation loss)
eval_rewards:
- name: "geneval" # Eval reward identifier
reward_model: "GenEval" # Reward model class (from registry)
diff --git a/examples/opd/lora/sd3_5/geneval_pickscore_ocr.yaml b/examples/opd/lora/sd3_5/geneval_pickscore_ocr.yaml
index 5e2e85393..66874f182 100644
--- a/examples/opd/lora/sd3_5/geneval_pickscore_ocr.yaml
+++ b/examples/opd/lora/sd3_5/geneval_pickscore_ocr.yaml
@@ -165,32 +165,7 @@ eval:
eval_freq: 10 # Run evaluation every N epochs (0 to disable)
seed: 42 # Eval random seed
-# Reward Model Configuration (monitoring only — NOT used by the distillation loss)
-# Routing via `applicable_datasets`: geneval -> geneval, ocr -> ocr,
-# pick_score -> all three sources (applicable_datasets omitted = applies to all).
-rewards:
- - name: "geneval" # Reward identifier (appears in log keys)
- reward_model: "GenEval" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: float32 # Precision for reward model inference
- applicable_datasets: ["geneval"] # Only fires on geneval source batches
-
- - name: "ocr" # Reward identifier (appears in log keys)
- reward_model: "OCR" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: bfloat16 # Precision for reward model inference
- applicable_datasets: ["ocr"] # Only fires on ocr source batches
-
- - name: "pick_score" # Reward identifier (appears in log keys)
- reward_model: "PickScore" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: bfloat16 # Precision for reward model inference
- # applicable_datasets omitted -> applies to ALL training sources (geneval, pickscore, ocr)
-
-# Eval Reward Models (same routing as training rewards)
+# Evaluation-only Reward Models (used by evaluate(); never by the distillation loss)
eval_rewards:
- name: "geneval" # Eval reward identifier
reward_model: "GenEval" # Reward model class (from registry)
diff --git a/examples/opd/lora/sd3_5/geneval_pickscore_ocr_x0_norm.yaml b/examples/opd/lora/sd3_5/geneval_pickscore_ocr_x0_norm.yaml
index 8830bed76..245841414 100644
--- a/examples/opd/lora/sd3_5/geneval_pickscore_ocr_x0_norm.yaml
+++ b/examples/opd/lora/sd3_5/geneval_pickscore_ocr_x0_norm.yaml
@@ -165,32 +165,7 @@ eval:
eval_freq: 10 # Run evaluation every N epochs (0 to disable)
seed: 42 # Eval random seed
-# Reward Model Configuration (monitoring only — NOT used by the distillation loss)
-# Routing via `applicable_datasets`: geneval -> geneval, ocr -> ocr,
-# pick_score -> all three sources (applicable_datasets omitted = applies to all).
-rewards:
- - name: "geneval" # Reward identifier (appears in log keys)
- reward_model: "GenEval" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: float32 # Precision for reward model inference
- applicable_datasets: ["geneval"] # Only fires on geneval source batches
-
- - name: "ocr" # Reward identifier (appears in log keys)
- reward_model: "OCR" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: bfloat16 # Precision for reward model inference
- applicable_datasets: ["ocr"] # Only fires on ocr source batches
-
- - name: "pick_score" # Reward identifier (appears in log keys)
- reward_model: "PickScore" # Reward model class (from registry)
- batch_size: 32 # Inference batch size for this reward model
- device: "cuda" # Device to run reward model on
- dtype: bfloat16 # Precision for reward model inference
- # applicable_datasets omitted -> applies to ALL training sources (geneval, pickscore, ocr)
-
-# Eval Reward Models (same routing as training rewards)
+# Evaluation-only Reward Models (used by evaluate(); never by the distillation loss)
eval_rewards:
- name: "geneval" # Eval reward identifier
reward_model: "GenEval" # Reward model class (from registry)
diff --git a/guidance/algorithms.md b/guidance/algorithms.md
index 67c2bca26..7f90b6124 100644
--- a/guidance/algorithms.md
+++ b/guidance/algorithms.md
@@ -762,7 +762,7 @@ The dynamics support matrix is:
| `v` | Yes | No | None |
| `x0` | Yes | No | None |
-`v` and `x0` fail fast under non-ODE dynamics because the target conversion assumes the ODE relation `mu = x_t + v * dt`. The `xt` target remains valid for Flow-SDE, Dance-SDE, and CPS; after optional self-normalization it is divided by the scheduler transition variance. No target uses the historical `0.5` multiplier. Rewards are used **only** for periodic eval monitoring (`evaluate()`), never in the distillation loss.
+`v` and `x0` fail fast under non-ODE dynamics because the target conversion assumes the ODE relation `mu = x_t + v * dt`. The `xt` target remains valid for Flow-SDE, Dance-SDE, and CPS; after optional self-normalization it is divided by the scheduler transition variance. No target uses the historical `0.5` multiplier. DiffusionOPD rejects training `rewards` because its execution contract has no feedback stage. Configure periodic monitoring only under `eval_rewards`; those scores are used by `evaluate()` and never enter the distillation loss.
### How it works (2-pass per epoch)
@@ -806,7 +806,7 @@ scheduler:
noise_level: 0.0
```
-Each teacher's `applicable_datasets` must reference declared `data.datasets[*].name` entries (validated at config load). The config schema allows several teachers to share a dataset for a future multi-teacher/ensemble trainer, but the current `DiffusionOPDTrainer` requires exactly one teacher per dataset and raises otherwise. See [`examples/opd/lora/sd3_5/`](../examples/opd/lora/sd3_5/) for two complete configs (`DiffusionOPD_aligned.yaml` to reproduce official results).
+Each teacher's `applicable_datasets` must reference declared `data.datasets[*].name` entries (validated at config load). The config schema allows several teachers to share a dataset for a future multi-teacher/ensemble trainer, but the current `DiffusionOPDTrainer` requires exactly one teacher per dataset and raises otherwise. See [`examples/opd/lora/sd3_5/`](../examples/opd/lora/sd3_5/) for three complete configs; `DiffusionOPD_aligned.yaml` reproduces the official setup.
## References
diff --git a/tests/trainers/test_distillation_evaluation.py b/tests/trainers/test_distillation_evaluation.py
index 98fd94141..728a04d90 100644
--- a/tests/trainers/test_distillation_evaluation.py
+++ b/tests/trainers/test_distillation_evaluation.py
@@ -114,11 +114,19 @@ def _distillation_config(trainer_type: str, **sections: Any) -> Arguments:
}
if trainer_type == "tdm-r1":
config["train"]["group_size"] = 2
+ if trainer_type == "diffusion-opd":
+ config["train"]["teachers"] = [
+ {
+ "name": "teacher",
+ "path": "teacher/lora",
+ "applicable_datasets": ["prompts"],
+ }
+ ]
config.update(sections)
return Arguments.from_dict(config)
-@pytest.mark.parametrize("trainer_type", ["dmd2", "tdm"])
+@pytest.mark.parametrize("trainer_type", ["dmd2", "tdm", "diffusion-opd"])
def test_reward_free_distillation_accepts_eval_rewards_only(trainer_type: str) -> None:
"""Quality monitoring is an evaluation concern; the loss stays reward-free."""
config = _distillation_config(
@@ -131,7 +139,7 @@ def test_reward_free_distillation_accepts_eval_rewards_only(trainer_type: str) -
assert config.eval_reward_args.get_by_name("quality").applicable_datasets == ["bench"]
-@pytest.mark.parametrize("trainer_type", ["dmd2", "tdm"])
+@pytest.mark.parametrize("trainer_type", ["dmd2", "tdm", "diffusion-opd"])
def test_reward_free_distillation_still_rejects_training_rewards(trainer_type: str) -> None:
"""An eval-only reward must not become a training signal by accident."""
with pytest.raises(ValueError, match="rewards"):
From cfecd731d3896b1696a77cda1916fbf199e5ced9 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:38:52 +0800
Subject: [PATCH 66/76] [docs] fix: align runtime docstrings with contracts
---
src/flow_factory/data_utils/dataset.py | 15 +++--
src/flow_factory/loading/backend.py | 13 +++-
src/flow_factory/loading/coordinator.py | 10 ++-
.../models/bagel/modeling/bagel/bagel.py | 66 +++++++++++++++++++
.../bagel/modeling/bagel/qwen2_navit.py | 27 ++++++++
.../models/minimax_h3/_chunking.py | 59 +++++++++++++++--
.../models/minimax_h3/adapters.py | 17 ++++-
.../models/minimax_h3/dependency.py | 2 +-
.../models/minimax_h3/workflow.py | 14 +++-
src/flow_factory/models/precision.py | 42 +++++++++++-
src/flow_factory/models/runtime/classic.py | 10 ++-
src/flow_factory/models/wan/_conditioning.py | 25 ++++++-
src/flow_factory/models/wan/wan2_i2v.py | 5 +-
src/flow_factory/optimizer/loader.py | 4 +-
src/flow_factory/samples/samples.py | 4 +-
src/flow_factory/trainers/loader.py | 25 +++----
.../trainers/multirole/backend.py | 6 +-
src/flow_factory/utils/dist.py | 8 +--
18 files changed, 306 insertions(+), 46 deletions(-)
diff --git a/src/flow_factory/data_utils/dataset.py b/src/flow_factory/data_utils/dataset.py
index cf7e09653..d24a1dd5a 100644
--- a/src/flow_factory/data_utils/dataset.py
+++ b/src/flow_factory/data_utils/dataset.py
@@ -551,12 +551,14 @@ def _preprocess_batch(
2. Load and prepare image inputs
3. Load and prepare video inputs
4. Load and prepare audio inputs
- 5. Call preprocess function
- 6. Move result tensors to CPU for caching
- 7. Pack non-preprocessed columns into ``metadata``
+ 5. Load ordered references and forward their canonical manifest sidecars
+ 6. Call preprocess function
+ 7. Move result tensors to CPU for caching
+ 8. Pack non-preprocessed columns into ``metadata``
Args:
batch: Dictionary with batch data.
+ indices: Source row indices aligned with batch rows, used in validation diagnostics.
image_dir: Directory containing images (``None`` skips image loading).
Per-sample paths are loaded as PIL Images and kept as a
``List[Image]``; the column-level ``images`` field is therefore
@@ -697,6 +699,7 @@ def _preprocess_batch(
audio_args["audios"].append(audios)
batch["audios"].append(audios)
+ # 5. Load ordered references and retain their canonical reconstruction sidecars.
reference_args: Dict[str, Any] = {}
if self._uses_ordered_references:
raw_references = batch.pop("references")
@@ -732,7 +735,7 @@ def _preprocess_batch(
if column in batch
}
- # 5. Call preprocess function with filtered kwargs
+ # 6. Call preprocess function with filtered kwargs
input_args = {
**prompt_args,
**image_args,
@@ -751,7 +754,7 @@ def _preprocess_batch(
f"{sorted(passthrough_collisions)!r}"
)
- # 6. Process results - move tensors to CPU for caching.
+ # 7. Process results - move tensors to CPU for caching.
# Image-valued adapter outputs (declared via `python_format_columns`)
# are stored as per-sample List[PIL] so HF serializes them via the Image
# feature; ragged image tensors (variable size/count, e.g. multi-ref I2I)
@@ -783,7 +786,7 @@ def _preprocess_batch(
# Case C: Other types (None, int, etc)
final_res[k] = v
- # 7. Prepare final results
+ # 8. Prepare final results
batch_dict = {**batch, **final_res}
if self._uses_ordered_references:
_validate_arrow_safe_ordered_result(batch_dict, len(batch["prompt"]))
diff --git a/src/flow_factory/loading/backend.py b/src/flow_factory/loading/backend.py
index 388ade55d..23ccd88fa 100644
--- a/src/flow_factory/loading/backend.py
+++ b/src/flow_factory/loading/backend.py
@@ -130,7 +130,18 @@ def load_scope(self, role: ComponentRole) -> Iterator[None]:
os.environ["FSDP_CPU_RAM_EFFICIENT_LOADING"] = previous
def prepare(self, *objects: Any) -> Any:
- """Prepare roots and apply adapter-requested FSDP2 communication policy."""
+ """Apply adapter-owned FSDP2 wrap, checkpoint, replay, and communication policies.
+
+ Args:
+ *objects: Model roots and related objects forwarded to ``Accelerator.prepare``.
+
+ Returns:
+ The prepared object or tuple returned by ``Accelerator.prepare``.
+
+ Raises:
+ TypeError: If an adapter-requested FSDP2 policy cannot be installed on the supplied
+ or prepared roots.
+ """
plugin = self.accelerator.state.fsdp_plugin
use_in_forward_checkpointing = (
(getattr(plugin, "fsdp_version", 1) or 1) >= 2
diff --git a/src/flow_factory/loading/coordinator.py b/src/flow_factory/loading/coordinator.py
index bfec60399..9da7b2e5c 100644
--- a/src/flow_factory/loading/coordinator.py
+++ b/src/flow_factory/loading/coordinator.py
@@ -96,7 +96,15 @@ def load_components(
*,
device: Any,
) -> None:
- """Load replicas and target-owned auxiliary remainders without moving targets."""
+ """Load replicas and target-owned auxiliary remainders without moving targets.
+
+ Args:
+ components: Logical components requested by the adapter lifecycle call.
+ device: Destination device forwarded to component materialization.
+
+ Returns:
+ None. Only replicas observed as materialized are finalized through the backend.
+ """
requested = self.adapter._resolve_component_names(components)
replicated = [
name
diff --git a/src/flow_factory/models/bagel/modeling/bagel/bagel.py b/src/flow_factory/models/bagel/modeling/bagel/bagel.py
index 06286fa99..dcffb21fb 100644
--- a/src/flow_factory/models/bagel/modeling/bagel/bagel.py
+++ b/src/flow_factory/models/bagel/modeling/bagel/bagel.py
@@ -301,6 +301,24 @@ def forward_cache_update_text(
key_values_lens: torch.IntTensor,
language_model_forward: Optional[Callable[..., Any]] = None,
):
+ """Append packed text tokens to the language-model cache.
+
+ Args:
+ past_key_values: Existing packed key/value cache.
+ packed_text_ids: Raw text token IDs to embed inside the routed model forward.
+ packed_text_position_ids: Packed position IDs for the appended text.
+ text_token_lens: Per-sample text query lengths.
+ packed_text_indexes: Packed query destination indices.
+ packed_key_value_indexes: Packed indices of the existing cache entries.
+ key_values_lens: Per-sample existing cache lengths.
+ language_model_forward: Optional outer language-model callable. Flow-Factory injects
+ the prepared/routed transformer so embedding, decoder, and final normalization
+ stay under the sharded root's forward hooks; standalone Bagel defaults to
+ ``self.language_model``.
+
+ Returns:
+ The updated packed key/value cache.
+ """
# Flow-Factory injects its prepared component route here. Standalone Bagel
# keeps the original physical language-model call as the default.
if language_model_forward is None:
@@ -419,6 +437,29 @@ def forward_cache_update_vit(
key_values_lens: torch.IntTensor,
language_model_forward: Optional[Callable[..., Any]] = None,
):
+ """Encode packed vision tokens and append them to the language-model cache.
+
+ Args:
+ past_key_values: Existing packed key/value cache.
+ packed_text_ids: Raw boundary-token IDs inserted into the packed query.
+ packed_text_indexes: Destination indices for the boundary-token embeddings.
+ packed_vit_tokens: Flattened image inputs consumed by the vision encoder.
+ packed_vit_token_indexes: Destination indices for encoded vision tokens.
+ packed_vit_position_ids: Packed vision position IDs.
+ vit_token_seqlens: Per-image vision token lengths.
+ packed_position_ids: Packed language-model query position IDs.
+ packed_seqlens: Per-sample query lengths.
+ packed_indexes: Packed query indices.
+ packed_key_value_indexes: Packed indices of the existing cache entries.
+ key_values_lens: Per-sample existing cache lengths.
+ language_model_forward: Optional outer language-model callable. Flow-Factory injects
+ the prepared/routed transformer so embedding, decoder, and final normalization
+ stay under the sharded root's forward hooks; standalone Bagel defaults to
+ ``self.language_model``.
+
+ Returns:
+ The updated packed key/value cache.
+ """
if language_model_forward is None:
language_model_forward = self.language_model
@@ -570,6 +611,31 @@ def forward_cache_update_vae(
packed_key_value_indexes: torch.Tensor,
language_model_forward: Optional[Callable[..., Any]] = None,
):
+ """Encode packed image latents and append them to the language-model cache.
+
+ Args:
+ vae_model: VAE used to encode the padded image batch.
+ past_key_values: Existing packed key/value cache.
+ padded_images: Padded image tensors for the active reference round.
+ patchified_vae_latent_shapes: Per-image patch-grid heights and widths.
+ packed_vae_position_ids: Packed latent position IDs.
+ packed_timesteps: Timesteps used by the latent time embedder.
+ packed_vae_token_indexes: Destination indices for encoded latent tokens.
+ packed_text_ids: Raw boundary-token IDs inserted into the packed query.
+ packed_text_indexes: Destination indices for the boundary-token embeddings.
+ packed_position_ids: Packed language-model query position IDs.
+ packed_seqlens: Per-sample query lengths.
+ packed_indexes: Packed query indices.
+ key_values_lens: Per-sample existing cache lengths.
+ packed_key_value_indexes: Packed indices of the existing cache entries.
+ language_model_forward: Optional outer language-model callable. Flow-Factory injects
+ the prepared/routed transformer so embedding, decoder, and final normalization
+ stay under the sharded root's forward hooks; standalone Bagel defaults to
+ ``self.language_model``.
+
+ Returns:
+ The updated packed key/value cache.
+ """
if language_model_forward is None:
language_model_forward = self.language_model
diff --git a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
index 9d0ad9328..186966970 100644
--- a/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
+++ b/src/flow_factory/models/bagel/modeling/bagel/qwen2_navit.py
@@ -1326,6 +1326,33 @@ def forward_inference(
packed_text_indexes=None,
packed_text_ids: Optional[torch.LongTensor] = None,
) -> BaseNavitOutputWithPast:
+ """Run one packed language-model forward from embeddings, token IDs, or both.
+
+ Args:
+ packed_query_sequence: Existing packed auxiliary sequence, or None for a text-only
+ cache update.
+ query_lens: Per-sample query lengths.
+ packed_query_position_ids: Packed query position IDs.
+ packed_query_indexes: Packed query indices.
+ past_key_values: Optional existing packed key/value cache.
+ key_values_lens: Per-sample existing cache lengths.
+ packed_key_value_indexes: Packed indices of existing cache entries.
+ update_past_key_values: Whether to append the query to the cache.
+ is_causal: Whether the query uses causal attention.
+ mode: Bagel language-model routing mode.
+ packed_vae_token_indexes: Optional positions of generative latent tokens.
+ packed_text_indexes: Destination indices when inserting raw text embeddings into an
+ existing packed sequence.
+ packed_text_ids: Raw token IDs to embed inside this prepared-root forward. Required
+ when ``packed_query_sequence`` is None.
+
+ Returns:
+ Packed model output with the updated cache when requested.
+
+ Raises:
+ ValueError: If neither input source is supplied, destination indices are missing, or
+ token and destination counts differ.
+ """
# Keep token embedding inside the outer language-model forward. Distributed
# wrappers attach their unshard hooks to this boundary, so callers must not
diff --git a/src/flow_factory/models/minimax_h3/_chunking.py b/src/flow_factory/models/minimax_h3/_chunking.py
index 3f346bd5c..efb16355c 100644
--- a/src/flow_factory/models/minimax_h3/_chunking.py
+++ b/src/flow_factory/models/minimax_h3/_chunking.py
@@ -325,8 +325,16 @@ def install_h3_feed_forward_chunking(
distributed wrapping. The mutation is idempotent and preserves every parameter
object.
+ Args:
+ transformer: Materialized H3 transformer containing both repeated block stacks.
+ max_tokens: Positive maximum token count processed by one feed-forward invocation.
+
Returns:
Number of feed-forward layers configured across both stacks.
+
+ Raises:
+ TypeError: If the materialized upstream feed-forward structure or hooks are unsupported.
+ ValueError: If ``max_tokens`` is invalid or conflicts with an earlier installation.
"""
max_tokens = _positive_int(max_tokens, "max_tokens")
blocks = tuple(_h3_repeated_blocks(transformer))
@@ -382,8 +390,18 @@ def install_h3_in_forward_block_checkpointing(transformer: nn.Module) -> int:
inputs before FSDP2 casts them for mixed-precision compute. Accelerate's generic
FSDP2 policy checkpoints each direct child separately. This instance-local
forward delegates the complete original block body to one non-reentrant
- checkpoint after the block's FSDP pre-forward hook has already run. Its saved
- BF16 block inputs live in pinned CPU memory until their backward replay.
+ checkpoint after the block's FSDP pre-forward hook has already run. Saved tensor inputs are
+ offloaded to CPU until backward replay, using pinned host memory for CUDA inputs. In the
+ intended BF16 FSDP2 path these are the post-cast BF16 block inputs.
+
+ Args:
+ transformer: Materialized H3 transformer containing both repeated block stacks.
+
+ Returns:
+ Number of repeated H3 blocks configured across both stacks.
+
+ Raises:
+ TypeError: If a block has an unsupported or conflicting forward installation.
"""
configured = 0
for block_name, block in _h3_repeated_blocks(transformer):
@@ -423,7 +441,19 @@ def install_h3_attention_norm_chunking(
*,
max_tokens: int = H3_MAX_ATTENTION_NORM_TOKENS,
) -> int:
- """Install bounded Q/K head normalization on both H3 repeated stacks."""
+ """Install bounded Q/K head normalization on both H3 repeated stacks.
+
+ Args:
+ transformer: Materialized H3 transformer containing both repeated block stacks.
+ max_tokens: Positive maximum token count normalized by one local operation.
+
+ Returns:
+ Number of Q/K RMSNorm modules configured across both stacks.
+
+ Raises:
+ TypeError: If the materialized upstream attention or RMSNorm structure is unsupported.
+ ValueError: If ``max_tokens`` is invalid or conflicts with an earlier installation.
+ """
max_tokens = _positive_int(max_tokens, "max_tokens")
configured = 0
for block_name, block in _h3_repeated_blocks(transformer):
@@ -478,8 +508,16 @@ def install_h3_lora_projection_chunking(
parameters, children, hooks, and state-dict paths remain owned by that same
module object.
+ Args:
+ transformer: Materialized H3 transformer after PEFT injection.
+ max_tokens: Positive maximum token count projected by one adapted invocation.
+
Returns:
Number of adapted Q/K/V/output projections configured across both stacks.
+
+ Raises:
+ TypeError: If the materialized attention or PEFT projection structure is unsupported.
+ ValueError: If ``max_tokens`` is invalid or conflicts with an earlier installation.
"""
max_tokens = _positive_int(max_tokens, "max_tokens")
configured = 0
@@ -563,7 +601,20 @@ def install_h3_rotary_chunking(
*,
max_tokens: int = H3_MAX_ROTARY_TOKENS,
) -> int:
- """Install the bounded processor on every standard H3 attention instance."""
+ """Install the bounded processor on every standard H3 attention instance.
+
+ Args:
+ transformer: Materialized H3 transformer containing both repeated block stacks.
+ max_tokens: Positive maximum token count rotated by one local operation.
+
+ Returns:
+ Number of H3 attention processors configured across both stacks.
+
+ Raises:
+ ImportError: If the pinned MiniMax H3 attention processor is unavailable.
+ TypeError: If the materialized attention processor structure is unsupported.
+ ValueError: If ``max_tokens`` is invalid or conflicts with an earlier installation.
+ """
max_tokens = _positive_int(max_tokens, "max_tokens")
symbols = require_minimax_h3_support()
processor_class = symbols.MiniMaxH3AttnProcessor
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index db5a6a633..719d1c927 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -425,7 +425,22 @@ def apply_lora(
components: Union[str, List[str]] = "transformer",
overwrite: bool = False,
) -> Any:
- """Apply PEFT and bound Ref2VA projection memory under FSDP2."""
+ """Apply PEFT and bound Ref2VA projection memory under FSDP2.
+
+ Args:
+ target_modules: PEFT target-module patterns forwarded to ``BaseAdapter.apply_lora``.
+ components: Canonical components that receive LoRA.
+ overwrite: Whether to replace an existing default adapter.
+
+ Returns:
+ The PEFT model, per-component PEFT mapping, or empty mapping returned by the base
+ method.
+
+ Raises:
+ TypeError: If Ref2VA FSDP2 projection chunking receives a nonstandard PEFT component,
+ a non-``nn.Module`` base model, or incompatible projection structure.
+ ValueError: If projection chunking conflicts with an earlier installation.
+ """
component_names = (components,) if isinstance(components, str) else tuple(components)
result = super().apply_lora(
target_modules=target_modules,
diff --git a/src/flow_factory/models/minimax_h3/dependency.py b/src/flow_factory/models/minimax_h3/dependency.py
index 936685ee3..f7675f695 100644
--- a/src/flow_factory/models/minimax_h3/dependency.py
+++ b/src/flow_factory/models/minimax_h3/dependency.py
@@ -55,7 +55,7 @@
@dataclass(frozen=True)
class MiniMaxH3Symbols:
- """Hold all upstream classes used by the shared H3 core."""
+ """Hold all upstream symbols used by the shared H3 core."""
ModularPipeline: Type[Any]
MiniMaxH3ModularPipeline: Type[Any]
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index d0a9b0ea6..b4fc9dd50 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -96,7 +96,19 @@ def load_h3_workflow_pipeline(
def build_h3_component_runtime(adapter: Any) -> ModularPipelineRuntime:
- """Wrap one pruned pipeline and materialize only its training transformer."""
+ """Build the pruned runtime and prepare the H3 transformer for bounded execution.
+
+ Args:
+ adapter: H3 adapter that declares the target transformer and loads the modular pipeline.
+
+ Returns:
+ Runtime with the training transformer materialized and its feed-forward and attention
+ normalization operations configured for bounded token chunks.
+
+ Raises:
+ ValueError: If the adapter targets components outside the H3 training contract.
+ TypeError: If the materialized transformer structure cannot accept the required chunking.
+ """
validate_h3_target_components(adapter)
runtime = ModularPipelineRuntime.from_adapter(adapter, adapter.load_pipeline())
runtime.materialize_components([adapter.transformer_component_name])
diff --git a/src/flow_factory/models/precision.py b/src/flow_factory/models/precision.py
index 4537b7fe8..a6f4fbcd9 100644
--- a/src/flow_factory/models/precision.py
+++ b/src/flow_factory/models/precision.py
@@ -125,7 +125,24 @@ def component_dtype_mapping(
text_encoder_names: Sequence[str],
manifest_declared_names: Sequence[str] | None = None,
) -> dict[str, torch.dtype]:
- """Resolve policies for concrete components, with wider adapter-manifest declarations."""
+ """Resolve user and adapter dtype policies for concrete components.
+
+ Args:
+ user_policy: User-selected dtype policy, which takes precedence.
+ manifest_policy: Adapter-declared fallback dtype policy.
+ component_names: Concrete component names to resolve into the returned mapping.
+ transformer_names: Concrete names matched by the ``transformers`` group selector.
+ text_encoder_names: Concrete names matched by the ``text_encoders`` group selector.
+ manifest_declared_names: Optional superset used only to validate adapter-manifest
+ selectors, including class-declared optional components. Resolution still iterates
+ ``component_names``. Defaults to ``component_names``.
+
+ Returns:
+ Concrete component names mapped to their resolved non-null dtypes.
+
+ Raises:
+ ValueError: If either policy contains a selector outside its declared namespace.
+ """
validate_dtype_policy_selectors(user_policy, declared_names=component_names)
validate_dtype_policy_selectors(
manifest_policy,
@@ -160,7 +177,28 @@ def build_component_load_dtype_kwargs(
requested_names: Sequence[str] | None = None,
preserve_unselected: bool = False,
) -> Dict[str, object]:
- """Build the one native-loader dtype argument for eager or selective loading."""
+ """Build the native-loader dtype argument for eager or selective loading.
+
+ Args:
+ user_policy: User-selected dtype policy, which takes precedence.
+ manifest_policy: Adapter-declared fallback dtype policy.
+ component_names: Concrete component names available to the loader.
+ transformer_names: Concrete names matched by the ``transformers`` group selector.
+ text_encoder_names: Concrete names matched by the ``text_encoders`` group selector.
+ manifest_declared_names: Optional superset used only to validate adapter-manifest
+ selectors, including class-declared optional components. Resolution still iterates
+ ``component_names``. Defaults to ``component_names``.
+ requested_names: Optional subset being loaded by this request.
+ preserve_unselected: Whether the loader mapping should retain an explicit null default so
+ unselected components keep their checkpoint dtype.
+
+ Returns:
+ Empty kwargs when no override applies, or one ``dtype`` keyword accepted by the native
+ component loader.
+
+ Raises:
+ ValueError: If either policy contains a selector outside its declared namespace.
+ """
if isinstance(user_policy, torch.dtype):
return {"dtype": user_policy}
if user_policy is None and isinstance(manifest_policy, torch.dtype):
diff --git a/src/flow_factory/models/runtime/classic.py b/src/flow_factory/models/runtime/classic.py
index 897d35a8c..6b1423e27 100644
--- a/src/flow_factory/models/runtime/classic.py
+++ b/src/flow_factory/models/runtime/classic.py
@@ -53,7 +53,15 @@ def _get_materialized_component(self, name: str) -> Any:
return getattr(self.pipeline, name, None)
def physical_route(self, name: str) -> tuple[str, tuple[str, ...]]:
- """Collapse logical aliases that reference one canonical module object."""
+ """Collapse logical aliases that reference one canonical module object.
+
+ Args:
+ name: Declared logical component name to resolve.
+
+ Returns:
+ Canonical root name and an empty nested path. A declared null component retains its
+ own logical root rather than aliasing through the ``None`` singleton.
+ """
self._validate_declared_names([name])
component = self.declared_components[name]
if component is None:
diff --git a/src/flow_factory/models/wan/_conditioning.py b/src/flow_factory/models/wan/_conditioning.py
index 47503ab55..29e13a7a4 100644
--- a/src/flow_factory/models/wan/_conditioning.py
+++ b/src/flow_factory/models/wan/_conditioning.py
@@ -239,7 +239,27 @@ def prepare_wan_i2v_condition_tensors(
device: torch.device,
last_image: Optional[torch.Tensor] = None,
) -> WanI2VConditionTensors:
- """Encode ordered input frames with posterior mode and build Wan condition channels."""
+ """Encode ordered input frames with posterior mode and build Wan condition channels.
+
+ Args:
+ adapter: Realized Wan I2V adapter owning the pipeline and VAE.
+ image: First-frame pixels shaped ``(B, 3, H, W)``.
+ height: Configured output height.
+ width: Configured output width.
+ num_frames: Configured output frame count.
+ dtype: Floating dtype for returned condition tensors.
+ device: Device used for VAE encoding and returned tensors.
+ last_image: Optional last-frame pixels with the same shape as ``image``.
+
+ Returns:
+ Encoded condition channels and, for expanded-timestep checkpoints, the separate
+ first-frame mask.
+
+ Raises:
+ TypeError: If the requested output dtype or VAE dtype is invalid.
+ ValueError: If input tensors, pixel/latent geometry, endpoint support, or temporal
+ divisibility are incompatible with the realized checkpoint.
+ """
if not isinstance(dtype, torch.dtype) or not dtype.is_floating_point:
raise TypeError(f"Wan I2V condition dtype must be floating, received {dtype!r}")
if not isinstance(image, torch.Tensor) or image.ndim != 4:
@@ -319,6 +339,9 @@ def prepare_wan_i2v_condition_tensors(
).to(device=device, dtype=dtype)
latent_condition = normalize_wan_video_latents(adapter, latent_condition)
+ # The condition VAE time axis belongs to the encoded source. Expanded-timestep checkpoints
+ # encode one frame and broadcast it across the target via first_frame_mask, so it must not be
+ # validated against the rollout's num_latent_frames.
condition_latent_frames = (video_condition.shape[2] - 1) // temporal_scale + 1
expected_latents = (
batch_size,
diff --git a/src/flow_factory/models/wan/wan2_i2v.py b/src/flow_factory/models/wan/wan2_i2v.py
index a08ddf66c..270d30998 100644
--- a/src/flow_factory/models/wan/wan2_i2v.py
+++ b/src/flow_factory/models/wan/wan2_i2v.py
@@ -122,7 +122,10 @@ def load_pipeline(self) -> WanImageToVideoPipeline:
)
def _resolve_pipeline_io_contract(self) -> PipelineIOContract:
- """Resolve checkpoint-specific first/last-frame cardinality."""
+ """Resolve exact-one for ordinary/expanded paths and exact-two for FLF2V weights.
+
+ Wan2.2's VAE-only path retains the class-level one-or-two input superset.
+ """
supports_endpoint_pair = False
if not self.pipeline.config.expand_timesteps:
transformer_configs = tuple(
diff --git a/src/flow_factory/optimizer/loader.py b/src/flow_factory/optimizer/loader.py
index eb3232890..bddc8d366 100644
--- a/src/flow_factory/optimizer/loader.py
+++ b/src/flow_factory/optimizer/loader.py
@@ -97,7 +97,9 @@ def build_optimizer(
One optimizer whose ``param_groups`` carry a ``role_name`` per group.
Raises:
- ValueError: If a configuration has no parameters or an unknown type.
+ ValueError: If a configuration has no parameters, uses an unknown optimizer type,
+ selects Muon without matrix parameters, or the current PyTorch build lacks
+ ``torch.optim.Muon``.
"""
adamw_groups: List[Dict[str, Any]] = []
muon_groups: List[Dict[str, Any]] = []
diff --git a/src/flow_factory/samples/samples.py b/src/flow_factory/samples/samples.py
index 86d2e7dd0..b1a1a1533 100644
--- a/src/flow_factory/samples/samples.py
+++ b/src/flow_factory/samples/samples.py
@@ -97,8 +97,8 @@ class BaseSample:
{"height", "width", "latent_index_map", "log_prob_index_map"}
)
- # Fields that must be transported whenever a concrete sample is reconstructed
- # from a partial cross-rank gather, even when no downstream consumer reads them.
+ # Fields ``gather_samples`` must transport to satisfy concrete-class reconstruction
+ # invariants, even when omitted from the consumer-requested field list.
reconstruction_required_fields: ClassVar[frozenset[str]] = frozenset()
# Denoiseing trajectory
diff --git a/src/flow_factory/trainers/loader.py b/src/flow_factory/trainers/loader.py
index 673e5790e..3b8f16199 100644
--- a/src/flow_factory/trainers/loader.py
+++ b/src/flow_factory/trainers/loader.py
@@ -60,29 +60,20 @@ def _requires_ddp_unused_parameter_detection(
def load_trainer(config: Arguments) -> BaseTrainer:
- """
- Factory function to instantiate trainer based on algorithm type.
-
- Uses registry pattern for automatic trainer discovery and loading.
- Supports both built-in trainers and custom algorithms via python paths.
+ """Instantiate the configured trainer after validating its execution plan.
Args:
- config: Configuration containing trainer_type and all hyperparameters
+ config: Parsed configuration containing the trainer, model, optimizer, and backend policy.
Returns:
- An instance of a BaseTrainer subclass
+ The initialized trainer selected by ``config.training_args.trainer_type``.
Raises:
- ImportError: If the trainer is not registered or cannot be imported
-
- Examples:
- # Using built-in trainer
- config.training_args.trainer_type = "grpo"
- trainer = load_trainer(config)
-
- # Using custom trainer
- config.training_args.trainer_type = "my_package.trainers.PPOTrainer"
- trainer = load_trainer(config)
+ ImportError: If the requested trainer cannot be resolved or imported.
+ TypeError: If the registry entry does not resolve to a trainer class.
+ ValueError: If the trainer/adapter contract or distributed optimizer/checkpoint plan is
+ unsupported.
+ RuntimeError: If an FSDP2 checkpoint plan lacks the required plugin state.
"""
# Resolve and validate algorithm semantics before constructing an Accelerator or
# loading model weights. A stale trainer/argument registry pairing must fail with
diff --git a/src/flow_factory/trainers/multirole/backend.py b/src/flow_factory/trainers/multirole/backend.py
index e877d9956..860d082db 100644
--- a/src/flow_factory/trainers/multirole/backend.py
+++ b/src/flow_factory/trainers/multirole/backend.py
@@ -73,17 +73,19 @@ def configure_checkpointing_backend_plan(
accelerator: Accelerator,
training_args: TrainingArguments,
) -> bool:
- """Select a checkpoint owner before model loading and distributed preparation.
+ """Select one checkpointing owner for the active distributed plan.
Args:
accelerator: Runtime backend whose checkpointing policy is being configured.
training_args: Parsed algorithm arguments containing the model checkpoint policy.
Returns:
- Whether a previously realized adapter must disable model checkpointing.
+ True if model-level checkpointing was disabled and a caller holding an already-realized
+ adapter must remove its checkpoint wrappers; otherwise False.
Raises:
ValueError: If FSDP2 is paired with a selective model checkpoint policy.
+ RuntimeError: If FSDP2 backend checkpointing is selected without an FSDP plugin.
"""
if accelerator.distributed_type != DistributedType.FSDP:
return False
diff --git a/src/flow_factory/utils/dist.py b/src/flow_factory/utils/dist.py
index 33af7e60d..51415a076 100644
--- a/src/flow_factory/utils/dist.py
+++ b/src/flow_factory/utils/dist.py
@@ -502,10 +502,10 @@ def gather_samples(
Args:
accelerator: Accelerator instance.
- samples: Local samples on this rank.
- field_names: Fields to gather. When ``'extra_kwargs'`` is included,
- each key inside the dict is gathered independently and reassembled.
- Concrete sample reconstruction fields are added automatically.
+ samples: Local samples on this rank. All entries must share one concrete sample class.
+ field_names: Consumer-requested fields to gather. When ``extra_kwargs`` is included, each
+ key is gathered independently. Fields declared by the concrete class in
+ ``reconstruction_required_fields`` are always added before reconstruction.
device: Target device for tensor fields in the returned samples.
Returns:
From 13dd5d2c3d88f60c54618488ca5f645ad9774b5a Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:39:14 +0800
Subject: [PATCH 67/76] [docs] update final GPU validation status
---
README.md | 28 +++++++++++++++-------------
examples/README.md | 15 ++++++++-------
guidance/datasets.md | 21 +++++++++++++--------
guidance/gpu_validation.md | 26 ++++++++++++++++++++++----
tests/docs/test_minimax_h3_docs.py | 21 ++++++++++++---------
5 files changed, 70 insertions(+), 41 deletions(-)
diff --git a/README.md b/README.md
index c62bb489f..4c49a7424 100644
--- a/README.md
+++ b/README.md
@@ -86,9 +86,9 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
| First/Last-Frame-to-Video | Wan2.1-FLF2V-14B-720P | 14B | wan2_i2v |
| Text-to-Audio-Video | LTX-2 | 19B | ltx2_t2av |
- | LTX-2.3 | 22B | ltx2_t2av |
+ | LTX-2.3 (Diffusers) | 22B | ltx2_t2av |
| Image-to-Audio-Video | LTX-2 | 19B | ltx2_i2av |
- | LTX-2.3 | 22B | ltx2_i2av |
+ | LTX-2.3 (Diffusers) | 22B | ltx2_i2av |
| Text-to-Audio-Video | MiniMax H3 T2VA | 33B | minimax-h3-t2va |
| First/Last-Frame-to-Audio-Video | MiniMax H3 FL2VA | 33B | minimax-h3-fl2va |
| Ordered-Reference-to-Audio-Video | MiniMax H3 Ref2VA | 33B | minimax-h3-ref2va |
@@ -104,11 +104,12 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
> offline DPO shares that exact realization across chosen and rejected arms. See the
> [offline model matrix](guidance/datasets.md#offline-model-support).
-> **MiniMax H3 status:** the T2VA debug and
-> [native-quality FSDP2](examples/grpo/lora/minimax_h3_t2va/quality_720p_fsdp2.yaml)
-> paths are real-weight
-> validated; a completed long-run reward trend is not claimed. FL2VA and Ref2VA remain
-> schema/API and local offline-path validated, pending the documented real-weight GPU matrix.
+> **MiniMax H3 status:** T2VA, FL2VA, and Ref2VA completed all 36 real-weight smoke cells in
+> the documented matrix: three workflows x DDP/DeepSpeed ZeRO-2/FSDP2 x
+> GRPO/SFT/offline DPO/TDM. The FL2VA first-plus-last SFT/offline-DPO gate also passed.
+> The T2VA [native-quality FSDP2](examples/grpo/lora/minimax_h3_t2va/quality_720p_fsdp2.yaml)
+> path has separate initialization, checkpoint, decode, and evaluation coverage. These results
+> do not claim a completed long-run reward trend, convergence, or numerical parity.
> H3 requires B=1, has no CFG, uses neutral guidance `1.0`, and
> keeps separate video/audio trajectories.
> Video uses shift 12, audio uses shift 3, and the model predicts data-ward velocity.
@@ -135,10 +136,10 @@ This experimental feature leverages `diffusers`'s `transformer.set_attention_bac
See [`Algorithm Guidance`](guidance/algorithms.md) for more information.
-> Models and algorithms are decoupled at the framework interface. Validation status varies by example.
-> Training-verified examples carry hardware and reward-trend evidence.
-> MiniMax H3 T2VA has real-weight LoRA validation; FL2VA, Ref2VA, and unlisted
-> combinations require separate training evidence.
+> Models and algorithms are decoupled at the framework interface. The documented ten-mode
+> real-weight smoke matrix completed all 120 model/backend/algorithm cells. Combinations outside
+> that matrix still require separate execution evidence, and smoke completion is not a claim of
+> reward improvement.
# 💾 Hardware Requirements
@@ -147,7 +148,7 @@ See [`Algorithm Guidance`](guidance/algorithms.md) for more information.
## Installation
```bash
-git clone https://github.com/Jayce-Ping/Flow-Factory.git
+git clone https://github.com/X-GenGroup/Flow-Factory.git
cd Flow-Factory
pip install -e .
```
@@ -244,7 +245,8 @@ Prompt and input-condition encodings are cached. Target, chosen, and rejected me
encoded on the fly; their VAE latents are never stored in the preprocessing cache. One offline
epoch is one complete dataloader traversal sharded by PyTorch's official `DistributedSampler`. See the
[dataset guide](guidance/datasets.md#offline-v2-records) for the full schema and cadence rules, and
-the [GPU validation plan](guidance/gpu_validation.md) for the 120-job model/backend/algorithm matrix.
+the [completed GPU validation matrix](guidance/gpu_validation.md) for the 120 main jobs and 24
+additional dynamic gates.
The [offline smoke builder](dataset/offline_smoke/README.md) reconstructs independent SFT and
offline-DPO mini datasets for every currently implemented image, video, and audio-video profile.
diff --git a/examples/README.md b/examples/README.md
index 2765c7800..7d93dbd34 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -88,13 +88,14 @@ The T2VA `debug.yaml` recipe is real-weight validated with the 61 GB checkpoint
(61.74 GiB transformer):
1 GPU and 16 GPUs across two nodes completed CPS rollout, video/audio decode,
CLAP reward, GRPO replay/backward/optimizer step, and LoRA checkpoint save/resume.
-Its 64x96 canvas is intentionally a correctness geometry. The quality-oriented T2VA
-default is now the shared-`vid_prompt`, LoRA-rank-64 baseline aligned with the LTX2
-T2AV recipe and uses both CLAP and ImageBind rewards. It is configuration/API
-validated; no completed long-run reward trend is claimed. FL2VA and Ref2VA are also
-schema/API and local offline-path validated, rather than claims of real-weight training stability
-or reward improvement. The complete follow-up campaign is defined in the
-[GPU validation plan](../guidance/gpu_validation.md).
+Its 64x96 canvas is intentionally a correctness geometry. The PR #220 real-weight smoke campaign
+subsequently completed all 36 H3 main cells: T2VA, FL2VA, and Ref2VA across
+DDP/DeepSpeed ZeRO-2/FSDP2 and GRPO/SFT/offline DPO/TDM. The FL2VA first-plus-last
+SFT/offline-DPO variant gate also passed. The quality-oriented T2VA default remains the
+shared-`vid_prompt`, LoRA-rank-64 baseline aligned with the LTX2 T2AV recipe and uses both CLAP
+and ImageBind rewards. These smoke results establish execution coverage, not a completed
+long-run reward trend, convergence, or numerical parity. See the
+[GPU validation matrix](../guidance/gpu_validation.md).
The T2VA `quality_720p_fsdp2.yaml` recipe is the active native-quality path:
768x1344, 124 frames, 24 denoising steps, LoRA rank 64 / alpha 128, and two
diff --git a/guidance/datasets.md b/guidance/datasets.md
index 81da9aaac..c4ebe0eea 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -256,14 +256,17 @@ output semantics.
|---|---|---|
| Supported | `sd3-5`, `flux1`, `flux1-kontext`, `flux2`, `flux2-klein`, `qwen-image`, `qwen-image-edit-plus`, `z-image`, `bagel`, `sensenova` | Image-output codecs with adapter-specific geometry and packing. SenseNova uses the existing grouped `images` input with within-type order, not heterogeneous references. |
| Supported | `wan2_t2v` | Video targets require `fps`; the codec resamples to configured frames/rate and samples the Wan VAE posterior on the fly. |
-| Supported | `wan2_i2v` | Input media binds a required `first_frame` image and an optional `last_frame` image. Condition pixels are cached at configured geometry, then encoded with VAE posterior mode once per batch. Expanded-timestep TI2V checkpoints accept the first frame only because official Diffusers ignores a last image in that mode. Video targets require `fps`. Offline execution is B=1. |
+| Supported | `wan2_i2v` | Input media uses a checkpoint-specific contract. Expanded-timestep TI2V and standard Wan2.1 I2V accept exactly `first_frame`; dedicated Wan2.1 FLF2V requires both `first_frame` and `last_frame`; Wan2.2 I2V-A14B accepts `first_frame` plus an optional VAE-only `last_frame`. Condition pixels are cached at configured geometry and encoded with VAE posterior mode once per batch. Video targets require `fps`; offline execution is B=1. |
| Supported | `ltx2_t2av`, `ltx2_i2av` | Every candidate is an exact ordered `(video, audio)` pair with required `fps` and `sample_rate`. Both streams are aligned to the official LTX2 clock and encoded/packed on the fly. I2AV requires the `first_frame` image slot, substitutes its posterior-mode first latent into each target, and excludes the pinned tokens with an active mask. |
| Supported | `minimax-h3-t2va`, `minimax-h3-fl2va`, `minimax-h3-ref2va` | Every candidate is an exact ordered `(video, audio)` pair. FL2VA accepts `first_frame`, `last_frame`, or both slots; Ref2VA accepts 1-12 globally ordered references and requires at least one image or video. Conditioned workflows realize one official prefix per batch, shared by both offline-DPO candidates and policy/reference forwards. H3 remains B=1. |
-Wan first/last semantics use generic semantic slots, not model-specific schema keys. The first
-frame is required; the last frame is optional. Unslotted input remains a positional convenience,
-but an explicit slot is recommended for sparse or generated manifests. The target is the complete
-generated video: its first frame, and its final frame when provided, correspond to the conditions.
+Wan endpoint semantics use generic semantic slots, not model-specific schema keys. The realized
+slot cardinality is checkpoint-specific: expanded-timestep TI2V and standard Wan2.1 I2V accept
+exactly `first_frame`; dedicated Wan2.1 FLF2V requires both endpoints; Wan2.2 I2V-A14B accepts
+`first_frame` plus an optional VAE-only `last_frame`. Unslotted input remains a positional
+convenience, but explicit slots are recommended for sparse or generated manifests. The target is
+the complete generated video, whose conditioned endpoint or endpoints must correspond to the
+supplied images.
```jsonl
{"schema_version":2,"input":{"prompt":"A paper boat crosses the pond.","media":[{"type":"image","path":"conditions/first.png","slot":"first_frame"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/first-only.mp4","fps":24.0}]}},"metadata":{}}
@@ -407,9 +410,11 @@ checkpoint save/resume. Its 64x96 canvas validates
correctness and memory fit, not visual quality or reward improvement.
`quality_720p_fsdp2.yaml` has real-weight initialization, checkpoint, native-resolution
decode, and evaluation coverage; no long-run reward trend is claimed. The aligned default uses
-the shared `dataset/vid_prompt` source, LoRA rank 64, and CLAP plus ImageBind rewards; it remains a
-configuration/API-validated baseline without a published long-run trend. FL2VA and Ref2VA remain
-schema/API-validated starting points.
+the shared `dataset/vid_prompt` source, LoRA rank 64, and CLAP plus ImageBind rewards. The PR #220
+smoke campaign completed all 36 H3 main cells: T2VA, FL2VA, and Ref2VA across
+DDP/DeepSpeed ZeRO-2/FSDP2 and GRPO/SFT/offline DPO/TDM. The FL2VA first-plus-last
+SFT/offline-DPO variant gate also passed. This is execution coverage, not a published long-run
+reward trend, convergence, or numerical-parity claim.
### T2VA: `minimax-h3-t2va`
diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md
index 23b8a6d37..8c6c6a3dd 100644
--- a/guidance/gpu_validation.md
+++ b/guidance/gpu_validation.md
@@ -1,8 +1,26 @@
-# GPU Validation Plan
+# GPU Validation Matrix
-This document is the handoff contract for real-weight GPU validation. It does
-not claim that an unexecuted combination is supported. A combination becomes
-validated only after its artifacts satisfy the acceptance criteria below.
+This document defines the real-weight GPU validation contract and records the PR #220 result.
+A combination is validated only after its execution evidence satisfies the acceptance criteria
+below.
+
+## PR #220 result
+
+The complete dynamic smoke scope passed **144/144** unique jobs:
+
+- **120/120 main jobs**: 10 semantic modes x DDP/DeepSpeed ZeRO-2/FSDP2 x
+ GRPO/SFT/offline DPO/TDM.
+- **22/22 checkpoint-variant jobs**: Wan 2.1/2.2 T2V and I2V variants, all six strict Wan A14B
+ dual-transformer routing jobs, LTX 2.3 T2AV/I2AV, and H3 FL2VA first-plus-last coverage.
+- **2/2 supplemental jobs**: Flux1-Kontext image-to-image SFT and offline DPO.
+
+All **132/132** checkpoint-variant backend/algorithm cells also passed static configuration and
+contract validation. Four positive Muon jobs passed for DDP/FSDP2 with SFT and mixed-role TDM;
+the DeepSpeed ZeRO-2 negative gate rejected Muon before model loading as intended.
+
+These are reduced-geometry, finite-length execution smokes. They establish model loading,
+distributed routing, forward/backward, optimizer, and finite-data behavior; they do not claim
+convergence, long-run reward improvement, quality parity, or numerical parity.
## Environment gate
diff --git a/tests/docs/test_minimax_h3_docs.py b/tests/docs/test_minimax_h3_docs.py
index cb4f5e0f2..edbe2af12 100644
--- a/tests/docs/test_minimax_h3_docs.py
+++ b/tests/docs/test_minimax_h3_docs.py
@@ -32,10 +32,10 @@ def test_readme_documents_h3_links_dependency_and_limits() -> None:
"The configurations under `examples/` have been verified to yield measurable "
"performance gains."
) not in text
- assert "Validation status varies by example" in text
- assert "hardware and reward-trend evidence" in text
- assert "MiniMax H3 T2VA has real-weight LoRA validation" in text
- assert "FL2VA and Ref2VA remain" in text
+ assert "all 120 model/backend/algorithm cells" in text
+ assert "smoke completion is not a claim of" in text
+ assert "all 36 real-weight smoke cells" in text
+ assert "FL2VA first-plus-last" in text
assert "T2VA is real-weight validated on 1 and 16 GPUs" not in text
assert text.count("30B | ") == 3
@@ -59,7 +59,7 @@ def test_readme_documents_h3_links_dependency_and_limits() -> None:
"N transitions",
"N + 1 states",
"30B",
- "completed long-run reward trend is not claimed",
+ "do not claim a completed long-run reward trend",
"[Datasets](guidance/datasets.md)",
):
assert required in text
@@ -71,9 +71,9 @@ def test_examples_readme_links_h3_and_separates_validation_levels() -> None:
relative_link = f"../{root_link}"
assert relative_link in text
assert (ROOT / root_link).is_file()
- assert "schema/API and local offline-path validated" in text
- assert "GPU validation plan" in text
- assert "hardware" in text
+ assert "all 36 H3 main cells" in text
+ assert "GPU validation matrix" in text
+ assert "execution coverage" in text
assert "reward" in text
assert "61 GB" in text
assert "ImageBind" in text
@@ -81,10 +81,13 @@ def test_examples_readme_links_h3_and_separates_validation_levels() -> None:
assert "NonCommercial" in text
-def test_gpu_validation_plan_declares_the_complete_smoke_matrix() -> None:
+def test_gpu_validation_matrix_declares_scope_and_completed_result() -> None:
text = _text("guidance/gpu_validation.md")
assert "10 x 3 x 4 = 120 jobs" in text
+ for result in ("144/144", "120/120", "22/22", "2/2", "132/132"):
+ assert result in text
+ assert "Four positive Muon jobs passed" in text
for mode in (
"sd35-t2i",
"bagel-mri2i",
From 3a6d414785b1683c594a27fc2c364d1a90e95284 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:49:19 +0800
Subject: [PATCH 68/76] [docs] clarify GPU evidence boundaries
---
guidance/gpu_validation.md | 24 +++++++++++++++---------
1 file changed, 15 insertions(+), 9 deletions(-)
diff --git a/guidance/gpu_validation.md b/guidance/gpu_validation.md
index 8c6c6a3dd..abe09dcd4 100644
--- a/guidance/gpu_validation.md
+++ b/guidance/gpu_validation.md
@@ -1,26 +1,32 @@
# GPU Validation Matrix
-This document defines the real-weight GPU validation contract and records the PR #220 result.
-A combination is validated only after its execution evidence satisfies the acceptance criteria
+This document defines the real-weight GPU validation contract and records the PR #220 execution
+result. The historical direct runs establish successful launch and training completion; a future
+formal campaign should additionally capture every artifact listed in the acceptance criteria
below.
## PR #220 result
-The complete dynamic smoke scope passed **144/144** unique jobs:
+The complete dynamic smoke scope reached its successful terminal marker for **144/144** unique
+jobs:
- **120/120 main jobs**: 10 semantic modes x DDP/DeepSpeed ZeRO-2/FSDP2 x
GRPO/SFT/offline DPO/TDM.
-- **22/22 checkpoint-variant jobs**: Wan 2.1/2.2 T2V and I2V variants, all six strict Wan A14B
- dual-transformer routing jobs, LTX 2.3 T2AV/I2AV, and H3 FL2VA first-plus-last coverage.
+- **22/22 formal variant-gate jobs**: Wan 2.1/2.2 T2V and I2V routing variants (including
+ corrected redundant FLF2V coverage), all six strict Wan A14B dual-transformer routing jobs,
+ LTX 2.3 T2AV/I2AV, and H3 FL2VA first-plus-last coverage.
- **2/2 supplemental jobs**: Flux1-Kontext image-to-image SFT and offline DPO.
All **132/132** checkpoint-variant backend/algorithm cells also passed static configuration and
contract validation. Four positive Muon jobs passed for DDP/FSDP2 with SFT and mixed-role TDM;
the DeepSpeed ZeRO-2 negative gate rejected Muon before model loading as intended.
-These are reduced-geometry, finite-length execution smokes. They establish model loading,
-distributed routing, forward/backward, optimizer, and finite-data behavior; they do not claim
-convergence, long-run reward improvement, quality parity, or numerical parity.
+These are reduced-geometry, finite-length execution smokes. They establish archived real-weight
+execution coverage for model loading, distributed backends, optimizer paths, and finite-data
+termination. Capability-sensitive instrumented gates additionally checked routing and updates,
+but earlier direct runs did not uniformly capture the complete command, environment, and metrics
+artifact set required below. These results do not claim convergence, long-run reward improvement,
+quality parity, or numerical parity.
## Environment gate
@@ -180,7 +186,7 @@ backends and all four algorithms.
| Family mode | Additional checkpoint variants |
|---|---|
-| Wan T2V | `Wan2.1-T2V-14B-Diffusers`, `Wan2.2-TI2V-5B-Diffusers`, `Wan2.2-T2V-A14B-Diffusers` |
+| Wan T2V/TI2V family | `Wan2.1-T2V-14B-Diffusers`, `Wan2.2-TI2V-5B-Diffusers`, `Wan2.2-T2V-A14B-Diffusers` |
| Wan I2V first-only | `Wan2.1-I2V-14B-480P-Diffusers`, `Wan2.1-I2V-14B-720P-Diffusers`, `Wan2.2-I2V-A14B-Diffusers` |
| Wan first/last | `Wan2.2-I2V-A14B-Diffusers` |
| LTX2 T2AV and I2AV | `dg845/LTX-2.3-Diffusers` |
From 8d69cb539472dc3523fa20494fc3a26a8092a59c Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:51:21 +0800
Subject: [PATCH 69/76] [optimizer] docs: correct Muon availability guidance
---
.agents/knowledge/dependencies.md | 13 ++++++++++---
.agents/knowledge/topics/component_variants.md | 5 +++--
guidance/workflow.md | 3 ++-
src/flow_factory/optimizer/loader.py | 2 +-
4 files changed, 16 insertions(+), 7 deletions(-)
diff --git a/.agents/knowledge/dependencies.md b/.agents/knowledge/dependencies.md
index d6d52ef18..79198cd3e 100644
--- a/.agents/knowledge/dependencies.md
+++ b/.agents/knowledge/dependencies.md
@@ -38,7 +38,7 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
| Package | Min Version | Purpose |
|---------|-------------|---------|
-| `torch` | >= 2.6.0 | PyTorch core |
+| `torch` | >= 2.6.0 | PyTorch core and AdamW baseline; Muon needs the optional API noted below |
| `torchvision` | >= 0.19.0 | Vision utilities |
| `torchaudio` | >= 2.4.0 | Audio I/O (audio / audio-video models, CLAP) |
| `transformers` | >= 4.57.1 | Text encoders, tokenizers |
@@ -56,6 +56,13 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
- Only **ZeRO-1** and **ZeRO-2** are supported. ZeRO-3 is broken for reward model sharding (constraint #10).
- DeepSpeed is optional — Accelerate alone handles most distributed scenarios.
+### Muon
+- The core `torch>=2.6.0` floor does not guarantee `torch.optim.Muon`. Selecting
+ `optimizer: muon` requires a build that exposes that API (included in standard releases from
+ PyTorch 2.9); runtime capability detection remains authoritative.
+- Muon is supported with DDP and FSDP2. The pre-load optimizer/backend validator rejects
+ DeepSpeed and FSDP1 before pretrained weights are loaded.
+
### diffusers
- Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule
may be used for upstream development, but must not silently override the declared runtime dependency.
@@ -101,5 +108,5 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
## Cross-refs
-- `constraints.md` #10 (DeepSpeed ZeRO-3 unsupported)
-- `architecture.md` "Configuration Hierarchy" (hparams structure)
+- UP: [`constraints.md` #10](constraints.md#10-deepspeed-zero-3-is-unsupported), [Architecture Configuration Hierarchy](architecture.md#configuration-hierarchy)
+- PEER: [Component Variants](topics/component_variants.md)
diff --git a/.agents/knowledge/topics/component_variants.md b/.agents/knowledge/topics/component_variants.md
index a6c1d5735..e3793f6e0 100644
--- a/.agents/knowledge/topics/component_variants.md
+++ b/.agents/knowledge/topics/component_variants.md
@@ -154,7 +154,8 @@ its children's groups as one list. An all-AdamW run still gets a plain
Muon therefore gives one variant **two** parameter groups, which is why
`OptimizationRole.optimizer_group_ids` is a tuple. It requires a PyTorch build that
-exposes `torch.optim.Muon`. DeepSpeed is rejected as unverified because it rebuilds
+exposes `torch.optim.Muon` (included in standard releases from PyTorch 2.9); the runtime feature
+check, rather than the version string alone, is authoritative. DeepSpeed is rejected as unverified because it rebuilds
its own optimizer wrapper, while FSDP1 is rejected because its flat parameters erase
the required matrix rank; the supported distributed plans are DDP and FSDP2.
@@ -192,5 +193,5 @@ family describes it.
## Cross-refs
-- UP: [`constraints.md` #10](../constraints.md), [`architecture.md` Component Management](../architecture.md#component-management)
+- UP: [`constraints.md` #10](../constraints.md#10-deepspeed-zero-3-is-unsupported), [`architecture.md` Component Management](../architecture.md#component-management)
- PEER: [Component Runtime](component_runtime.md), [Structured Trajectory](structured_trajectory.md), [Autocast and Parameter Swaps](autocast_param_swap.md)
diff --git a/guidance/workflow.md b/guidance/workflow.md
index a52e24c3d..3ad9f6f82 100644
--- a/guidance/workflow.md
+++ b/guidance/workflow.md
@@ -669,7 +669,8 @@ AdamW for its biases, normalization scales and embeddings, which the `fallback_`
fields configure. `optimizer/loader.py` wraps that pair in a `CompositeOptimizer` so
the framework still prepares exactly one root. An all-AdamW run gets a plain
`torch.optim.AdamW`, unchanged. Muon requires a PyTorch build that exposes
-`torch.optim.Muon` (2.10 or newer in supported environments). Muon combined with
+`torch.optim.Muon` (included in standard releases from PyTorch 2.9; runtime capability detection
+is authoritative). Muon combined with
DeepSpeed is refused at startup as unverified, and FSDP1 flattens matrices into
incompatible parameters; use DDP or FSDP2.
diff --git a/src/flow_factory/optimizer/loader.py b/src/flow_factory/optimizer/loader.py
index bddc8d366..71b78a153 100644
--- a/src/flow_factory/optimizer/loader.py
+++ b/src/flow_factory/optimizer/loader.py
@@ -39,7 +39,7 @@ def validate_muon_available() -> None:
if not hasattr(torch.optim, "Muon"):
raise ValueError(
f"torch.optim.Muon is unavailable in PyTorch {torch.__version__}. Install a "
- "PyTorch build that provides Muon (2.10 or newer in supported environments), "
+ "PyTorch build that provides Muon (standard releases include it from 2.9), "
"or select the adamw optimizer."
)
From 60201685fc89f75fff474a89bb4bbcaaecda4e89 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:53:01 +0800
Subject: [PATCH 70/76] [agents] docs: align knowledge with current contracts
---
.agents/knowledge/README.md | 5 +++
.agents/knowledge/architecture.md | 16 +++++-----
.agents/knowledge/constraints.md | 31 ++++++++++++++-----
.agents/knowledge/docs_maintenance.md | 12 +++++--
.../knowledge/topics/adapter_conventions.md | 5 +--
.../knowledge/topics/autocast_param_swap.md | 6 ++--
.agents/knowledge/topics/component_runtime.md | 22 ++++++++++++-
.agents/knowledge/topics/dtype_precision.md | 7 ++---
.agents/knowledge/topics/fix_patterns.md | 5 ++-
.agents/knowledge/topics/minimax_h3.md | 17 +++++-----
.agents/knowledge/topics/parity_testing.md | 4 +--
.agents/knowledge/topics/sample_lifecycle.md | 22 +++++++++----
.agents/knowledge/topics/samplers.md | 4 +--
.agents/knowledge/topics/timestep_sigma.md | 5 ++-
.../topics/train_inference_consistency.md | 5 ++-
.cursor/rules/agents-docs-maintenance.mdc | 2 +-
AGENTS.md | 6 ++--
17 files changed, 115 insertions(+), 59 deletions(-)
diff --git a/.agents/knowledge/README.md b/.agents/knowledge/README.md
index ba09d48d4..60f7eaa77 100644
--- a/.agents/knowledge/README.md
+++ b/.agents/knowledge/README.md
@@ -7,13 +7,18 @@
| Touching dtype/precision, mixed precision config, debugging NaN/overflow | `topics/dtype_precision.md` |
| Editing a trainer `optimize()` loop / autocast scope, ref/EMA/named param swaps | `topics/autocast_param_swap.md` |
| Adding or modifying a model adapter | `topics/adapter_conventions.md` |
+| Changing offline `PipelineIOContract`, condition-state preparation, output codecs, or output geometry | `topics/adapter_conventions.md`, `topics/component_runtime.md` |
| Changing component discovery, loading, lifecycle, or distributed preparation | `topics/component_runtime.md` |
+| Changing FSDP loading or activation-checkpoint ownership | `topics/component_runtime.md` |
| Touching rollout collection, replay bridges, index maps, or multi-component order | `topics/structured_trajectory.md` |
| Adding an algorithm that trains several model copies at once, or changing per-variant LoRA/full storage or variant checkpointing | `topics/component_variants.md` |
+| Changing optimizer roles, Muon, `CompositeOptimizer`, or optimizer/backend compatibility | `topics/component_variants.md`, `dependencies.md` |
| Adding adapter, upgrading diffusers, debugging output quality | `topics/parity_testing.md` |
+| Touching MiniMax H3 workflows, dependency pins, ordered references, or H3 memory policies | `topics/minimax_h3.md` |
| Touching `TimeSampler`, `adapter.forward(t=...)`, `timestep_range`, `flow_match_sigma` | `topics/timestep_sigma.md` |
| Editing `data_utils/sampler*`, hparams sampler/batch fields | `topics/samplers.md` |
| Touching `sample()`/`optimize()` data flow, debugging `sample()`/`optimize()` OOM, adding high-resolution / video example configs | `topics/sample_lifecycle.md` |
+| Changing `BaseSample`, partial sample gathering, or concrete-sample reconstruction | `topics/sample_lifecycle.md` |
| After completing a bug fix | `topics/fix_patterns.md` |
| Changing `pyproject.toml`, deps, install commands | `dependencies.md` |
| Adding or editing `.agents/` documentation | `docs_maintenance.md` |
diff --git a/.agents/knowledge/architecture.md b/.agents/knowledge/architecture.md
index 84baee949..226bae25a 100644
--- a/.agents/knowledge/architecture.md
+++ b/.agents/knowledge/architecture.md
@@ -89,12 +89,14 @@ All four registries map string keys → lazy import paths. Resolution: registry
| `nft` | `DiffusionNFTTrainer` | Decoupled | `BaseTrainer` |
| `awm` | `AWMTrainer` | Decoupled | `BaseTrainer` |
| `crd` | `CRDTrainer` | Decoupled | `BaseTrainer` |
-| `diffusion-opd` | `DiffusionOPDTrainer` | Distillation (on-policy) | `BaseTrainer` |
-| `dmd2` | `DMD2Trainer` | Distillation | `BaseTrainer` |
-| `tdm` | `TDMTrainer` | Distillation | `BaseTrainer` |
-| `tdm-r1` | `TDMR1Trainer` | Distillation + reward | `BaseTrainer` |
+| `diffusion-opd` | `DiffusionOPDTrainer` | Distillation, generation + no feedback | `BaseTrainer` |
+| `dmd2` | `DMD2Trainer` | Distillation, generation + no feedback, ODE | `BaseTrainer` |
+| `tdm` | `TDMTrainer` | Distillation, generation + no feedback, ODE | `BaseTrainer` |
+| `tdm-r1` | `TDMR1Trainer` | Decoupled, generation + runtime reward, ODE | `TDMTrainer` |
-**Flat hierarchy**: New trainers inherit from `BaseTrainer` directly. The sanctioned exceptions are `GRPOGuardTrainer → GRPOTrainer` and `DPPOTrainer → GRPOTrainer` (strict GRPO loss variants; see constraint #11).
+**Flat hierarchy**: New trainers inherit from `BaseTrainer` directly. The sanctioned existing
+extensions are `GRPOGuardTrainer → GRPOTrainer`, `DPPOTrainer → GRPOTrainer`, and
+`TDMR1Trainer → TDMTrainer`; see constraint #11.
**Model Adapters** (`models/registry.py`):
| Key | Class | Task |
@@ -288,13 +290,13 @@ Details: `topics/component_variants.md`.
### Reward Processing
`RewardProcessor` dispatches by model type:
-- **Pointwise**: batch by `batch_size`
+- **Pointwise**: applicable sub-batches of at most `batch_size`
- **Groupwise**: group by `unique_id` (local or distributed path)
- **Multi-reward**: weighted aggregation
- **Async**: optional non-blocking computation
### Advantage Computation
-`AdvantageProcessor` (`advantage/advantage_processor.py`): communication-aware, auto-selects gather vs local path. Strategies: `"sum"` (GRPO) and `"gdpo"`. All reward-based trainers delegate to `self.advantage_processor.compute_advantages()`; the distillation trainer `diffusion-opd` is the exception (its `prepare_feedback()` is a no-op — no reward/advantage stage).
+`AdvantageProcessor` (`advantage/advantage_processor.py`): communication-aware, auto-selects gather vs local path. Strategies: `"sum"` (GRPO) and `"gdpo"`. Runtime-reward trainers delegate to `self.advantage_processor.compute_advantages()`. Feedback-`none` trainers (`diffusion-opd`, DMD2, and TDM) bypass reward and advantage stages structurally.
### Configuration Hierarchy
```
diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md
index 0e506e9f0..77fcc746b 100644
--- a/.agents/knowledge/constraints.md
+++ b/.agents/knowledge/constraints.md
@@ -12,7 +12,11 @@ These constraints MUST NOT be violated. Consult this file before making any code
The four registries (`_TRAINER_REGISTRY`, `_MODEL_ADAPTER_REGISTRY`, `_REWARD_MODEL_REGISTRY`, `_ACCELERATOR_REGISTRY`) map string identifiers to **fully qualified Python class paths** for lazy import. If you move, rename, or restructure a class, the corresponding registry entry MUST be updated, or `ImportError` will occur at runtime.
### 2. Registry Identifier Convention
-Registry keys are **case-insensitive** (lowered at lookup). Model adapter keys use lowercase with hyphens (e.g., `flux1-kontext`). Trainer keys use lowercase (e.g., `grpo-guard`). Reward keys use lowercase (e.g., `pickscore`). New entries must follow the same convention.
+Registry keys are **case-insensitive** (lowered at lookup). Preserve the canonical registered
+spelling: most model adapter keys use lowercase with hyphens (for example `flux1-kontext`), while
+the public Wan/LTX2 keys retain underscores (for example `wan2_t2v` and `ltx2_t2av`). Trainer and
+reward keys use lowercase (for example `grpo-guard` and `pickscore`). New entries must choose one
+canonical lowercase spelling and use it consistently across registries, arguments, and examples.
### 3. Dynamic Import Fallback
All four registries support a **direct Python path** fallback (e.g., `my_package.models.CustomAdapter`). If an identifier is not found in the registry, it is treated as a fully qualified import path. Do not break this two-mode resolution logic.
@@ -47,8 +51,8 @@ either axis from batch fields or make it user-configurable independently of `tra
### 7. Coupled vs Decoupled Paradigm
- **Coupled** (GRPO, GRPO-Guard, DPPO): Training timesteps are coupled with SDE-based sampling. Requires log-probability computation. Must use SDE dynamics (`Flow-SDE`, `Dance-SDE`, `CPS`).
-- **Decoupled** (SFT, offline DPO, online DPO, NFT, AWM, DGPO, CRD): Training timesteps are decoupled from sampling. Can use any dynamics including `ODE`.
-- **Distillation** (`diffusion-opd`): On-policy multi-teacher distillation; dynamics-agnostic (ODE or SDE) and has no reward/advantage stage.
+- **Decoupled** (SFT, offline DPO, online DPO, NFT, AWM, DGPO, CRD, TDM-R1): Training timesteps are decoupled from sampling. They may use ODE subject to algorithm-specific rules; TDM-R1 requires ODE.
+- **Distillation** (`diffusion-opd`, DMD2, TDM): Generated acquisition with no runtime reward/advantage stage. DiffusionOPD supports ODE or SDE; DMD2 and TDM require ODE.
Mixing paradigms (e.g., using `ODE` dynamics with `GRPO`) will produce incorrect gradients silently.
@@ -75,7 +79,7 @@ offline epoch means one complete traversal of that resulting finite loader.
Checkpoints are written and read for **trainable members only** — components whose `target_module_map[name]` is non-empty (`adapter.trainable_component_names`). Frozen-but-shardable bundle members (e.g. Wan2.2's `transformer_2`, kept in `target_components` only to be FSDP-sharded for memory; see #9) map to `None` and are skipped by both `save_checkpoint` and `_load_lora`/`_load_full_model`. Loaders MUST iterate `trainable_component_names`, not `target_components`, or resume logs a spurious error for a per-component subdir that was never written. `resume_type='state'` restores via `accelerator.load_state` into the prepared bundle root and is therefore keyed to bundle membership — resuming into a different `target_components` / bundle composition will mismatch.
### 10. DeepSpeed ZeRO-3 Is Unsupported
-Supported distributed plans are DDP, FSDP, and DeepSpeed ZeRO-1/2. Reward model sharding under ZeRO-3 is broken even with DeepSpeed's own `zero.GatheredParameters` context manager, and parameter sharding also breaks frozen-component synchronization. `validate_supported_distributed_plan` (`trainers/abc.py`) rejects it at `BaseTrainer.__init__`, before any weights load, and `config/deepspeed/` ships no ZeRO-3 profile. Multi-role training narrows this further: `_validate_multirole_backend` requires ZeRO-1/2 and, under FSDP2, `use_orig_params=True`.
+Supported distributed plans are DDP, FSDP, and DeepSpeed ZeRO-1/2. Reward model sharding under ZeRO-3 is broken even with DeepSpeed's own `zero.GatheredParameters` context manager, and parameter sharding also breaks frozen-component synchronization. `validate_supported_distributed_plan` is defined in `trainers/multirole/backend.py`; `trainers/loader.py` calls it before model construction, while `BaseTrainer.__init__` repeats the check defensively. `config/deepspeed/` ships no ZeRO-3 profile. Multi-role training narrows this further: `_validate_multirole_backend` requires ZeRO-1/2 and, under FSDP2, `use_orig_params=True`. Muon narrows the plan independently to DDP/FSDP2 and requires a build exposing `torch.optim.Muon`; `validate_optimizer_backend_plan` rejects DeepSpeed, FSDP1, and an unavailable Muon API before weights load.
---
@@ -94,7 +98,14 @@ preparation, rewards, and advantage processing.
`optimize_batch()`. Online `DPOTrainer` forms pairs at `optimize()` entry. Offline DPO consumes
dataset pairs directly.
-**Trainer hierarchy**: New trainers MUST inherit directly from `BaseTrainer`. The only sanctioned exceptions are strict behavioral variants of GRPO that change only the per-step loss while reusing GRPO's sampling/advantage/eval machinery: `GRPOGuardTrainer → GRPOTrainer` (adds ratio-normalization) and `DPPOTrainer → GRPOTrainer` (replaces the PPO ratio-clip with a KL trust-region mask). Trainer-to-trainer inheritance creates fragile coupling; when in doubt, inherit from `BaseTrainer` and extract shared logic into helper methods. All reward-based trainers delegate advantage computation to `self.advantage_processor.compute_advantages()`; the distillation trainer `diffusion-opd` is the exception (its `prepare_feedback()` is a no-op with no reward/advantage stage).
+**Trainer hierarchy**: New trainers MUST inherit directly from `BaseTrainer`. The sanctioned
+existing behavioral extensions are `GRPOGuardTrainer → GRPOTrainer`, `DPPOTrainer → GRPOTrainer`,
+and `TDMR1Trainer → TDMTrainer`; TDM-R1 reuses TDM's deterministic trajectory and multi-role
+runtime while restoring runtime reward feedback. Trainer-to-trainer inheritance creates fragile
+coupling; when in doubt, inherit from `BaseTrainer` and extract shared logic into helper methods.
+Every runtime-reward trainer delegates advantage computation to `AdvantageProcessor`. Trainers
+with feedback `none` (`diffusion-opd`, DMD2, and TDM) bypass reward/advantage structurally; their
+no-op `prepare_feedback()` overrides are compatibility shims, not the execution mechanism.
### 12. BaseAdapter Abstract Methods
Subclasses of `BaseAdapter` MUST implement these **4 abstract methods**:
@@ -123,14 +134,20 @@ policy stays explicit at their semantic boundaries. Unsupported adapters declare
**Adapter hierarchy**: All model adapters MUST inherit directly from `BaseAdapter` — never from another adapter. Shared logic between adapters for the same model family should use private helper functions, code duplication, or mixins — not adapter-to-adapter inheritance. Adapter subclassing creates fragile coupling where changes to a parent adapter silently break child adapters, and makes the 4-abstract-method contract harder to verify (the 4 per-modality encoders have no-op defaults, so a fresh subclass of `BaseAdapter` is always valid; chained inheritance hides which encoder a model actually overrides).
### 13. BaseRewardModel Paradigm Split
-- `PointwiseRewardModel.__call__` receives batches of size `batch_size`, returns rewards of shape `(batch_size,)`
-- `GroupwiseRewardModel.__call__` receives all samples in a group (size `group_size`), returns rewards of shape `(group_size,)`
+- `PointwiseRewardModel.__call__` receives an applicable sub-batch of at most configured
+ `batch_size` and returns one reward per received sample.
+- `GroupwiseRewardModel.__call__` receives one complete applicable `unique_id` group and returns
+ one reward per group member in input order.
The `RewardProcessor` dispatches differently based on the model type. Do not change the calling convention.
### 14. Sample Dataclass Hierarchy
`BaseSample` → `T2ISample`, `ImageConditionSample`, `T2VSample`, `T2AVSample`, etc. The `_shared_fields` class variable determines which fields are NOT stacked across a batch. Incorrect `_shared_fields` causes silent data corruption during collation.
+`reconstruction_required_fields` separately names fields required to instantiate a concrete sample
+after a partial distributed gather, even when the downstream consumer did not request them. See
+[`topics/sample_lifecycle.md`](topics/sample_lifecycle.md#partial-gather-reconstruction).
+
**Two-layer hierarchy**: Task-level samples (`T2ISample`, `I2VSample`, `I2AVSample`, ...) are defined in `samples/samples.py` and inherit from `BaseSample` or its condition mixins (`ImageConditionSample`, `VideoConditionSample`). Model-specific samples (`LTX2Sample`, `LTX2I2AVSample`, ...) MUST inherit from the appropriate task-level sample — never from another model-specific sample across files. This mirrors the flat adapter hierarchy: `LTX2I2AVSample(I2AVSample)`, NOT `LTX2I2AVSample(LTX2Sample)`.
Legacy trajectory fields remain authoritative when `BaseSample.trajectory is None`. Structured
diff --git a/.agents/knowledge/docs_maintenance.md b/.agents/knowledge/docs_maintenance.md
index edcdee136..48df8cd3f 100644
--- a/.agents/knowledge/docs_maintenance.md
+++ b/.agents/knowledge/docs_maintenance.md
@@ -9,8 +9,8 @@
The knowledge system uses a 3-layer design with bidirectional cross-references:
```
-Root: AGENTS.md — project identity, behavioral principles (one-liners)
-Tier 1: philosophy.md, constraints.md, — core thesis + concise indexes (always read)
+Root: AGENTS.md — project identity and universal operating guide
+Tier 1: philosophy.md, constraints.md, — always-read current authority + concise indexes
architecture.md
Routing: README.md — trigger-based table from Tier 1 to Tier 2
Leaves: topics/*.md — self-contained detail (read when triggered)
@@ -19,7 +19,13 @@ Skills: skills/*/SKILL.md — workflows with downward refs to
## Node Roles
-**Non-leaf** (root + Tier 1): State core thesis in 1-3 lines, then index to leaf docs via tables or pointers. No inline explanations, no code examples, no checklists — those belong in leaves.
+**Non-leaf** (root + Tier 1): Keep universal operating rules and current architecture summaries
+concise, then route specialized detail to leaves. `AGENTS.md` may retain universal commands and the
+commit workflow; `constraints.md` may include the minimum rationale needed to make a hard rule
+enforceable; `architecture.md` may retain the current module graph, registries, extension points,
+and short design summaries. Model-specific recipes, deep implementation walkthroughs, new
+chronological fix records, and specialized checklists belong in leaves. Do not add Tier-1 detail
+when a leaf pointer is sufficient.
**Leaf** (`topics/*.md`): Self-contained, concise, essential knowledge. Include code refs, checklists, numbered gotchas. No filler prose, no introductory fluff, no restating what parent docs already say. Format examples: `adapter_conventions.md`, `train_inference_consistency.md`.
diff --git a/.agents/knowledge/topics/adapter_conventions.md b/.agents/knowledge/topics/adapter_conventions.md
index 595e281d5..3be52baa6 100644
--- a/.agents/knowledge/topics/adapter_conventions.md
+++ b/.agents/knowledge/topics/adapter_conventions.md
@@ -241,5 +241,6 @@ LTX2 packs `[video|audio]` into one `(B, Seq, C)` sequence, so it resolves as PA
## Cross-refs
-- UP: `architecture.md` "Adapter Pattern", `constraints.md` #5 #11-12
-- PEER: `train_inference_consistency.md`, `parity_testing.md`, `ff-new-model` Pitfall #6
+- UP: [`constraints.md` #5](../constraints.md#5-adapter-component-runtime-contract), [`constraints.md` #11](../constraints.md#11-basetrainer-execution-contract), [`constraints.md` #12](../constraints.md#12-baseadapter-abstract-methods), [Architecture Adapter Pattern](../architecture.md#adapter-pattern-models)
+- PEER: [Train/Inference Consistency](train_inference_consistency.md), [Parity Testing](parity_testing.md)
+- WORKFLOW: [`ff-new-model`](../../skills/ff-new-model/SKILL.md)
diff --git a/.agents/knowledge/topics/autocast_param_swap.md b/.agents/knowledge/topics/autocast_param_swap.md
index 6daea4265..746192a18 100644
--- a/.agents/knowledge/topics/autocast_param_swap.md
+++ b/.agents/knowledge/topics/autocast_param_swap.md
@@ -35,6 +35,6 @@ Ref/EMA/named snapshots share **one** model and swap in place via `EMAModuleWrap
## Cross-refs
-- `constraints.md` #20a, #20, #10
-- `topics/dtype_precision.md`, `train_inference_consistency.md`
-- `models/abc.py` `use_ref_parameters` / `use_ema_parameters` / `use_named_parameters`; `ema/ema.py` `EMAModuleWrapper`
+- UP: [`constraints.md` #20a](../constraints.md#20a-autocast-weight-cache-must-not-span-a-forward), [`constraints.md` #20](../constraints.md#20-mixed-precision-consistency), [`constraints.md` #10](../constraints.md#10-deepspeed-zero-3-is-unsupported)
+- PEER: [Dtype and Precision](dtype_precision.md), [Train/Inference Consistency](train_inference_consistency.md)
+- CODE: [`BaseAdapter` parameter-swap contexts](../../../src/flow_factory/models/abc.py), [`EMAModuleWrapper`](../../../src/flow_factory/ema/ema.py)
diff --git a/.agents/knowledge/topics/component_runtime.md b/.agents/knowledge/topics/component_runtime.md
index 06729b22d..a93ca6783 100644
--- a/.agents/knowledge/topics/component_runtime.md
+++ b/.agents/knowledge/topics/component_runtime.md
@@ -70,6 +70,25 @@ instance attribute.
- A target-owned composite root may still contain frozen auxiliary siblings. Pseudo runtimes move
only that remainder and exclude every prepared target route.
+## Activation-checkpoint ownership
+
+`configure_checkpointing_backend_plan()` resolves one owner before model loading; direct trainer
+construction repeats the same plan defensively after adapter realization.
+
+| FSDP plan | Activation-checkpoint owner |
+|---|---|
+| FSDP1 with train-level model checkpointing | Model-side owner; a simultaneously enabled backend owner is disabled |
+| FSDP1 + TDM-R1 | Both owners are disabled because reference/snapshot forwards invalidate FSDP1 recomputation |
+| FSDP2 with no train-level model policy | The configured backend policy is retained; otherwise checkpointing stays off |
+| FSDP2 with a full train-level policy | Accelerate/FSDP2 backend wrappers own recomputation; the train-level model policy is disabled |
+| FSDP2 with a selective train-level policy | Rejected because backend wrappers cannot preserve selective model boundaries |
+| FSDP2 with an adapter-owned in-forward capability | The adapter installs block-local boundaries after the FSDP input cast; the duplicate plugin owner is disabled during preparation |
+
+MiniMax H3 Ref2VA currently owns the adapter-specific exception. Its prepared FSDP units also
+receive a replay-time public `unshard()` hook for the affected PyTorch lifecycle, where a second
+checkpoint graph can replay after another graph resharded the same unit. These flags are
+model-specific memory policy and must not be copied to another adapter without evidence.
+
## Fix records
### Repeated-block metadata must survive the distributed bundle boundary
@@ -107,8 +126,9 @@ instance attribute.
- [ ] `adapter.pipeline`, `adapter.scheduler`, and public lifecycle hooks remain compatible.
- [ ] Scheduler names equal `trajectory_component_order`.
- [ ] Distributed preparation still routes through `ModelBundle`/`RoutedComponentProxy`.
+- [ ] Exactly one activation-checkpoint owner remains active for the selected FSDP plan.
## Cross-refs
- UP: [`constraints.md` #5](../constraints.md#5-adapter-component-runtime-contract), [`architecture.md` Component Management](../architecture.md#component-management)
-- PEER: [Structured Trajectory](structured_trajectory.md), [Component Variants](component_variants.md), [Adapter Conventions](adapter_conventions.md)
+- PEER: [Structured Trajectory](structured_trajectory.md), [Component Variants](component_variants.md), [Adapter Conventions](adapter_conventions.md), [MiniMax H3](minimax_h3.md)
diff --git a/.agents/knowledge/topics/dtype_precision.md b/.agents/knowledge/topics/dtype_precision.md
index 0280fdbd9..207d08cbc 100644
--- a/.agents/knowledge/topics/dtype_precision.md
+++ b/.agents/knowledge/topics/dtype_precision.md
@@ -96,8 +96,5 @@ The round-trip ensures that the precision of stored latents matches what trainin
## Cross-refs
-- `constraints.md` #18 (all-rank synchronization — precision errors may manifest differently per rank)
-- `constraints.md` #20 (mixed precision consistency)
-- `topics/autocast_param_swap.md` (#20a)
-- `train_inference_consistency.md` (log_prob mismatch from precision)
-- `topics/timestep_sigma.md` (scheduler math always float32)
+- UP: [`constraints.md` #18](../constraints.md#18-all-rank-synchronization-points), [`constraints.md` #20](../constraints.md#20-mixed-precision-consistency)
+- PEER: [Autocast and Parameter Swaps](autocast_param_swap.md), [Train/Inference Consistency](train_inference_consistency.md), [Timestep and Sigma](timestep_sigma.md)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index b5dd4c067..198877972 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -742,6 +742,5 @@ Based on the fix type, write the fix entry to the appropriate document:
## Cross-refs
-- `constraints.md` (archival target for constraint violations)
-- `architecture.md` (archival target for data-flow misunderstandings)
-- `ff-debug/SKILL.md` Phase 5 (knowledge capture workflow)
+- UP: [Hard Constraints](../constraints.md), [Architecture](../architecture.md)
+- WORKFLOW: [`ff-debug` Phase 5](../../skills/ff-debug/SKILL.md#5-capture-the-fix)
diff --git a/.agents/knowledge/topics/minimax_h3.md b/.agents/knowledge/topics/minimax_h3.md
index d4f3048da..bde33e60c 100644
--- a/.agents/knowledge/topics/minimax_h3.md
+++ b/.agents/knowledge/topics/minimax_h3.md
@@ -143,11 +143,13 @@ element-weighted reducer.
## Verification boundary
All workflows have pinned API/schema/no-weight verification and local offline codec/forward
-coverage. T2VA additionally completed
-real-weight LoRA rollout, decode, reward, replay, backward, checkpoint, and resume tests on one
-GPU and with FSDP2 on 16 GPUs. The native-resolution path completed initialization, checkpoint,
-decode, and evaluation. FL2VA/Ref2VA SFT and offline DPO still require the documented real-weight
-GPU matrix. Do not claim long-run reward improvement, convergence, or numerical parity.
+coverage. T2VA additionally completed real-weight LoRA rollout, decode, reward, replay, backward,
+checkpoint, and resume tests on one GPU and with FSDP2 on 16 GPUs. The native-resolution path
+completed initialization, checkpoint, decode, and evaluation. The
+[PR #220 matrix](../../../guidance/gpu_validation.md#pr-220-result) then completed all 36 H3
+real-weight smoke cells: T2VA, FL2VA, and Ref2VA across DDP/DeepSpeed ZeRO-2/FSDP2 and
+GRPO/SFT/offline DPO/TDM. The FL2VA first-plus-last SFT/offline-DPO variant gate also passed. Do
+not claim long-run reward improvement, convergence, quality parity, or numerical parity.
## Upgrade checklist
@@ -157,9 +159,10 @@ GPU matrix. Do not claim long-run reward improvement, convergence, or numerical
- [ ] Run H3 scheduler/runtime/registry/reference tests in the pinned environment.
- [ ] Parse all H3 examples through `Arguments.load_from_yaml`.
- [ ] Run the T2VA output-codec and common SFT/offline-DPO structured-state tests.
-- [ ] Rerun the documented T2VA real-weight smoke before changing support or memory claims.
+- [ ] Rerun the affected H3 cells in the documented real-weight matrix before changing support or
+ memory claims.
## Cross-refs
-- UP: [`constraints.md` #5](../constraints.md#5-adapter-component-runtime-contract), [`constraints.md` #14](../constraints.md#14-sample-dataclass-hierarchy), [`architecture.md` MiniMax H3](../architecture.md#minimax-h3)
+- UP: [`constraints.md` #5](../constraints.md#5-adapter-component-runtime-contract), [`constraints.md` #14](../constraints.md#14-sample-dataclass-hierarchy), [`architecture.md` registered components](../architecture.md#registered-components)
- PEER: [Component Runtime](component_runtime.md), [Structured Trajectory](structured_trajectory.md), [Parity Testing](parity_testing.md)
diff --git a/.agents/knowledge/topics/parity_testing.md b/.agents/knowledge/topics/parity_testing.md
index a14b8d1e4..f639092dd 100644
--- a/.agents/knowledge/topics/parity_testing.md
+++ b/.agents/knowledge/topics/parity_testing.md
@@ -58,5 +58,5 @@ def compare_tensors(name: str, a: torch.Tensor, b: torch.Tensor, atol: float = 1
## Cross-refs
-- `adapter_conventions.md` (inference/forward identity, upstream alignment rules)
-- `dtype_precision.md` (tensor dtype for comparison tolerance)
+- UP: [`constraints.md` #12](../constraints.md#12-baseadapter-abstract-methods), [Architecture Adapter Pattern](../architecture.md#adapter-pattern-models)
+- PEER: [Adapter Conventions](adapter_conventions.md), [Dtype and Precision](dtype_precision.md)
diff --git a/.agents/knowledge/topics/sample_lifecycle.md b/.agents/knowledge/topics/sample_lifecycle.md
index c4e59c456..4a0f4a3d8 100644
--- a/.agents/knowledge/topics/sample_lifecycle.md
+++ b/.agents/knowledge/topics/sample_lifecycle.md
@@ -102,11 +102,21 @@ Effect of the offload pipeline: `sample.to('cpu')` and `sample.to(device)` both
If a future custom adapter stores large GPU tensors in `extra_kwargs`, either handle them adapter-side or refactor `BaseSample.to` to delegate to `move_tensors_to_device(value, device, max_depth=1)` in an independent PR (note: that refactor will start moving `extra_kwargs['advantage']` together with the sample, which is benign for the current data flow but is a contract change).
+## Partial-gather reconstruction
+
+`gather_samples()` may transport only fields requested by a distributed consumer, then reconstruct
+the original concrete `BaseSample` subclass. A field that is required by the concrete class's
+`__post_init__`, identity normalization, or constructor invariant must therefore be included in the
+class-level `reconstruction_required_fields`, even when the reward or trainer does not consume it.
+Always union with the inherited set; `OrderedReferenceConditionSample` adds
+`reference_manifest` this way.
+
+This contract is independent of both `_shared_fields` (batch collation) and a reward model's
+`required_fields` (distributed groupwise consumer data). Add a regression that gathers a narrower
+field selection and proves the reconstructed concrete sample still carries every required field.
+
## Cross-refs
-- `constraints.md` #11 (BaseTrainer hook order: `sample()` → `prepare_feedback()` → `optimize()`)
-- `constraints.md` #14 (BaseSample dataclass hierarchy and `_shared_fields`)
-- `constraints.md` #15 + `.cursor/rules/examples-yaml-sync.mdc` (the three-tier strategy is an intentional deviation)
-- `topics/train_inference_consistency.md` item #4 (EMA swap without restore — preserved by per-batch interleave)
-- `topics/dtype_precision.md` (device-move never changes dtype; orthogonal to autocast)
-- `topics/autocast_param_swap.md` (#20a)
+- UP: [`constraints.md` #11](../constraints.md#11-basetrainer-execution-contract), [`constraints.md` #14](../constraints.md#14-sample-dataclass-hierarchy), [`constraints.md` #15](../constraints.md#15-pydantic-hparams-synchronization)
+- PEER: [Train/Inference Consistency](train_inference_consistency.md), [Dtype and Precision](dtype_precision.md), [Autocast and Parameter Swaps](autocast_param_swap.md)
+- RULE: [Examples YAML Synchronization](../../../.cursor/rules/examples-yaml-sync.mdc)
diff --git a/.agents/knowledge/topics/samplers.md b/.agents/knowledge/topics/samplers.md
index c53b0e546..7e341de18 100644
--- a/.agents/knowledge/topics/samplers.md
+++ b/.agents/knowledge/topics/samplers.md
@@ -332,6 +332,4 @@ data:
## Cross-refs
-- `constraints.md` #9, #9a (accelerator prepare scope, sampler geometric constraints)
-- `architecture.md` "Execution Pipelines" (generation acquisition)
-- `architecture.md` "Advantage Computation" (communication path depends on sampler type)
+- UP: [`constraints.md` #9](../constraints.md#9-accelerator-prepare-scope), [`constraints.md` #9a](../constraints.md#9a-sampler-geometric-constraints), [Architecture Execution Pipelines](../architecture.md#execution-pipelines), [Architecture Advantage Computation](../architecture.md#advantage-computation)
diff --git a/.agents/knowledge/topics/timestep_sigma.md b/.agents/knowledge/topics/timestep_sigma.md
index 8abbe5544..4a87b790f 100644
--- a/.agents/knowledge/topics/timestep_sigma.md
+++ b/.agents/knowledge/topics/timestep_sigma.md
@@ -44,6 +44,5 @@ Throughout the codebase, two related but distinct scales are used for time:
## Cross-refs
-- `constraints.md` #7 (coupled/decoupled paradigm — affects which timestep sampling is valid)
-- `topics/train_inference_consistency.md` (same `t` must produce same output in rollout vs training)
-- `topics/adapter_conventions.md` (adapter encapsulates timestep-to-model conversion)
+- UP: [`constraints.md` #7](../constraints.md#7-coupled-vs-decoupled-paradigm), [Architecture Timestep and Sigma Convention](../architecture.md#timestep--sigma-convention)
+- PEER: [Train/Inference Consistency](train_inference_consistency.md), [Adapter Conventions](adapter_conventions.md)
diff --git a/.agents/knowledge/topics/train_inference_consistency.md b/.agents/knowledge/topics/train_inference_consistency.md
index fda66c548..4894632be 100644
--- a/.agents/knowledge/topics/train_inference_consistency.md
+++ b/.agents/knowledge/topics/train_inference_consistency.md
@@ -67,6 +67,5 @@ If rollout and training `forward()` diverge, `ratio` deviates from 1.0 at epoch
## Cross-refs
-- `constraints.md` #7 (coupled/decoupled paradigm)
-- `dtype_precision.md` (precision boundaries, cast_latents)
-- `adapter_conventions.md` (inference/forward identity rule)
+- UP: [`constraints.md` #7](../constraints.md#7-coupled-vs-decoupled-paradigm), [Architecture Execution Pipelines](../architecture.md#execution-pipelines)
+- PEER: [Dtype and Precision](dtype_precision.md), [Adapter Conventions](adapter_conventions.md)
diff --git a/.cursor/rules/agents-docs-maintenance.mdc b/.cursor/rules/agents-docs-maintenance.mdc
index 1aeb85e9c..19907717a 100644
--- a/.cursor/rules/agents-docs-maintenance.mdc
+++ b/.cursor/rules/agents-docs-maintenance.mdc
@@ -10,7 +10,7 @@ The `.agents/knowledge/` system uses layered design with cross-references. Follo
## Layered structure
-- **Non-leaf** (root, Tier 1): State thesis in 1-3 lines, then index to leaf docs via tables/pointers. No inline explanations, code examples, or checklists — those live in leaves.
+- **Non-leaf** (root, Tier 1): Keep universal operating rules and current architecture summaries concise, then index specialized detail via tables/pointers. `AGENTS.md` may retain universal commands and the commit workflow; `constraints.md` may include the minimum rationale needed to enforce a hard rule; `architecture.md` may retain the current module graph, registries, extension points, and short design summaries. Model-specific recipes, deep implementation walkthroughs, new chronological fix records, and specialized checklists live in leaves.
- **Leaf** (`topics/*.md`): Self-contained, concise detail. Include code refs, checklists, gotchas. No filler prose, no restating parent content.
- No duplication across layers. If detail exists in a leaf, the parent points to it.
diff --git a/AGENTS.md b/AGENTS.md
index ce79bf0e8..2ebfbf74d 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -5,7 +5,7 @@
Flow-Factory is a unified **online and offline fine-tuning framework** for diffusion/flow-matching models. It provides a modular architecture where trainers, model adapters, data acquisition, and reward models are independently extensible through typed contracts and registries.
- **Algorithms**: SFT, offline DPO, online DPO, GRPO, GRPO-Guard, DPPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1
-- **Models**: FLUX.1 (+Kontext), FLUX.2 (+Klein), SD3.5, Qwen-Image (+Edit-Plus), Z-Image, Wan2 (T2V/I2V), LTX2 (T2AV/I2AV), Bagel, SenseNova-U1 (1.0/1.5; T2I + ordered multi-reference I2I)
+- **Models**: FLUX.1 (+Kontext), FLUX.2 (+Klein), SD3.5, Qwen-Image (+Edit-Plus), Z-Image, Wan2 (T2V/I2V), LTX2 (T2AV/I2AV), MiniMax H3 (T2VA/FL2VA/Ref2VA), Bagel, SenseNova-U1 (1.0/1.5; T2I + ordered multi-reference I2I)
- **Rewards**: PickScore (+Rank), CLIP, CLAP, ImageBind, OCR, GenEval/GenEval2, HPSv2, VLM-Evaluate, rational-rewards, and custom rewards
- **Python**: >=3.10 | **PyTorch**: >=2.6.0 | **License**: Apache-2.0
@@ -23,8 +23,8 @@ On session start, read **Tier 1** (see `.agents/knowledge/README.md`):
## Core Operating Principles
1. **Constraints first** — Read `constraints.md` + `architecture.md` before changes; search codebase before attempting fixes.
-2. **Cross-component awareness** — Changes to `abc.py` affect ALL subclasses; verify across algorithms (GRPO + NFT/AWM).
-3. **Plan before implement** — Multi-file tasks -> TodoWrite. Plan must state which skills apply.
+2. **Cross-component awareness** — Changes to base classes or typed contracts affect every registry-resolved implementation; verify the affected coupled-reward, decoupled-reward, no-feedback, and dataset-acquisition paths.
+3. **Plan before implement** — Multi-file tasks require an explicit task plan using the agent's supported planning mechanism. The plan must state which skills apply.
4. **Challenge first, execute second** — Spot logic flaws or simpler alternatives? Raise before executing.
5. **Escalation** — After three failed approaches, document findings and request review.
6. **Fix capture** — After every bug fix, generate summary per `.agents/knowledge/topics/fix_patterns.md` template.
From 002df5f2b1d3f08087f0ebfd7bdbb864fcf9ad9e Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 07:53:10 +0800
Subject: [PATCH 71/76] [agents] docs: upgrade framework development skills
---
.agents/skills/ff-debug/SKILL.md | 263 ++++++++-------
.agents/skills/ff-develop/SKILL.md | 253 +++++++++-----
.agents/skills/ff-new-algorithm/SKILL.md | 413 +++++++++++------------
.agents/skills/ff-new-model/SKILL.md | 295 +++++++++-------
.agents/skills/ff-new-reward/SKILL.md | 174 +++++-----
.agents/skills/ff-review/SKILL.md | 270 +++++++++------
6 files changed, 941 insertions(+), 727 deletions(-)
diff --git a/.agents/skills/ff-debug/SKILL.md b/.agents/skills/ff-debug/SKILL.md
index b98ef698b..bc5d3f9cc 100644
--- a/.agents/skills/ff-debug/SKILL.md
+++ b/.agents/skills/ff-debug/SKILL.md
@@ -1,128 +1,153 @@
---
name: ff-debug
-description: "Bug fixing and debugging for ANY error, crash, loss divergence, gradient explosion, distributed hang, NaN, or unexpected behavior. Covers quick fixes and full protocol with 5-phase investigation. Trigger: 'fix bug', 'fix error', 'broken', 'crash', 'doesn't work', 'fails with', 'loss NaN', 'training hangs', 'OOM'."
+description: "Debug Flow-Factory crashes, hangs, OOMs, numerical failures, finite-loader errors, component routing, distributed preparation, optimizer roles, and checkpoint/resume mismatches. Use for bug fixing or unexpected training behavior."
---
# Debug Workflow
-## Related Topics (read for numerical / consistency issues)
-
-- NaN, loss divergence, wrong gradients -> `topics/train_inference_consistency.md`
-- Dtype mismatch, overflow, precision -> `topics/dtype_precision.md`
-- Frozen/flat loss or KL ≈ 0 -> `topics/autocast_param_swap.md` (#20a)
-
-## Two Pathways
-
-### Quick Path (obvious root cause)
-
-Use when: Error message clearly points to the issue (typo, missing import, wrong type).
-
-1. Reproduce the error
-2. Check `.agents/knowledge/constraints.md` for relevant constraints
-3. Write targeted fix
-4. Verify with test
-5. Run `/ff-review`, commit
-
-If not resolved in 15 min -> switch to Full Protocol.
-
-### Full Protocol (complex issues)
-
-Use when:
-- Distributed training bugs (deadlocks, rank mismatches)
-- Numerical issues (NaN, loss divergence, wrong gradients)
-- Silent failures (training runs but produces garbage)
-- Multiple failed fix attempts
-
-## Full Protocol — Five Phases
-
-### Phase 1: Root Cause Investigation
-
-1. **Read complete error messages** — Full stack traces matter, don't skim
-2. **Consult constraints** — Check `.agents/knowledge/constraints.md`
-3. **Reproduce consistently** — Isolate the exact trigger condition
-4. **Trace execution path** — Follow through the 6-stage pipeline
-5. **Check recent changes** — `git log --oneline -10` — what changed recently?
-
-#### Distributed-Specific Checklist
-- Does the error appear on all ranks or just one?
-- Is `accelerator.wait_for_everyone()` missing before the failure point?
-- Are frozen components synchronized across ranks? (Constraint #19)
-- Is ZeRO-3 being used? (Constraint #10 — unsupported)
-
-### Phase 2: Pattern Analysis
-
-1. **Find working examples** — Compare with a similar model/algorithm that works
-2. **Diff analysis** — What's different between working and broken paths? Compare **completely** — diff line by line, not skim. Include config YAML and environment vars.
-3. **Isolate variables** — Change one thing at a time
-4. **Check dependencies** — Different diffusers version? Different PyTorch version?
-
-### Phase 3: Hypothesis Testing
-
-1. **One hypothesis per iteration** — Formulate a single falsifiable hypothesis
-2. **Minimal test case** — Reproduce with smallest possible config
-3. **Low confidence (<80%)?** — Add debug logging before applying fix
-
-**Red flags — STOP and restart from Phase 1:**
-- "Let me just try changing X and see what happens"
-- "Quick fix for now, clean up later"
-- "It probably works, let me move on"
-
-**Verification gate** — before acting on a conclusion, check:
-- Does the evidence actually support this cause, or just correlate?
-- Could a different root cause produce the same symptoms?
-- What observation would disprove this hypothesis? Have you looked for it?
-
-### Phase 4: Fix Implementation
-
-1. **Write failing test first** (if possible)
-2. **Implement targeted fix** — Only fix the bug, don't refactor
-3. **Check cross-algorithm impact** — Does this fix break GRPO? NFT? AWM?
-4. **Check cross-model impact** — Test with at least two model adapters
-5. Before committing: run `/ff-review` skill.
-
-### Phase 5: Knowledge Capture
-
-After fix is verified:
-- Update `constraints.md` if a new constraint was discovered
-- Add regression test if applicable
-- Document the root cause in the commit message
-- Follow fix archival process in `topics/fix_patterns.md`
+## Load the Relevant Contracts
+
+Always read Tier 1: `../../knowledge/constraints.md`, `../../knowledge/architecture.md`, and
+`../../knowledge/philosophy.md`. Then route by symptom:
+
+| Symptom or area | Also read |
+|---|---|
+| Wrong gradients, ratio drift, bad output | `../../knowledge/topics/train_inference_consistency.md`, `../../knowledge/topics/parity_testing.md` |
+| Dtype mismatch, overflow, flat loss or KL | `../../knowledge/topics/dtype_precision.md`, `../../knowledge/topics/autocast_param_swap.md` |
+| Missing component, wrong device, lazy load, wrap/OOM | `../../knowledge/topics/component_runtime.md` |
+| Multi-component rollout or replay | `../../knowledge/topics/structured_trajectory.md` |
+| Variant, role cadence, optimizer group, Muon, role checkpoint | `../../knowledge/topics/component_variants.md` |
+| Finite dataset, target encoding, SFT/offline DPO | `../../../guidance/workflow.md`, `../../../guidance/datasets.md` |
+
+## Classify the Execution Path First
+
+Resolve the trainer and its algorithm-specific `TrainingArguments`; their immutable
+`ExecutionContract` values must be equal.
+
+| Composition | Expected path |
+|---|---|
+| `generation + runtime_reward` | `sample()` -> feedback/reward -> advantage -> `optimize()` |
+| `generation + none` | `sample()` -> `optimize()`; no training reward/advantage stage |
+| `dataset + none` | Exhaust one official finite `DistributedSampler` loader through `optimize_batch()` |
+
+Do not infer execution mode from batch fields. Track `optimizer_step` independently from
+`rollout_iteration` or `data_epoch`; a dataset epoch advances only after clean loader exhaustion.
+
+## Quick Path
+
+Use when the failure is deterministic and the stack trace identifies one local contract breach.
+
+1. Reproduce it with the smallest representative test or config.
+2. Trace the owning boundary and relevant constraint.
+3. Add a regression that fails for the same reason.
+4. Apply the narrow fix and run affected contract tests.
+5. Run `/ff-review` before committing.
+
+If the cause is uncertain, distributed, numerical, or survives one focused attempt, use the full
+protocol.
+
+## Full Protocol
+
+### 1. Establish the Failure Boundary
+
+- Read every rank's complete traceback and first causal error.
+- Record the resolved trainer, adapter, execution contract, model I/O contract, backend, optimizer
+ types, finetune type, dtype policy, and checkpoint mode.
+- Compare against a working path one variable at a time, including YAML and backend config.
+- Identify when failure occurs: preflight, native component load, preprocessing, bundle prepare,
+ proxy-routed forward, acquisition, optimizer step, save, or resume.
+- Check recent changes with a focused file/commit diff; do not assume temporal correlation is cause.
+
+### 2. Check Ownership Invariants
+
+#### Dataset acquisition
+
+- Uses PyTorch's official `DistributedSampler`, even at one rank, and calls
+ `set_epoch(data_epoch)`.
+- Uses explicit positive `gradient_accumulation_steps`; rank-local batch count closes every
+ accumulation window without an implicit flush.
+- Caches prompt/input conditions only. Target, chosen, and rejected media is decoded and encoded on
+ demand.
+- Calls `prepare_condition_state()` once per batch. Offline DPO reuses that object, schedule, noise,
+ and one reference scope across both arms.
+- Applies the adapter's complete `offline_training_forward_overrides`, not rollout CFG semantics.
+
+#### Component runtime and loading
+
+- Resolve membership with `has_component`, `get_component`, or `_require_component`, never
+ `hasattr(adapter, name)`.
+- Keep canonical components, prepared/replacement overrides, declared specs, materialized modules,
+ optional `None`, and pseudo aliases distinct.
+- `materialize_components(None)` means already-materialized modules, not all lazy declarations.
+- Trace logical names to physical roots through `ModelLoadCoordinator`; adapters/trainers must not
+ reproduce FSDP loading state or broadcast weights directly.
+- All target and frozen-but-shardable routes enter one `ModelBundle` with one optimizer prepare
+ root. Adapter forwards after prepare route through `RoutedComponentProxy`.
+- For FSDP OOM, inspect bundle-exposed `_no_split_modules`/`_repeated_blocks`, wrap classes, adapter
+ memory capabilities, checkpoint replay, unshard stream, and backward prefetch before shrinking
+ the workload.
+
+#### Variants, optimizers, and distributed plans
+
+- A temporal reference/EMA/snapshot is not a live component variant.
+- Variants are declared before prepare; role parameter and optimizer-group ownership is disjoint and
+ exhaustive.
+- Do not assume one group per role: Muon uses matrices plus an optional AdamW remainder inside one
+ `CompositeOptimizer` root.
+- Reject ZeRO-3. Reject Muon before model loading when its PyTorch API is absent or the backend is
+ DeepSpeed/FSDP1. Multi-role DeepSpeed requires ZeRO-1/2. Multi-role FSDP2 requires
+ `use_orig_params=True`; after prepare, registry and optimizer references must point to the
+ DTensor-backed parameters owned by the prepared model root.
+- Activation checkpointing has one owner. FSDP2 full policy is normalized to backend ownership;
+ selective model checkpointing is rejected. Inspect adapter-owned in-forward checkpointing only
+ when the adapter explicitly opts in.
+
+#### Checkpoint and exact resume
+
+- Model-only export scope and resumable state scope are different. Resumable multi-role saves include
+ training-only roles, role counters, optimizer ownership, and variant snapshots.
+- Validate variant/runtime metadata before Accelerate mutates prepared state.
+- Exact identity covers changed objective, model/backend/optimizer semantics, realized data order,
+ and replayed evaluation configuration; cadence and resume location remain operational controls.
+- Preserve all-rank phase symmetry and atomic publication ordering.
+
+#### Numerical and reward paths
+
+- Wrap each policy/reference/EMA forward in its own autocast region when weights can change in place.
+- Consume trajectories only through adapter bridge methods and authoritative component order.
+- A partial cross-rank gather must union the concrete sample class's
+ `reconstruction_required_fields` before reconstruction. This is independent of reward
+ `required_fields` and collation `_shared_fields`.
+- Pointwise rewards return one finite value per actual input chunk, which may be smaller than
+ `batch_size`; groupwise rewards preserve complete `unique_id` order.
+- Per-dataset reward applicability is framework-owned; model NaN/Inf is an error, not a routing
+ sentinel.
+
+### 3. Test One Falsifiable Hypothesis
+
+State the proposed cause, the observation that would disprove it, and the smallest experiment that
+separates it from alternatives. Add instrumentation when confidence is below 80%. Avoid speculative
+fallbacks or several behavioral changes in one experiment.
+
+### 4. Fix and Verify in Scope
+
+- Write the regression first when practical.
+- Fix the authoritative owner rather than patching downstream consumers.
+- Verify the narrow unit contract, then affected compositions:
+ - execution kernel: GRPO, a reward-free generation trainer, and SFT/offline DPO as applicable;
+ - adapter/trajectory: legacy single-component and structured multimodal paths;
+ - runtime/loading: the affected classic/modular/pseudo runtime and supported backends;
+ - variants/optimizer: single-role AdamW plus multi-role/Muon cases when touched;
+ - checkpoint: model-only round trip and exact-state resume when touched.
+- Test at least two adapters only when the changed abstraction is shared across adapters.
+
+### 5. Capture the Fix
+
+Follow `../../knowledge/topics/fix_patterns.md`: record symptom, root cause, fix, lesson, related
+constraint, test evidence, and commit. Update Tier 1 only when the fix establishes a durable
+cross-module invariant.
## Three-Strike Rule
-If the same approach fails three times:
-1. **HALT** all fix attempts
-2. Question whether the underlying approach/architecture is wrong
-3. Step back and re-examine: are you solving the right problem?
-4. Report to user with analysis before continuing
-
-## Common Issue Categories
-
-### Training Loop Issues
-- [ ] Stage ordering violated? (Constraint #6)
-- [ ] Coupled/decoupled paradigm mismatch? (Constraint #7)
-- [ ] Component not on correct device? (Constraint #8)
-- [ ] Dataloader incorrectly prepared via accelerator? (Constraint #9)
-
-### Model Adapter Issues
-- [ ] `load_pipeline()` returning wrong type? (Constraint #5)
-- [ ] `target_module_map` mapping incorrect components?
-- [ ] `_shared_fields` causing data corruption? (Constraint #14)
-- [ ] Preprocessing modules not offloaded after Stage 1?
-
-### Reward Issues
-- [ ] Pointwise/Groupwise confusion? (Constraint #13)
-- [ ] Wrong reward shape returned?
-- [ ] `required_fields` not set correctly?
-- [ ] Device mismatch between reward model and generated samples?
-
-### Configuration Issues
-- [ ] YAML key doesn't match the dataclass field name? (Constraint #17)
-- [ ] Algorithm-specific args using wrong subclass? (Constraint #16)
-- [ ] Registry key doesn't match? (Constraint #1)
-
-### Distributed Issues
-- [ ] Missing synchronization barrier? (Constraint #18)
-- [ ] FSDP frozen components uninitialized on Rank > 0? (Constraint #19)
-- [ ] Mixed precision casting order incorrect? (Constraint #20) — see also `topics/dtype_precision.md` for precision diagnosis checklist
-- [ ] Using ZeRO-3? (Constraint #10 — not supported)
+After three failed approaches to the same cause, stop patching, document evidence and rejected
+hypotheses, reassess the ownership model, and request review before continuing.
diff --git a/.agents/skills/ff-develop/SKILL.md b/.agents/skills/ff-develop/SKILL.md
index 59239311a..f9e2cc2f5 100644
--- a/.agents/skills/ff-develop/SKILL.md
+++ b/.agents/skills/ff-develop/SKILL.md
@@ -1,98 +1,167 @@
---
name: ff-develop
-description: "Feature development with cross-module impact analysis. Covers trainer hierarchy, model adapters, reward pipeline, config system, sample dataclasses, and distributed training paths. Trigger: 'add feature', 'implement', 'refactor', 'reorganize', 'new capability'."
+description: "Develop or refactor Flow-Factory features with cross-module impact analysis across execution contracts, typed data/model I/O, component loading, prepared bundles, rewards, variants, optimizers, distributed backends, and checkpoints."
---
# Feature Development Workflow
-## Related Topics (read if your change touches these areas)
-
-- Adapter changes -> `topics/adapter_conventions.md`
-- Trainer/scheduler changes -> `topics/train_inference_consistency.md`
-- Trainer `optimize()` loop, autocast scope, ref/EMA/named param swaps -> `topics/autocast_param_swap.md`
-- Precision changes -> `topics/dtype_precision.md`
-
-## Impact Analysis Checklist
-
-Before implementing features or refactoring, analyze impacts across these areas:
-
-### 1. Trainer Hierarchy (`constraints.md` #11)
-- Changes to `BaseTrainer` affect all 9 concrete trainers (grpo, grpo-guard, dppo, nft, awm, dgpo, dpo, crd, diffusion-opd); changed abstract methods must be implemented on every one
-- Changes to `AdvantageProcessor` affect all reward-based trainers (`architecture.md` "Advantage Computation"; `diffusion-opd` skips it)
-- Check: Does your change alter `_initialization()`, `_init_reward_model()`, or `_init_dataloader()`?
-
-### 2. Model Adapter Hierarchy (`constraints.md` #12)
-- Changes to `BaseAdapter` affect ALL model adapters
-- Check: Does your change modify component management, LoRA logic, or mode switching?
-- **Adding a new modality** (e.g. audio): prefer non-abstract no-op default + opt-in override (R7 pattern). Don't add `@abstractmethod` to a new encoder; that forces stub edits on every existing concrete adapter. The 4 abstract methods (`load_pipeline`, `decode_latents`, `forward`, `inference`) are intentionally minimal — encoders are opt-in by modality.
-
-### 3. Reward Pipeline (`constraints.md` #13)
-- Changes to `BaseRewardModel` or `RewardProcessor` affect all reward models
-- Check: Does your change alter the Pointwise/Groupwise dispatch?
-
-### 4. Configuration System (`constraints.md` #15–17)
-- Check: Did you rename, remove, or **add** fields? ALL configs in `examples/` must be updated
-
-### 5. Sample Dataclasses (`constraints.md` #14)
-- Changes to `BaseSample` or its subclasses affect data flow through all 6 stages
-- Check: Did you change `_shared_fields` or add new fields?
-
-### 6. Distributed Training Paths (`constraints.md` #9, #18–20)
-- Changes may behave differently under Accelerate vs DeepSpeed
-- Check: Does your change involve `accelerator.prepare()`, gradient accumulation, or model sharding?
-
-## Refactoring Safety Rules
-
-1. **Establish baseline** — Run tests before making changes
-2. **One at a time** — ONE structural change → update ALL callers → verify → commit
-3. **Never combine** — Don't combine multiple refactoring steps in one commit
-
-## Workflow Steps
-
-1. **Understand scope**
- - Read relevant `abc.py` base classes
- - Identify all affected subclasses and callers
- - Read related `guidance/` docs
-
-2. **Plan changes**
- - List all files that need modification
- - Document expected behavior changes
- - Identify test scenarios
-
-3. **Implement methodically**
- - Make ONE change at a time
- - Update ALL callers/subclasses
- - Run tests after each change
-
-4. **Cross-algorithm verification**
- - Test with GRPO (coupled paradigm; also covers GRPO-Guard / DPPO variants)
- - Test with NFT or AWM (decoupled paradigm; also DGPO / CRD / DPO)
- - If the change touches the sample/optimize path, also test `diffusion-opd` (distillation; no reward/advantage stage)
- - Verify with at least two different model adapters
-
-## Documentation
-
-Before committing, check if the change requires documentation updates:
-
-- **New/changed API** -> update relevant `guidance/` doc
-- **New/changed config fields** -> update ALL example configs in `examples/`
-- **Architecture change** -> update `.agents/knowledge/architecture.md`
-- **New constraint discovered** -> add to `.agents/knowledge/constraints.md`
-- **Bug fix experience?** -> follow `.agents/knowledge/topics/fix_patterns.md` archival process
-
-## When to Delegate
-
-- **Adding a new model** → `/ff-new-model`
-- **Adding a new reward** → `/ff-new-reward`
-- **Adding a new algorithm** → `/ff-new-algorithm`
-- **Debugging a bug** → `/ff-debug`
-- **Pre-commit review** → `/ff-review`
-
-## Pre-Commit Checks
-
-- [ ] Impact analysis completed for all 6 areas
-- [ ] All callers/subclasses updated
-- [ ] Tests pass
-- [ ] Code formatted with Black and isort
-- [ ] YAML configs in `examples/` updated: new fields added, renamed fields updated, removed fields cleaned up
-- [ ] License header present on new files
+## Read by Change Area
+
+Always read Tier 1. Add only the topic docs relevant to the change:
+
+| Change area | Read |
+|---|---|
+| Adapter semantics or parity | `../../knowledge/topics/adapter_conventions.md`, `../../knowledge/topics/parity_testing.md` |
+| Trainer forward/replay | `../../knowledge/topics/train_inference_consistency.md`, `../../knowledge/topics/autocast_param_swap.md` |
+| Dtype or mixed precision | `../../knowledge/topics/dtype_precision.md` |
+| Component discovery, lifecycle, load, prepare | `../../knowledge/topics/component_runtime.md` |
+| Multi-component rollout/replay | `../../knowledge/topics/structured_trajectory.md` |
+| Variants, roles, optimizers, role checkpoints | `../../knowledge/topics/component_variants.md` |
+| Dataset acquisition/offline objective | `../../../guidance/workflow.md`, `../../../guidance/datasets.md` |
+| Acceleration plugin | `../../../guidance/acceleration.md` |
+
+## Plan Around Ownership
+
+Before editing, state:
+
+- the authoritative owner and public contract being changed;
+- affected execution compositions (`generation + runtime_reward`, `generation + none`,
+ `dataset + none`);
+- registries, arguments, examples, docs, and checkpoint compatibility surfaces;
+- affected component runtime types and distributed backends;
+- unit, integration, and GPU evidence required.
+
+Prefer one coherent invariant per commit. A contract change may need code, tests, docs, examples, and
+knowledge updates atomically; do not split it merely to minimize file count.
+
+## Impact Analysis
+
+### 1. Execution Kernel and Trainers
+
+- Derive current trainers from both trainer and TrainingArguments registries; do not maintain a
+ hard-coded subclass list.
+- Trainer and algorithm-specific arguments declare the same immutable `ExecutionContract`.
+- `BaseTrainer.start()` owns acquisition dispatch, periodic boundaries, progress, and cycle hooks.
+ Generation implements `optimize(samples)`; dataset acquisition implements
+ `optimize_batch(batch)`.
+- Keep `optimizer_step`, `rollout_iteration`, and `data_epoch` independent. Exact runtime identity
+ must include any changed objective, data, backend, optimizer, or replayed evaluation semantics.
+- New trainers inherit directly from `BaseTrainer`. Existing sanctioned strict extensions are
+ GRPO-Guard/DPPO from GRPO and TDM-R1 from TDM; another trainer-to-trainer extension requires an
+ explicit architectural justification.
+
+### 2. Data Acquisition and Schemas
+
+- Generation uses framework grouped loaders and adapter inference. Dataset acquisition uses an
+ official finite `DistributedSampler`, calls `set_epoch(data_epoch)`, and does not prepare its
+ training loader through Accelerator.
+- Dataset source weights stay `1`; gradient accumulation is explicit and every rank-local epoch
+ closes cleanly without an implicit partial-window flush.
+- Strict V2 supervision records and collators remain model-neutral. Cache prompt/input conditions,
+ never target/chosen/rejected pixels or latents.
+- Evaluation remains generation-based, including for SFT and offline DPO.
+
+### 3. Adapter and Model I/O Contracts
+
+- Preserve the four abstract adapter methods and opt-in modality encoders. New sample classes use
+ the task-level hierarchy, not another model-specific class.
+- Assess the immutable `PipelineIOContract`, checkpoint-realized specialization, condition-state
+ preparer, output-state codec, exact geometry validation, offline forward overrides, and objective
+ reduction whenever offline capability changes.
+- Public boundary-owning adapter wrappers validate shared semantics; extend their protected hooks
+ instead of overriding them.
+- Structured multimodal trajectories use adapter-owned component order and bridge APIs; trainer code
+ never branches on legacy vs structured storage.
+- Fields required by concrete sample reconstruction belong in inherited
+ `reconstruction_required_fields`; do not conflate this with reward `required_fields` or collator
+ `_shared_fields`.
+
+### 4. Component Runtime and Loading
+
+- Preserve canonical, override, declared, materialized, optional, and alias boundaries across
+ `ClassicPipelineRuntime`, `ModularPipelineRuntime`, and `PseudoPipelineRuntime`.
+- Resolve component membership through the runtime, not adapter attributes. Lazy materialization
+ names required components explicitly.
+- Public adapter lifecycle hooks remain the trainer-facing seam. `ModelLoadCoordinator` compiles
+ logical declarations into exactly-once physical-root ownership; trainer/model code does not
+ duplicate backend loading policy.
+- Loading dtype is applied during native materialization; frozen/trainable dtype policy also reaches
+ components materialized later.
+
+### 5. Prepared Model Ownership
+
+- All target components, frozen-but-shardable siblings, and live variant routes form one
+ `ModelBundle`; prepare it with one optimizer root.
+- After prepare, canonical adapter forwards route through `RoutedComponentProxy`. Preserve stable
+ member names and `_no_split_modules`/`_repeated_blocks` metadata needed by FSDP policy discovery.
+- Do not manually move, wrap, or offload a prepared route or individual trainable variant.
+- Checkpoint save/load iterates trainable component ownership symmetrically; frozen bundle members
+ are not given nonexistent per-component artifacts.
+
+### 6. Variants, Roles, Optimizers, and Resume
+
+- Spatial live trainable copies use `ComponentVariantRegistry`; temporal references, EMA, and old
+ policies use named/ref snapshots.
+- Algorithms own role names and update cadence. Variants are declared before prepare; role
+ parameter and optimizer-group ownership is disjoint and exhaustive.
+- The top-level `optimizers:` list has one entry per trainable role. Do not assume one group per
+ role: Muon splits matrices from an optional AdamW remainder and still participates in one
+ `CompositeOptimizer` root.
+- Muon is supported only when `torch.optim.Muon` exists and the backend is DDP or FSDP2. Multi-role
+ DeepSpeed requires ZeRO-1/2. Multi-role FSDP2 requires `use_orig_params=True`; the registry is
+ rebound after prepare and optimizer references must point to the prepared root's DTensor-backed
+ parameters.
+- Resumable checkpoints include training-only roles, role counters, optimizer layout, and variant
+ snapshots; validate metadata before prepared state mutation.
+
+### 7. Rewards, Advantages, and Acceleration
+
+- Pointwise/groupwise routing, per-dataset applicability, async execution, and train/eval model
+ deduplication are shared reward contracts.
+- `feedback=none` bypasses training reward/advantage structurally; do not emulate it with incidental
+ no-op overrides. Evaluation rewards remain independently configurable.
+- Reward-based algorithms delegate advantage communication to `AdvantageProcessor`.
+- Acceleration entries preserve ordered application and declared safety/stage. Lossy rollout-only
+ acceleration is incompatible with coupled trainers.
+
+### 8. Distributed Precision and Checkpointing
+
+- Validate unsupported backend and optimizer plans before model weights load. ZeRO-3 remains
+ unsupported.
+- Activation checkpointing has one owner. FSDP2 normalizes full model checkpointing to backend
+ ownership and rejects selective model policies; adapter-owned in-forward boundaries require an
+ explicit capability.
+- Wrap each forward in its own autocast region when optimizer steps or parameter swaps can occur.
+- Preserve synchronization at preprocessing, evaluation, checkpoint, and publication boundaries.
+
+## Implementation and Verification
+
+1. Establish a passing baseline for the affected tests.
+2. Change the authoritative contract and update all discovered callers/subclasses.
+3. Add focused invariant tests before broad integration tests.
+4. Verify only the affected matrix, choosing representatives from:
+ - coupled generated feedback: GRPO, plus GRPO-Guard/DPPO if loss behavior changed;
+ - decoupled generated feedback: online DPO or NFT/AWM/DGPO/CRD;
+ - generated no-feedback: DiffusionOPD or DMD2/TDM;
+ - dataset no-feedback: SFT and offline DPO;
+ - multi-role: DMD2/TDM/TDM-R1;
+ - legacy single-component and structured multimodal adapters.
+5. Cover DDP, ZeRO-2, and FSDP2 only where the changed abstraction reaches those backends. Add
+ Muon positive/negative coverage when optimizer selection is touched.
+6. Run `/ff-review` before commit.
+
+## Documentation and Examples
+
+- API/workflow changes update the matching `guidance/` document.
+- User-facing config fields update every affected example with defaults/options; example paths
+ follow `examples/{algorithm}/{finetune_type}/{model_type}/{variant}.yaml`.
+- Architecture changes update the appropriate knowledge layer. New durable invariants go in
+ `constraints.md`; detailed discoveries and fix history stay in topic docs.
+- Bug fixes follow `../../knowledge/topics/fix_patterns.md`.
+
+## Pre-Commit Gate
+
+- All affected registries, arguments, callers, tests, examples, and docs agree.
+- Changed Python files pass Black/isort; new source files carry the license header.
+- Public methods are typed and use English Google-style docstrings.
+- No silent fallback weakens a typed contract or ownership boundary.
diff --git a/.agents/skills/ff-new-algorithm/SKILL.md b/.agents/skills/ff-new-algorithm/SKILL.md
index ecb46f37f..b57a32366 100644
--- a/.agents/skills/ff-new-algorithm/SKILL.md
+++ b/.agents/skills/ff-new-algorithm/SKILL.md
@@ -1,266 +1,241 @@
---
name: ff-new-algorithm
-description: "Complete workflow for adding an online or offline training algorithm. Covers execution-contract and paradigm selection, TrainingArguments subclass, trainer implementation, registry, example config, and verification. Trigger: 'add algorithm', 'new trainer', 'new training method', 'implement algorithm'."
+description: "Add an online, offline, or distillation training algorithm to Flow-Factory. Use for a new trainer, objective, execution contract integration, or multi-role training method."
---
# New Training Algorithm Integration
-> **Authoritative reference**: `guidance/algorithms.md`
-
-## Prerequisites
-
-Determine your algorithm's characteristics:
-- **Acquisition**: Generated rollouts or a finite dataset? (`generation` / `dataset`)
-- **Feedback**: Runtime reward/advantage or none? (`runtime_reward` / `none`)
-- **Paradigm**: Coupled or Decoupled? (`constraints.md` #7)
-- **Dynamics**: Which SDE/ODE formulation? (`Flow-SDE`, `Dance-SDE`, `CPS`, `ODE`)
-- **Supervision**: Prompt-only generation, demonstrations, preference pairs, or a new typed record?
-- **Advantage**: If feedback is enabled, how are advantages computed? (Most reward-based algorithms can delegate to `AdvantageProcessor`)
-- **Loss**: What is the policy optimization objective?
-
-## Phase 1: Design
-
-1. **Study existing implementations**:
- - Coupled example: `trainers/rl/grpo.py` (GRPO)
- - Decoupled example: `trainers/rl/nft.py` (DiffusionNFT) or `trainers/rl/awm.py` (AWM)
- - Finite demonstration example: `trainers/offline/sft.py` (SFT)
- - Finite preference example: `trainers/offline/offline_dpo.py` (offline DPO)
-2. **Identify what's shared vs unique** (`constraints.md` #11):
- - Shared: the cycle loop (`BaseTrainer.start`), acquisition dispatch, progress counters,
- adapter interface, checkpoint/eval boundaries, role optimization, and exact-resume identity
- - Conditional: runtime rewards, `AdvantageProcessor`, `prepare_feedback`, and
- `compute_advantages` exist only when the feedback contract requests them
- - Unique: the loss function and the algorithm-specific hyperparameters. Never restate the loop
- - Generation hook order: `sample()` → optional `prepare_feedback()` → `optimize()`
- - Dataset hook order: official finite loader traversal → `optimize_batch(batch)`
-3. **Declare one immutable execution contract**:
- - Online RL: `ONLINE_EXECUTION_CONTRACT` (`generation + runtime_reward`)
- - Generation without rewards: `ONLINE_NO_FEEDBACK_EXECUTION_CONTRACT`
- - Finite offline training: `OFFLINE_EXECUTION_CONTRACT` (`dataset + none`)
-
-Keep execution semantics orthogonal to `PipelineIOContract`. The algorithm owns how examples are
-acquired and optimized; the adapter owns accepted input/output media, geometry, and output-state
-encoding. A new offline record shape belongs in the typed data layer, never in model-specific loss
-branches.
-
-## Phase 2: Configuration
-
-### Step 1 — Define Algorithm-Specific Arguments
-
-Create a new file `src/flow_factory/hparams/training_args/my_algo.py`:
+Read `../../../guidance/algorithms.md`, `../../../guidance/workflow.md`, Tier 1, and
+`../../knowledge/topics/train_inference_consistency.md`. Also read
+`../../knowledge/topics/component_variants.md` for more than one live trainable copy and
+`../../../guidance/datasets.md` for dataset acquisition.
+
+## 1. Classify the Algorithm
+
+Decide these independent properties before writing code:
+
+- **Acquisition**: generated collection or finite dataset (`generation` / `dataset`).
+- **Feedback**: runtime reward/advantage or none (`runtime_reward` / `none`).
+- **Paradigm**: `coupled`, `decoupled`, or `distillation`.
+- **Dynamics**: coupled objectives require SDE transition densities; decoupled/distillation
+ objectives are solver-agnostic unless their own math narrows the choice.
+- **Supervision**: prompts, demonstrations, preference pairs, or a new typed model-neutral record.
+- **State**: legacy single-component or structured multi-component trajectory.
+- **Ownership**: one live policy, several live trainable roles, or temporal reference/EMA snapshots.
+
+Use a predefined immutable contract when it matches:
+
+| Composition | Constant | Examples |
+|---|---|---|
+| `generation + runtime_reward` | `ONLINE_EXECUTION_CONTRACT` | GRPO, online DPO |
+| `generation + none` | `ONLINE_NO_FEEDBACK_EXECUTION_CONTRACT` | DiffusionOPD, DMD2, TDM |
+| `dataset + none` | `OFFLINE_EXECUTION_CONTRACT` | SFT, offline DPO |
+
+Do not infer these axes from batch fields or expose them as user-configurable fields. A genuinely new
+composition requires an explicit execution-contract/driver design, not a trainer-local branch.
+
+Study the closest direct-`BaseTrainer` implementation: GRPO for coupled replay, NFT/AWM for
+decoupled generation, DiffusionOPD or DMD2/TDM for reward-free generation, and SFT/offline DPO for
+finite data. New trainers default to direct `BaseTrainer` inheritance; the existing sanctioned
+strict extensions are GRPO-Guard/DPPO from GRPO and TDM-R1 from TDM.
+
+## 2. Add Algorithm-Specific Arguments
+
+Create `src/flow_factory/hparams/training_args/my_algo.py`. The arguments class and trainer class
+must declare the same class-level contract; it is not serialized or user-overridable.
```python
-from __future__ import annotations
from dataclasses import dataclass, field
+from typing import ClassVar, Literal
-from ._base import TrainingArguments
+from ...contracts.execution import OFFLINE_EXECUTION_CONTRACT, ExecutionContract
+from ._offline import OfflineFlowMatchingTrainingArguments
@dataclass
-class MyAlgoTrainingArguments(TrainingArguments):
- """Training arguments specific to MyAlgo."""
+class MyAlgoTrainingArguments(OfflineFlowMatchingTrainingArguments):
+ """Configure MyAlgo over finite dataset acquisitions."""
+
+ execution_contract: ClassVar[ExecutionContract] = OFFLINE_EXECUTION_CONTRACT
+ trainer_type: Literal["my-algo"] = field(default="my-algo")
my_specific_param: float = field(
default=0.1,
- metadata={"help": "Description of param."},
- )
- another_param: int = field(
- default=10,
- metadata={"help": "Description of param."},
+ metadata={"help": "Describe the objective parameter."},
)
+
+ def __post_init__(self) -> None:
+ """Validate the fixed trainer identity and objective parameters."""
+ super().__post_init__()
+ if self.trainer_type != "my-algo":
+ raise ValueError("MyAlgoTrainingArguments requires trainer_type='my-algo'")
```
-If the algorithm uses a different CFG `guidance_scale` at optimize time than at sampling/rollout time (e.g., `kl_cfg` for a reference-model branch), override `get_preprocess_guidance_scale()` so the data preprocessing stage encodes negative prompts:
+Dataset algorithms inherit `OfflineFlowMatchingTrainingArguments` so finite-loader cadence,
+explicit accumulation, and offline flow-matching fields remain shared. Generated algorithms
+inherit `TrainingArguments` and use the corresponding online contract constant. If an
+optimize/reference branch uses stronger CFG than rollout, override
+`get_preprocess_guidance_scale()` so negative conditions are encoded. Keep algorithm-owned
+objective validation in this class.
-```python
-def get_preprocess_guidance_scale(self) -> float:
- """Ensure negative prompts are encoded when optimize-time CFG needs them."""
- return max(self.guidance_scale, self.my_optimize_cfg)
-```
+Register and re-export the class in:
-See `topics/adapter_conventions.md` "Classifier-Free Guidance (CFG) Convention" for the full two-stage CFG contract.
+- `hparams/training_args/_registry.py`
+- `hparams/training_args/__init__.py`
+- `hparams/__init__.py`
-### Step 2 — Register in Argument Resolver
+## 3. Implement the Trainer Hook, Not the Loop
-Update three files in `src/flow_factory/hparams/training_args/`:
+`BaseTrainer.start()` owns seeding, acquisition dispatch, periodic save/eval boundaries, progress,
+EMA cadence, and `_after_acquisition_cycle()`. Do not override it or manually advance progress.
+Place the implementation under `src/flow_factory/trainers//my_algo.py` so the relative
+imports below match the existing `rl`, `distillation`, and `offline` package depth.
-**a)** Add import + registry entry in `_registry.py`:
+### Generated acquisition
```python
-from .my_algo import MyAlgoTrainingArguments
+from typing import ClassVar, List, Literal
-_TRAINING_ARGS_REGISTRY: Dict[str, Type[TrainingArguments]] = {
- ...
- 'my_algo': MyAlgoTrainingArguments, # Add this
-}
-```
+from ...contracts import ONLINE_EXECUTION_CONTRACT, ExecutionContract
+from ...samples import BaseSample
+from ..abc import BaseTrainer
-**b)** Add re-export in `__init__.py`:
-```python
-from .my_algo import MyAlgoTrainingArguments
-# Also add to __all__
-```
+class MyOnlineTrainer(BaseTrainer):
+ """Optimize MyAlgo from generated examples and runtime feedback."""
-**c)** Add re-export in `src/flow_factory/hparams/__init__.py`:
+ paradigm: ClassVar[Literal["decoupled"]] = "decoupled"
+ execution_contract: ClassVar[ExecutionContract] = ONLINE_EXECUTION_CONTRACT
-```python
-from .training_args import MyAlgoTrainingArguments
-# Also add to __all__
+ def sample(self) -> List[BaseSample]:
+ """Generate the state required by the objective."""
+ return self.generate_samples(
+ reward_buffer=self.reward_buffer,
+ compute_log_prob=False,
+ trajectory_indices=[-1],
+ )
+
+ def optimize(self, samples: List[BaseSample]) -> None:
+ """Apply the algorithm objective to one generated collection."""
+ ...
```
-## Phase 3: Trainer Implementation
+Prefer `generate_samples()` because it owns rollout mode, source/metadata propagation, reward
+buffering, acceleration context, and loader iteration. Override acquisition internals only when the
+algorithm truly needs a different collection shape and preserve those contracts.
-### Step 3 — Create Trainer Class
+Generated no-feedback trainers select `ONLINE_NO_FEEDBACK_EXECUTION_CONTRACT`, pass no training
+reward buffer, and do not simulate absent feedback with reward no-ops. Override `_run_training_step`
+only for a different grouping of sample/optimize work, while retaining the outer shared loop.
-For a generated online algorithm:
+### Finite dataset acquisition
```python
-# src/flow_factory/trainers/rl/my_online_algo.py
-from ...contracts import ONLINE_EXECUTION_CONTRACT
-from ..abc import BaseTrainer
-from ..registry import register_trainer
-
-@register_trainer("my-online-algo")
-class MyOnlineAlgoTrainer(BaseTrainer):
- """My generated-acquisition algorithm."""
-
- execution_contract = ONLINE_EXECUTION_CONTRACT
-
- # Do NOT define start(). BaseTrainer.start() owns the acquisition loop: reseed,
- # periodic boundaries, acquisition dispatch, EMA, and _after_acquisition_cycle().
- # evaluate(), prepare_feedback(), and compute_advantages() are concrete base methods.
- # A generation trainer implements sample() and optimize(samples).
- #
- # Vary behavior through hooks instead of restating the loop:
- # sampling_context() - wrap the rollout (e.g. install a snapshot's weights)
- # _run_training_step() - replace the sample -> feedback -> optimize middle
- # _after_gradient_step() - run right after each optimizer step
- # _after_acquisition_cycle() - run once per rollout iteration or data epoch
- # _declare_model_variants() - declare several trainable copies (see component_variants.md)
-
- def sample(self):
- """Stages 2-3: K-repeat sampling + trajectory generation."""
- # Use self.adapter.inference() for trajectory generation
- pass
-
- def optimize(self, samples):
- """Stage 6: Policy update."""
- # Use self.adapter.forward() for single-step denoising.
- # Per-forward autocast — never one outer autocast around the loop (#20a).
- # Compute loss, backprop, step
- pass
-```
+from typing import Any, ClassVar, Dict, Literal, Tuple
-For a finite offline algorithm:
+from torch.utils.data import DataLoader
-```python
-# src/flow_factory/trainers/offline/my_offline_algo.py
-from ...contracts import OFFLINE_EXECUTION_CONTRACT
+from ...contracts import OFFLINE_EXECUTION_CONTRACT, ExecutionContract
+from ...data_utils.offline_train_data import build_offline_train_dataloader
from ..abc import BaseTrainer
-from ..registry import register_trainer
-@register_trainer("my-offline-algo")
-class MyOfflineAlgoTrainer(BaseTrainer):
- """My finite-dataset algorithm."""
- paradigm = "decoupled"
- execution_contract = OFFLINE_EXECUTION_CONTRACT
+class MyOfflineTrainer(BaseTrainer):
+ """Optimize MyAlgo over one complete finite loader per data epoch."""
- def _build_train_dataloader(self):
- """Build a finite official-DistributedSampler loader for one typed schema."""
- # Reuse build_offline_train_dataloader when the supervision type matches;
- # otherwise extend the typed schema/collator first.
- ...
+ paradigm: ClassVar[Literal["decoupled"]] = "decoupled"
+ execution_contract: ClassVar[ExecutionContract] = OFFLINE_EXECUTION_CONTRACT
+
+ def _build_train_dataloader(self) -> Tuple[DataLoader, Dict[str, DataLoader]]:
+ """Build the official-distributed-sampler loader for the typed schema."""
+ return build_offline_train_dataloader(...), {}
- def optimize_batch(self, batch):
+ def optimize_batch(self, batch: Any) -> None:
"""Apply one gradient-accumulation microstep from a dataset batch."""
- # Decode output media in the dataset and encode it on demand through
- # adapter.encode_output_state(); never add target VAE latents to the cache.
...
```
-> **Note**: `AdvantageProcessor` is relevant only to `runtime_reward` feedback.
-> Reward-based trainers delegate via `self.advantage_processor.compute_advantages()` — see
-> `architecture.md` "Advantage Computation". `none` feedback bypasses rewards structurally; do not
-> emulate that by overriding reward methods with incidental no-ops.
-
-### Step 4 — Register in Trainer Registry
-
-Add to `_TRAINER_REGISTRY` in `src/flow_factory/trainers/registry.py`:
-
-```python
-'my_algo': 'flow_factory.trainers.rl.my_algo.MyAlgoTrainer',
-```
-
-## Phase 4: Configuration & Examples
-
-Create example config `examples/my_algo/lora/flux1/default.yaml`:
-
-```yaml
-model:
- model_type: "flux1"
- model_name_or_path: "black-forest-labs/FLUX.1-dev"
- finetune_type: "lora"
- target_components: ["transformer"]
-
-train:
- trainer_type: "my_algo"
- my_specific_param: 0.1
- group_size: 4
-
- num_inference_steps: 28
-
-scheduler:
- dynamics_type: "ODE" # Or appropriate dynamics
-
-data:
- datasets:
- - name: default
- dataset_dir: "path/to/dataset" # Folder with train.jsonl / test.jsonl
- train:
- weight: 1
- max_dataset_size: 1024
- eval: {}
-
-rewards:
- - name: "pickscore"
- reward_model: "pickscore"
- weight: 1.0
- batch_size: 16
-
-optimizers:
- - name: default
- learning_rate: 1e-6
- weight_decay: 1e-4
- max_grad_norm: 1.0
-```
-
-## Phase 5: Verification
-
-- [ ] `MyAlgoTrainingArguments` correctly parsed from YAML
-- [ ] `get_training_args_class('my_algo')` returns correct subclass
-- [ ] `get_trainer_class('my_algo')` loads `MyAlgoTrainer`
-- [ ] `execution_contract` matches the argument class and implemented optimization hook
-- [ ] Training runs end-to-end for ≥2 acquisition cycles without errors
-- [ ] Dataset acquisition defines one epoch as one complete finite dataloader traversal
-- [ ] Loss values are numerically reasonable (not NaN, decreasing)
-- [ ] Rewards improve over training when feedback is `runtime_reward`
-- [ ] Offline supervision media is encoded on the fly and excluded from preprocessing caches
-- [ ] Checkpoint save/load works correctly
-- [ ] Works with at least two different model adapters
-- [ ] Coupled algorithms only use SDE dynamics
-- [ ] Decoupled algorithms work with both SDE and ODE dynamics
-
-## Common Pitfalls
-
-1. **Not subclassing `TrainingArguments`** — algorithm-specific params won't be parsed from YAML
-2. **Forgetting `_registry.py` + `__init__.py` updates** — falls back to base `TrainingArguments`, losing custom params
-3. **Using ODE with coupled paradigm** — no log-probabilities available, silent incorrect gradients
-4. **Not calling `self.should_continue_training()`** — infinite loop if `max_epochs` is set
-5. **Duplicating `_initialization()` logic** — already called in `BaseTrainer.__init__`; don't re-prepare modules
-6. **Reimplementing advantage gather/scatter** — use `self.advantage_processor.compute_advantages()` instead; it handles both sampler topologies automatically
-7. **Extending `GRPOTrainer` unnecessarily** — unless your algorithm extends GRPO's PPO-clipped loss, extend `BaseTrainer` directly (as NFT and AWM do)
-8. **Optimizer-time CFG without `get_preprocess_guidance_scale()`** — if your algorithm calls `adapter.forward(guidance_scale=X)` where X > 1.0 but `training_args.guidance_scale` ≤ 1.0, negative prompts won't be encoded at preprocessing time and CFG silently falls back to no-CFG. Override `get_preprocess_guidance_scale()` in your TrainingArguments subclass to return `max(guidance_scale, your_optimize_cfg)`. See DGPO's `kl_cfg` for a real example.
-9. **Inferring online/offline behavior from batch keys** — declare `ExecutionContract`; keep acquisition and feedback independent from the model I/O schema.
-10. **Using `optimize()` for finite data** — dataset acquisition calls `optimize_batch(batch)` and advances `data_epoch` only after clean loader exhaustion.
-11. **Caching target/chosen/rejected latents** — cache prompt/input conditions only; output media is decoded and encoded on demand through the adapter output codec.
+Reuse the demonstration/preference loader when its schema matches; otherwise extend the typed data
+layer first. The driver calls `set_epoch(data_epoch)` and advances only after clean exhaustion.
+Offline optimization must:
+
+- cache input conditions only and encode outputs on demand;
+- call `adapter.prepare_condition_state()` once per batch;
+- use `adapter.encode_output_state()` and adapter-owned offline forward overrides;
+- preserve shared schedule/noise/reference scope for preference arms;
+- use explicit positive gradient accumulation with no partial-window flush.
+
+Evaluation remains generation-based.
+
+## 4. Register the Trainer
+
+Add the canonical lazy path to `_TRAINER_REGISTRY` in `trainers/registry.py`. Registry keys are
+lowercase. A decorator is optional but never replaces the static entry. Verify trainer and argument
+registry keys together and preserve direct Python-path fallback.
+
+## 5. Add Multi-Role Training Only When Required
+
+One live policy uses the default base variant and one optimizer entry. If several trainable copies
+must coexist:
+
+1. Declare ordered role names through `TrainingArguments.required_trainable_roles`; the first owns
+ canonical base routes.
+2. Return a `RoleUpdatePlan` when roles have different cadence.
+3. Let `BaseTrainer._declare_model_variants()` materialize variants before prepare. Algorithm names
+ remain in the trainer layer, not `models/`.
+4. Run forwards under `adapter.use_component_variant(role)` and updates through
+ `RoleOptimizationCoordinator` or existing role-runtime helpers.
+5. Use temporal ref/named/EMA snapshots for frozen or time-shifted weights, never a live variant.
+6. Add one top-level `optimizers:` entry per role. Muon may contribute two parameter groups for one
+ role, so store group tuples rather than assuming one-to-one ownership.
+7. Preserve one prepared `ModelBundle` and one optimizer root. Muon requires a PyTorch build with
+ `torch.optim.Muon` and is supported on DDP/FSDP2, not DeepSpeed/FSDP1. Multi-role DeepSpeed is
+ ZeRO-1/2 only. Multi-role FSDP2 requires `use_orig_params=True`; the framework rebinds the
+ registry after prepare, and optimizer references must point to the prepared root's
+ DTensor-backed parameters.
+8. Declare checkpoint runtime children before exact resume. Resumable saves include all training
+ roles, role counters, optimizer ownership, and variant snapshots.
+
+Use DMD2/TDM/TDM-R1 and `../../knowledge/topics/component_variants.md` as references.
+
+## 6. Configuration and Documentation
+
+Create `examples/{algorithm}/{finetune_type}/{model_type}/default.yaml`:
+
+- runtime rewards only for `runtime_reward`; evaluation rewards are independent;
+- dataset acquisition uses unit source weights, `sampler_type: auto`, `max_epochs`, and explicit
+ accumulation;
+- generation acquisition uses valid grouped-sampler geometry;
+- coupled algorithms use SDE dynamics;
+- `optimizers:` has one named entry per trainable role, including `max_grad_norm` and
+ `optimizer: adamw|muon`;
+- document the objective and selection in `guidance/algorithms.md` and workflow/data changes in the
+ matching guides.
+
+## 7. Verification
+
+- Trainer/argument registries resolve and their contracts match before heavyweight loading.
+- Hook validation selects `optimize` or `optimize_batch` correctly.
+- Run at least two complete acquisition cycles; for dataset mode, test clean exhaustion and a failed
+ batch that must not advance `data_epoch`.
+- Verify numerical objective invariants and legacy/structured state when applicable.
+- Verify model-only save/load and exact-state resume when state semantics change.
+- Test at least two adapters when the objective is model-neutral; dataset algorithms require an
+ offline-capable adapter.
+- Cover DDP and an affected sharded backend. Multi-role changes cover DDP/ZeRO-2/FSDP2; Muon adds
+ positive DDP/FSDP2 and early negative availability/DeepSpeed/FSDP1 cases.
+- Run `/ff-review` before commit.
+
+## Common Failures
+
+- Missing or mismatched argument/trainer `execution_contract` or `paradigm`.
+- Overriding `start()`, manually advancing progress, or duplicating `_initialization()`.
+- Inferring online/offline behavior from a batch or using `optimize()` for finite data.
+- Calling adapter inference directly while losing source, reward, or acceleration invariants.
+- Adding trainer-to-trainer inheritance without matching a sanctioned strict extension and
+ documenting why direct `BaseTrainer` plus shared helpers is insufficient.
+- Reimplementing advantage gather/scatter.
+- Caching target/chosen/rejected state or redrawing candidate-specific input conditions.
+- Using rollout CFG semantics for offline flow matching.
+- Treating a frozen reference as a variant, preparing roles separately, or assuming one group per
+ role.
+- Omitting training roles/runtime children from resumable checkpoints.
diff --git a/.agents/skills/ff-new-model/SKILL.md b/.agents/skills/ff-new-model/SKILL.md
index 18ff636ca..01a7f6e07 100644
--- a/.agents/skills/ff-new-model/SKILL.md
+++ b/.agents/skills/ff-new-model/SKILL.md
@@ -1,120 +1,187 @@
---
name: ff-new-model
-description: "Complete workflow for adding a new model adapter. Covers analysis, sample dataclass, adapter implementation (4 abstract methods + per-modality encoder overrides), registry, example YAML, and verification. Trigger: 'add model', 'support new model', 'integrate model', 'new adapter'."
+description: "Add a Flow-Factory model adapter, including component runtime, model I/O, online trajectory, offline output-state, distributed loading, checkpointing, registry, examples, and parity verification."
---
# New Model Adapter Integration
-> **Authoritative reference**: `guidance/new_model.md` — read it first.
-
-## Prerequisites
-
-Before starting, ensure you understand:
-1. The target model's diffusers pipeline (or that you'll need a pseudo-pipeline)
-2. The task type: Text-to-Image, Image-to-Image, Text-to-Video, Image-to-Video
-3. Which Sample dataclass to extend
-
-## Phase 1: Analysis
-
-1. **Identify the diffusers pipeline** for the target model
- - Check if it exists in `diffusers`: `from diffusers import `
- - If not, you'll need a pseudo-pipeline (see `guidance/new_model.md` advanced section)
-2. **Study an existing adapter** of the same task type:
- - T2I: `models/flux/flux1.py` or `models/stable_diffusion/sd3_5.py`
- - I2I: `models/flux/flux1_kontext.py` or `models/qwen_image/qwen_image_edit_plus.py`
- - T2V: `models/wan/wan2_t2v.py`
- - I2V: `models/wan/wan2_i2v.py`
-3. **Map pipeline components** to adapter responsibilities:
- - Text encoders → `encode_prompt()`, `preprocessing_modules`
- - VAE → `encode_image()` / `decode_latents()`, `preprocessing_modules`
- - Audio encoder/VAE (if any) → `encode_audio()`, `preprocessing_modules`
- - Transformer/UNet → `forward()`, `default_target_modules` (LoRA target layer names), `inference_modules`
-4. **Also read**: `topics/adapter_conventions.md` for upstream alignment rules; `topics/dtype_precision.md` for precision handling in `cast_latents()`.
-
-## Phase 2: Implementation
-
-### Step 1 — Define Sample Dataclass
-
-```python
-# src/flow_factory/models//.py
-@dataclass
-class MyModelSample(T2ISample): # or appropriate base
- _shared_fields: ClassVar[frozenset[str]] = frozenset({})
- # Add model-specific fields if needed
-```
-
-### Step 2 — Create Adapter Class
-
-```python
-class MyModelAdapter(BaseAdapter):
-
- @property
- def preprocessing_modules(self) -> List[str]:
- return ["text_encoder", "vae"] # Components for Stage 1
-
- @property
- def inference_modules(self) -> List[str]:
- return ["vae"] # Components needed at inference time
-
- @property
- def default_target_modules(self) -> List[str]:
- # LoRA target module names used when YAML sets `target_modules: default`.
- # Override only if your transformer uses non-standard attention layer names.
- return ["to_q", "to_k", "to_v", "to_out.0"]
-```
-
-> Which components are trainable is **config-driven**: the YAML `target_components` / `target_modules` fields are resolved by `BaseAdapter._parse_target_modules()` into `self.target_module_map` (set in `__init__`). Adapters do **not** override `target_module_map`.
-
-### Step 3 — Implement Required Methods
-
-| Method | Purpose | Stage | Abstract? |
-|--------|---------|-------|-----------|
-| `load_pipeline()` | Load diffusers pipeline | Init | Yes |
-| `decode_latents()` | Latents → pixels | 3 | Yes |
-| `inference()` | Full multi-step denoising | 3 | Yes |
-| `forward()` | Single-step denoising loss | 6 | Yes |
-| `encode_prompt()` | Text → embeddings | 1 | No (no-op default; override if your model consumes text) |
-| `encode_image()` | Image → latents | 1 | No (no-op default; override if your model consumes images) |
-| `encode_video()` | Video frames → latents | 1 | No (no-op default; override if your model consumes videos) |
-| `encode_audio()` | Audio → embeddings/features | 1 | No (no-op default; override if your model consumes audio) |
-| `preprocess_func()` | Raw inputs → cached tensors (dispatches to the 4 encoders) | 1 | No (concrete, override only for cross-modal preprocessing) |
-
-### Step 4 — Register
-
-Add to `_MODEL_ADAPTER_REGISTRY` in `src/flow_factory/models/registry.py`:
-```python
-'my-model': 'flow_factory.models...MyModelAdapter',
-```
-
-## Phase 3: Configuration
-
-Create example YAML config in `examples/grpo/lora//default.yaml`:
-```yaml
-model:
- model_type: "my-model"
- model_name_or_path: "org/model-name"
- finetune_type: "lora"
- target_components: ["transformer"]
-```
-
-## Phase 4: Verification
-
-Also read: `topics/parity_testing.md` for the 4-layer verification protocol.
-
-- [ ] `load_pipeline()` successfully loads the model
-- [ ] `preprocess_func()` produces correct cached tensors
-- [ ] `inference()` generates valid images/videos
-- [ ] `forward()` computes loss without errors
-- [ ] Training runs end-to-end with GRPO for ≥2 steps
-- [ ] LoRA weights save and reload correctly
-- [ ] Registry entry resolves correctly: `get_model_adapter_class('my-model')`
-- [ ] Example YAML config is valid and complete
-
-## Common Pitfalls
-
-1. **Forgetting to set `preprocessing_modules`** — causes text encoder to stay on GPU, OOM during training
-2. **Wrong `target_components` / `target_modules` (or `default_target_modules`)** — LoRA applied to wrong components/layers, no training effect
-3. **Mismatched `_shared_fields`** — data corruption during batch collation
-4. **Not handling `enable_preprocess=False`** — encoding components not loaded at inference time
-5. **Inconsistent custom field types across samples** — if a custom sample field is `Tensor` on some samples and `List[Tensor]` on others, `gather_samples` will fall back to slow pickle-based `gather_object`. Always canonicalize to a single type in `__post_init__`; prefer `List[Tensor]` for variable-length data.
-6. **Wrong `images`/`condition_images`/`audios` convention** — `preprocess_func()`, `encode_image()`, `encode_video()`, `encode_audio()`, and `inference()` all operate at **batch level**: `images` is `List[List[Image.Image]]` (`MultiImageBatch`), `condition_images` is `List[List[Tensor(C,H,W)]]` (or `List[List[PIL.Image]]` for adapters that declare `python_format_columns`, e.g. Bagel and SenseNova), and `audios` is `List[List[Tensor]]` (`MultiAudioBatch`), where the outer list indexes samples in the batch and the inner list holds each sample's items. Empty samples contribute `[]` (never `None`); single-item samples contribute `[item]` (never a bare element). Never pass a flat `List[Image]` / `List[Tensor]` or unwrap single-element lists — that breaks Arrow's homogeneous-column requirement and forces every downstream consumer to handle three input shapes. For single-condition models, `_standardize_image_input` / `_standardize_video_input` must detect the nested format with `is_multi_image_batch` / `is_multi_video_batch`, extract the first element per sample (`[batch[0] for batch in images]`), and warn if extra conditions are discarded (e.g. `Wan2_I2V._standardize_image_input`, `LTX2_I2AV._standardize_image_input`). See `topics/adapter_conventions.md` Gotcha #5 and #6.
+Read `../../../guidance/new_model.md`, Tier 1,
+`../../knowledge/topics/adapter_conventions.md`, `../../knowledge/topics/component_runtime.md`, and
+`../../knowledge/topics/parity_testing.md`. Read
+`../../knowledge/topics/structured_trajectory.md` for multi-component models and
+`../../../guidance/datasets.md` when adding offline support.
+
+## 1. Design the Adapter Contracts
+
+Decide before implementation:
+
+- task inputs and outputs: image, video, audio-video, first/last frame, or ordered heterogeneous
+ references;
+- eager Diffusers, lazy modular, or explicit pseudo component runtime;
+- canonical logical components, physical ownership roots, optional components, and aliases;
+- legacy single-component or structured multi-component trajectory;
+- online-only or lossless SFT/offline-DPO output-state support;
+- class-level I/O superset and any checkpoint-realized specialization;
+- supported finetune types, batch semantics, dtype policy, and FSDP2 capabilities.
+
+Study a reference by contract rather than name alone: Flux/SD3 for classic image pipelines, Wan for
+conditioned video, LTX2 for structured AV, MiniMax H3 for modular ordered AV, and Bagel/SenseNova for
+pseudo multi-reference runtimes.
+
+All adapters inherit directly from `BaseAdapter`. Model-specific samples inherit the matching
+task-level sample (`T2ISample`, `I2VSample`, `T2AVSample`, `I2AVSample`, etc.), never another
+model-specific sample.
+
+## 2. Build the Component Runtime
+
+All adapters implement `load_pipeline()`. The default `build_component_runtime()` wraps it in
+`ClassicPipelineRuntime`. Lazy modular or explicit-container adapters override
+`build_component_runtime()` and retain `adapter.pipeline` as the compatibility alias.
+
+Declare component behavior precisely:
+
+- canonical lookup owns component identity;
+- overrides hold prepared proxies or checkpoint/LoRA replacements;
+- declared specs are available for explicit lookup and role discovery;
+- lifecycle enumeration contains materialized canonical `torch.nn.Module` values only;
+- optional `None`, lazy-only specs, pseudo aliases, and prepared overrides are not implicit
+ lifecycle roots;
+- `materialize_components(None)` means already materialized modules. Name lazy requirements
+ explicitly.
+
+Use `has_component`, `get_component`, and `_require_component`; never gate component behavior with
+`hasattr(adapter, name)`. Declare `preprocessing_modules` and `inference_modules`. Condition/output
+encoding components are declared by their preparer/codec rather than loaded inside an encode call.
+
+Do not call Accelerator prepare, FSDP wrapping, rank broadcasts, or manual target movement in an
+adapter. `ModelLoadCoordinator` maps logical names to physical roots, and trainer initialization
+prepares one `ModelBundle` plus one optimizer. After prepare, adapter access routes through
+`RoutedComponentProxy`.
+
+## 3. Implement the Core Adapter Surface
+
+`BaseAdapter` keeps four abstract methods:
+
+| Method | Contract |
+|---|---|
+| `load_pipeline()` | Return the native pipeline/container used by the runtime. |
+| `decode_latents()` | Decode generated latent state to media. |
+| `inference()` | Run the full denoising/generation loop. |
+| `forward()` | Run one model step; this is the train-inference parity boundary. |
+
+`encode_prompt`, `encode_image`, `encode_video`, and `encode_audio` are opt-in no-op encoders.
+`preprocess_func()` dispatches them and should be overridden only for cross-modal preprocessing.
+Preserve exact batch nesting: the outer list indexes samples and inner lists hold each sample's
+media items; empty samples contribute `[]`, never `None` or a bare singleton.
+
+Configure training through `default_target_modules`; YAML `target_components`/`target_modules`
+builds `target_module_map`. Do not override the realized map.
+
+`forward()` and `inference()` must agree on all generation-affecting inputs, scheduler state,
+precision, and component order. Use `cast_latents()` symmetrically. Keep algorithm-specific logic
+out of the adapter.
+
+## 4. Add Structured State When Components Differ
+
+For independently shaped or scheduled latent components:
+
+- declare immutable `trajectory_component_order`;
+- build a `SchedulerGroup` with exactly the same names and a canonical primary scheduler;
+- emit `StructuredTrajectory` only, leaving legacy trajectory fields `None`;
+- retain per-component schedules, state/log-prob index maps, callbacks, and active masks;
+- extend protected state/bridge/reduction hooks so trainers consume terminal state, replay steps,
+ forward-process noise, and reductions without inspecting storage format;
+- derive scheduler/RNG/reduction order from the declared tuple, never mapping iteration.
+
+Single-component adapters may retain legacy storage. Do not generalize their reduction order unless
+the change preserves parity.
+
+If a new concrete sample field is required by `__post_init__`, identity normalization, or
+constructor invariants after a partial distributed gather, union it into the inherited
+`reconstruction_required_fields`. That transport contract is separate from collator
+`_shared_fields` and reward-model `required_fields`.
+
+## 5. Declare Offline Output-State Support Explicitly
+
+An adapter that claims SFT/offline-DPO support must provide every boundary below:
+
+1. An immutable `PipelineIOContract` describing model-neutral ordered input/output media,
+ semantic slots/cardinality, rates, geometry owner, and batch capability.
+2. `_resolve_pipeline_io_contract()` only when checkpoint metadata narrows a class-level superset.
+3. A declaration-only `OutputStateCodec` from `build_output_state_codec()`. It declares logical
+ required components but cannot materialize, load, move, replace, or recast them.
+4. `_validate_encoded_output_geometry()` comparing codec output against adapter-owned facts.
+5. A declaration-only `ConditionStatePreparer` only when cached conditions are not the exact
+ forward/output-codec condition.
+6. One `PreparedConditionState` reused across every candidate and policy/reference forward.
+7. A complete immutable `offline_training_forward_overrides` mapping. These values define finite-
+ data model conditioning and are independent of rollout CFG settings.
+8. An offline flow-matching objective reducer only when modality aggregation differs from online
+ trajectory reduction.
+
+Input-condition caches never contain target/chosen/rejected pixels or latent states. Shared numeric
+transforms may serve condition and output paths, but posterior `sample` versus `argmax` policy stays
+explicit at the semantic boundary. Candidate output context cannot overwrite input-owned fields.
+
+If the adapter cannot represent output state losslessly, set a non-empty actionable
+`output_state_codec_unavailable_reason`. Dataset acquisition then fails before model weights load;
+online construction remains valid.
+
+Do not override public boundary-owning wrappers such as `prepare_condition_state`,
+`encode_output_state`, `forward_state`, or the shared reducers. Extend the protected hooks named by
+their errors/docs.
+
+## 6. Preserve Distributed and Checkpoint Ownership
+
+- Declare every target component and frozen-but-shardable sibling needed in the prepared bundle.
+- Preserve `_no_split_modules` or `_repeated_blocks` metadata so FSDP wrap discovery survives the
+ bundle boundary.
+- Opt into `supports_fsdp2_cpu_efficient_loading` only when selective rank-zero/meta target
+ construction is correct. Treat additional wrap classes, default-stream unshard, backward-prefetch
+ opt-out, and in-forward checkpointing as explicit model capabilities.
+- Do not enable activation checkpointing inside `load_pipeline()`. The early backend plan selects
+ one owner: FSDP2 moves a full model policy to backend ownership and rejects selective model
+ policies; adapter-owned in-forward block boundaries require explicit opt-in.
+- Apply component load dtype during native load, then frozen/trainable storage policy. Ensure lazy
+ materialization receives both policies.
+- Save/load trainable components symmetrically. Frozen-but-shardable bundle members do not own
+ checkpoint artifacts. Test through prepared proxies as well as unprepared model-only load.
+
+## 7. Register and Add Examples
+
+Add a lowercase canonical key and lazy class path to `_MODEL_ADAPTER_REGISTRY`; preserve direct
+Python-path fallback. Follow
+`examples/{algorithm}/{finetune_type}/{model_type}/{variant}.yaml` and update path references.
+
+Add a generation example. For declared offline support, add SFT/offline-DPO examples or tested
+fixtures whose strict V2 media matches the effective checkpoint contract. Document supported modes,
+batch limits, geometry, rates, and intentionally unsupported paths.
+
+## 8. Verification
+
+- Runtime construction and lifecycle tests for the chosen classic/modular/pseudo runtime.
+- Native pipeline config, components, one-step forward, final latent, and visual parity.
+- Rollout/training `forward()` parity and initial on-policy ratio for coupled support.
+- Legacy or structured trajectory bridge, active-mask, noise, and reducer tests.
+- Selective-field distributed gather reconstructs every concrete sample with its inherited
+ `reconstruction_required_fields` intact.
+- DDP, ZeRO-2, and FSDP2 initialization/step where supported; verify bundle/proxy routing and FSDP
+ wrap/checkpoint policy.
+- LoRA/full and model-only checkpoint round trips only for finetune types claimed.
+- One complete SFT epoch and offline-DPO epoch for every declared offline I/O mode, including exact
+ geometry and prepared-condition reuse.
+- An online-only adapter's offline selection fails before heavyweight loading.
+- Registry and example parsing tests. Run `/ff-review` before commit.
+
+## Common Failures
+
+- Membership through adapter attributes, eager materialization of all lazy specs, or alias double
+ movement.
+- Preparing components outside the bundle or dropping repeated-block wrap metadata.
+- Scheduler order from a mapping or multimodal state written into legacy fields.
+- Adapter-to-adapter or model-sample-to-model-sample inheritance.
+- Codec/preparer construction with materialization or device/dtype side effects.
+- Using the class I/O superset instead of the checkpoint-effective contract.
+- Redrawing an input condition per preference candidate or leaking sampling CFG into offline loss.
+- Overriding public contract wrappers instead of protected hooks.
+- Enabling both model and backend activation checkpointing.
diff --git a/.agents/skills/ff-new-reward/SKILL.md b/.agents/skills/ff-new-reward/SKILL.md
index 911d7d2dd..564258841 100644
--- a/.agents/skills/ff-new-reward/SKILL.md
+++ b/.agents/skills/ff-new-reward/SKILL.md
@@ -1,118 +1,136 @@
---
name: ff-new-reward
-description: "Complete workflow for adding a new reward model. Covers pointwise vs groupwise design, __call__ contract, registration, YAML config, multi-reward setup, and verification. Trigger: 'add reward', 'new reward model', 'custom reward', 'scoring function'."
+description: "Add a Flow-Factory reward model with pointwise/groupwise dispatch, media conversion, per-dataset routing, async execution, backend-safe loading, registry, configuration, and verification."
---
# New Reward Model Integration
-> **Authoritative reference**: `guidance/rewards.md` — read it first.
-> **Template**: `src/flow_factory/rewards/my_reward.py`
+Read `../../../guidance/rewards.md`, `../../knowledge/constraints.md` #13, and the current
+`rewards/abc.py`, `rewards/reward_processor.py`, and `rewards/loader.py` contracts.
-## Prerequisites
+## 1. Choose the Dispatch Contract
-Determine your reward type:
-- **Pointwise**: Each sample scored independently (e.g., aesthetic score, CLIP similarity)
-- **Groupwise**: Scores depend on comparison within a group (e.g., ranking, preference)
+- **Pointwise**: each input is scored independently. A call receives a non-empty chunk whose length
+ is at most configured `batch_size`; tail chunks and per-dataset applicability gating can make it
+ smaller.
+- **Groupwise**: one call receives a complete `unique_id` group, either local or reconstructed
+ across ranks. Configured pointwise `batch_size` does not define this call.
-## Phase 1: Design
+Return exactly one finite score per input passed to the call, in the same order. Do not return NaN
+as a “not applicable” marker: `RewardProcessor` owns applicability masks, NaN padding outside model
+calls, and `sample.applicable_rewards`.
-1. **Choose base class**: `PointwiseRewardModel` or `GroupwiseRewardModel`
-2. **Identify required inputs**: What fields from `Sample` does your reward need?
- - Common: `prompt`, `image`, `video`, `condition_images`, `condition_videos`
- - Set `required_fields` tuple accordingly
-3. **Input format**: PIL Images (default) or Tensors?
- - Set `use_tensor_inputs = True` if your model needs raw tensors
+Declare only fields needed from `BaseSample` in `required_fields`. Common fields include `prompt`,
+`image`, `video`, `audio`, `condition_images`, and `condition_videos`. Additional JSONL metadata is
+available as one JSON-encoded `metadata` string and must be parsed explicitly.
-## Phase 2: Implementation
+`required_fields` selects reward-consumer data; it does not replace a concrete sample class's
+`reconstruction_required_fields`. If adding a sample field that its constructor or `__post_init__`
+needs after a partial gather, update that inherited class contract separately. Collator
+`_shared_fields` is a third, independent concern.
-### Create the reward model file
+Set `use_tensor_inputs` deliberately:
+
+- `False`: images/video frames arrive as PIL and audio as NumPy;
+- `True`: media arrives as tensors.
+
+Condition media retains nested per-sample item structure.
+
+## 2. Implement the Model
```python
-# src/flow_factory/rewards/.py
-from .abc import PointwiseRewardModel, RewardModelOutput
-from ..hparams import RewardArguments
+from typing import Any, List, Optional
+
+import torch
from accelerate import Accelerator
-from typing import Optional, List
from PIL import Image
-import torch
+
+from ..hparams import RewardArguments
+from .abc import PointwiseRewardModel, RewardModelOutput
+
class MyRewardModel(PointwiseRewardModel):
+ """Score prompt-image alignment for each generated sample."""
+
required_fields = ("prompt", "image")
use_tensor_inputs = False
- def __init__(self, config: RewardArguments, accelerator: Accelerator):
+ def __init__(self, config: RewardArguments, accelerator: Accelerator) -> None:
+ """Load the reward network using the configured device and dtype."""
super().__init__(config, accelerator)
- # Load your model, processor, etc.
- # Use self.device and self.dtype from base class
+ ...
@torch.no_grad()
def __call__(
self,
prompt: List[str],
image: Optional[List[Image.Image]] = None,
- video: Optional[List[List[Image.Image]]] = None,
- audio: Optional[List[torch.Tensor]] = None,
- condition_images=None,
- condition_videos=None,
- **kwargs,
+ **kwargs: Any,
) -> RewardModelOutput:
- # Compute rewards — shape must be (batch_size,) for Pointwise
- # or (group_size,) for Groupwise
- rewards = torch.zeros(len(prompt), device=self.device)
- return RewardModelOutput(rewards=rewards)
+ """Return one finite score per received prompt-media pair."""
+ scores = ...
+ return RewardModelOutput(rewards=scores)
```
-### Key constraints for `__call__`:
-- **Pointwise**: Input length = `config.batch_size`. Return rewards shape `(batch_size,)`
-- **Groupwise**: Input length = `group_size`. You handle batching yourself. Return rewards shape `(group_size,)`
-- Always use `@torch.no_grad()` decorator
-- Return `RewardModelOutput` (not raw tensors)
+Use `self.device` and `self.dtype`; do not hardcode CUDA. Keep public signatures typed and
+Google-style. `@torch.no_grad()` is required for inference-only scoring.
-## Phase 3: Register
+Construction runs inside `ModelLoadCoordinator`'s REWARD load scope and reward resources remain
+full per-rank replicas. A reward implementation must not call `accelerator.prepare()`, enter
+target-only FSDP loading state, or mutate trainer component ownership.
-Add to `_REWARD_MODEL_REGISTRY` in `src/flow_factory/rewards/registry.py`:
-```python
-'my_reward': 'flow_factory.rewards..MyRewardModel',
-```
+## 3. Register and Configure
-## Phase 4: Configuration
+Add a lowercase canonical lazy path to `_REWARD_MODEL_REGISTRY`. Preserve direct Python-path
+fallback; a decorator is optional and does not replace the static entry.
-Use in YAML config:
```yaml
rewards:
- - name: "my_reward"
- reward_model: "my_reward" # Must match registry key
- model_path: "org/model-name" # HuggingFace model path (if applicable)
- dtype: "bfloat16"
- device: "cuda"
+ - name: my_reward
+ reward_model: my_reward
+ model_path: org/model-name
+ dtype: bfloat16
+ device: cuda
batch_size: 16
+ applicable_datasets: [alignment]
+ weight: 1.0
+ async_reward: false
+ num_workers: 1
```
-Multi-reward setup:
-```yaml
-rewards:
- - name: "aesthetic"
- reward_model: "PickScore"
- weight: 0.7
- - name: "custom"
- reward_model: "my_reward"
- weight: 0.3
-```
-
-## Phase 5: Verification
-
-- [ ] `__init__` loads model without errors
-- [ ] `__call__` returns correct reward shape
-- [ ] Rewards are numerically reasonable (not all zeros, no NaN/Inf)
-- [ ] Works with `RewardProcessor` dispatch (Pointwise/Groupwise routing)
-- [ ] Works in multi-reward setup with weight aggregation
-- [ ] Device placement correct (respects `config.device`)
-- [ ] Registry entry resolves: `get_reward_model_class('my_reward')`
-
-## Common Pitfalls
-
-1. **Wrong return shape** — Pointwise must return `(batch_size,)`, Groupwise `(group_size,)`
-2. **Forgetting `@torch.no_grad()`** — causes reward computation to build unnecessary graph, OOM
-3. **Hardcoding device** — use `self.device` from base class, not `torch.device('cuda')`
-4. **Not setting `required_fields`** — `RewardProcessor` won't pass the right data to your model
-5. **Mixing paradigms** — don't inherit `PointwiseRewardModel` if your reward needs group context
+`name` identifies the configured reward stream; `reward_model` resolves its implementation.
+`applicable_datasets` is resolved to dataset names/source IDs, and `weight` may be one scalar or a
+per-dataset mapping. Training and `eval_rewards:` are independent configurations. Runtime training
+rewards are valid only for an execution contract with `feedback=runtime_reward`; reward-free and
+offline trainers may still use evaluation rewards.
+
+The loader deduplicates train/eval entries with the same model identity. Do not keep mutable
+per-config-name state inside a shared model call.
+
+When `async_reward` is enabled, calls run in worker threads and may use dedicated CUDA streams.
+Implementations must be thread-safe or require `num_workers: 1`; tail pointwise batches still run
+during finalize.
+
+## 4. Verification
+
+- Direct pointwise calls at full and tail lengths, or one complete groupwise call.
+- Exact result shape/order and rejection of NaN/Inf.
+- Source-gated partial and no-applicable subsets through `RewardProcessor`.
+- Groupwise local and distributed reconstruction when applicable.
+- Selective groupwise gathers retain every concrete sample `reconstruction_required_fields` entry.
+- PIL/NumPy and tensor media conversion for every declared field.
+- Scalar and per-dataset weighted multi-reward aggregation.
+- Sync and async execution, including async tail flush and thread-safety assumptions.
+- Train/eval deduplication and exactly one REWARD load scope per unique model.
+- Registry lookup, direct-path fallback, device placement, and config parsing.
+- Run `/ff-review` before commit.
+
+## Common Failures
+
+- Assuming every pointwise call has exactly `config.batch_size` inputs.
+- Returning a raw scalar, wrong ordering, non-finite values, or framework-owned applicability NaNs.
+- Omitting a required field or flattening nested condition media.
+- Hardcoding a device or preparing/sharding a reward as a target component.
+- Mutating shared model state by reward configuration name.
+- Using a pointwise base for a group-dependent score, or enabling async workers for non-thread-safe
+ code.
diff --git a/.agents/skills/ff-review/SKILL.md b/.agents/skills/ff-review/SKILL.md
index 83f37bbf7..da6de83dc 100644
--- a/.agents/skills/ff-review/SKILL.md
+++ b/.agents/skills/ff-review/SKILL.md
@@ -1,116 +1,176 @@
---
name: ff-review
-description: "Mandatory pre-commit code review gate. Checks constraint violations, cross-module consistency, and implementation quality. Trigger proactively when changes span multiple files or touch shared infrastructure. Trigger: 'review', 'check before commit'."
+description: "Review Flow-Factory changes before commit or merge for contract violations, cross-module drift, distributed/checkpoint safety, docs consistency, implementation quality, and test evidence."
---
# Code Review Workflow
-## Process Overview
+## 1. Capture the Exact Scope
-```
-1. Capture changes → git diff
-2. Load constraints → .agents/knowledge/constraints.md
-3. Review against constraints and architecture
-4. Route by verdict:
- ✓ Safe → Proceed with commit
- ⚠ Needs-attention → Fix issues, then commit
- ✗ Risky → Halt and report
-```
-
-## Step 1: Capture Changes
+For a commit review, inspect staged and unstaged scopes separately:
```bash
-git diff HEAD # All changes
-git status # Modified files
+git status --short
+git diff --check
+git diff --cached --stat
+git diff --cached
+git diff
```
-## Step 2: Load Context
-
-- Read `.agents/knowledge/constraints.md` — All hard constraints
-- Reference `.agents/knowledge/architecture.md` — Module dependencies
-- Identify which modules are affected by the changes
-
-## Step 3: Review Checklist
-
-### Constraint Compliance
-- [ ] No constraint violations found
-- [ ] Registry entries updated if classes moved/renamed (#1–4)
-- [ ] Pipeline order preserved (#6)
-- [ ] Coupled/decoupled paradigm respected (#7)
-- [ ] Base class interfaces not broken (#11–13)
-- [ ] Config fields synchronized with YAML examples (#15–17)
-
-### Cross-Module Consistency
-- [ ] Changes to `abc.py` base classes reflected in ALL subclasses (grpo, grpo-guard, dppo, nft, awm, dgpo, dpo, crd, diffusion-opd)
-- [ ] Changes to `hparams/` reflected in ALL example configs
-- [ ] Changes to `AdvantageProcessor` compatible with all trainers
-- [ ] Registry keys match actual import paths
-- [ ] Sample dataclass `_shared_fields` consistent
-
-### Implementation Quality
-- [ ] No hardcoded devices (use `self.device` or `accelerator.device`)
-- [ ] `@torch.no_grad()` on reward model `__call__`
-- [ ] Proper synchronization barriers for distributed code
-- [ ] No ZeRO-3 usage
-- [ ] Type annotations on public methods
-
-### Code Style
-- [ ] Black formatting (`line-length=100`)
-- [ ] isort compliance (`profile="black"`)
-- [ ] English comments and docstrings
-- [ ] Apache 2.0 license header on new files
-- [ ] No unnecessary wildcard imports (except `hparams`)
-- [ ] **Top-level imports only** (constraint #22) — see that file for the three sanctioned exceptions (optional deps via `try/except ImportError`, backend-gated runtime feature checks like DeepSpeed/FSDP, unresolvable circular imports).
-
-### Documentation
-- [ ] `guidance/` docs updated if behavior changed
-- [ ] New config fields added to ALL example configs with defaults and `# Options:` comments
-- [ ] PR title follows format: `[{modules}] {type}: {description}`
-
-## Step 4: Route by Verdict
-
-### ✓ Safe
-No issues found. Proceed with commit.
-
-### ⚠ Needs-Attention
-Issues found but fixable:
-1. List each issue with file and line
-2. Fix identified problems
-3. Re-stage and re-review
-
-### ✗ Risky
-Potential breaking changes:
-1. Halt commit
-2. Report findings with severity
-3. Await explicit user approval
-
-## After Commit
-
-- Run `black --check src/ && isort --check src/` to confirm formatting compliance.
-- Verify PR title follows `[{modules}] {type}: {description}` format.
-- If this was a bug fix, follow `topics/fix_patterns.md` archival process.
-
-## Pre-Review Reading
-
-Before reviewing, always read Tier 1: `constraints.md`, `architecture.md`, `philosophy.md`.
-
-Additionally, read based on diff scope:
-
-| Diff touches... | Also read |
-|----------------|-----------|
-| `models/` | `topics/adapter_conventions.md`, `topics/parity_testing.md` |
-| `trainers/` | `topics/train_inference_consistency.md`, `topics/autocast_param_swap.md` |
-| `scheduler/` | `topics/train_inference_consistency.md`, `topics/dtype_precision.md` |
-| New adapter | `topics/adapter_conventions.md`, `topics/parity_testing.md` |
-| dtype/precision | `topics/dtype_precision.md`, `topics/autocast_param_swap.md` |
-
-## Common Issues Found in Review
-
-1. **Registry path stale** — Class moved but registry not updated
-2. **Config field renamed** — YAML examples still use old name
-3. **New config field not in examples** — Users won't discover it; add with default value and `# Options:` comment
-4. **Base class change not propagated** — Subclass override now has wrong signature
-5. **Missing `wait_for_everyone()`** — Distributed deadlock risk
-6. **Reward shape mismatch** — Pointwise returning wrong batch dim
-7. **License header missing** — New files without Apache 2.0 header
-8. **Autocast spans a forward** — flat loss / KL ≈ 0 (fp32 master); see #20a / `topics/autocast_param_swap.md`
+For a PR-wide review, record the base and commit range and inspect that range in addition to local
+changes. Do not assume `git diff HEAD` represents every intended change.
+
+Read Tier 1 and derive affected trainers, adapters, rewards, accelerators, and argument classes from
+their registries; avoid hard-coded component lists.
+
+## 2. Load Scope-Specific References
+
+| Diff touches | Also read |
+|---|---|
+| Execution contracts, trainer loop, offline data | `../../../guidance/workflow.md`, `../../../guidance/algorithms.md`, `../../../guidance/datasets.md` |
+| Adapter semantics or model parity | `../../knowledge/topics/adapter_conventions.md`, `../../knowledge/topics/parity_testing.md` |
+| Component runtime, loading, bundle | `../../knowledge/topics/component_runtime.md` |
+| Trajectory, sample, scheduler group | `../../knowledge/topics/structured_trajectory.md`, `../../knowledge/topics/train_inference_consistency.md` |
+| Variant, role optimizer, Muon | `../../knowledge/topics/component_variants.md` |
+| Dtype, autocast, parameter swaps | `../../knowledge/topics/dtype_precision.md`, `../../knowledge/topics/autocast_param_swap.md` |
+| Gradient checkpoint/FSDP memory | `../../../guidance/new_model.md` checkpointing contract |
+| Reward processing | `../../../guidance/rewards.md` |
+| Acceleration | `../../../guidance/acceleration.md` |
+| `.agents/` | `../../knowledge/docs_maintenance.md`, agent-doc maintenance rule |
+
+## 3. Review Contract Boundaries
+
+### Registries and configuration
+
+- Static registry keys, lazy import paths, direct-path fallback, and argument registry agree.
+- Keys follow canonical naming; moved/renamed classes update every registry/export.
+- Trainer and algorithm-specific arguments declare the same immutable `ExecutionContract`; users
+ cannot configure its axes.
+- Added/renamed/removed fields update all affected examples and consumers. Example paths follow the
+ project convention.
+
+### Execution and data acquisition
+
+- `BaseTrainer.start()` and acquisition drivers remain authoritative. The selected hook is
+ `optimize(samples)` for generation or `optimize_batch(batch)` for dataset acquisition.
+- `generation + runtime_reward`, `generation + none`, and `dataset + none` execute only their
+ declared stages.
+- Dataset training uses an official finite `DistributedSampler`, explicit accumulation, full clean
+ traversal, unit source weights, and no Accelerator preparation of the training loader.
+- Only input conditions enter preprocessing cache; output supervision is encoded on demand.
+- `optimizer_step`, `rollout_iteration`, and `data_epoch` advance at their own boundaries.
+- Exact runtime identity includes changed objective/data/backend/optimizer and replayed evaluation
+ semantics.
+
+### Adapter, pipeline I/O, and trajectories
+
+- Base adapter keeps four abstract methods; adapters and model-specific samples retain flat/task-
+ level inheritance contracts.
+- Offline support declares a valid effective `PipelineIOContract`, declaration-only condition
+ preparer/output codec, exact geometry validation, one prepared condition per request, complete
+ offline forward overrides, and an explicit blocker when unsupported.
+- Public boundary-owning wrappers are not overridden; protected hooks carry specialization.
+- `forward()`/`inference()` preserve inputs, precision, scheduler state, and parity.
+- Structured trajectories own component order, maps, callbacks, masks, noise, and reductions;
+ trainers consume only bridge APIs. Legacy single-component behavior remains unchanged.
+- Concrete sample fields required after partial gather are inherited through
+ `reconstruction_required_fields`; reward `required_fields` and collator `_shared_fields` remain
+ separate contracts.
+
+### Component runtime, loading, and prepared ownership
+
+- Canonical, override, declared, materialized, optional, and alias paths remain distinct.
+- Membership uses runtime APIs, not adapter attributes; omitted lazy materialization does not load
+ all specs.
+- Logical-to-physical loading goes through `ModelLoadCoordinator`; auxiliary/reward roots remain
+ replicas and target-only backend state does not leak.
+- Every target/frozen-shardable/variant route enters one `ModelBundle` prepared with one optimizer;
+ canonical forwards route through `RoutedComponentProxy` afterward.
+- Stable names and `_no_split_modules`/`_repeated_blocks` metadata survive the bundle boundary.
+- Save/load iterates trainable ownership symmetrically and skips frozen-only checkpoint artifacts.
+
+### Variants, optimizers, distributed plans, and checkpoints
+
+- Algorithm vocabulary and role cadence stay in trainers. Temporal ref/EMA/old state is not modeled
+ as a live trainable variant.
+- Variants are declared before prepare; parameter and optimizer ownership is disjoint/exhaustive.
+- No code assumes one group per role. Muon matrices and AdamW fallback groups share one
+ `CompositeOptimizer` root.
+- ZeRO-3 is rejected. Muon availability and DeepSpeed/FSDP1 incompatibility fail before model load;
+ multi-role DeepSpeed is ZeRO-1/2. Multi-role FSDP2 requires `use_orig_params=True`; registry and
+ optimizer references must point to the replacement DTensor-backed parameters owned by the
+ prepared model root.
+- Activation checkpointing has one owner. FSDP2 full policy moves to backend ownership, selective
+ model policy is rejected, and adapter-owned in-forward boundaries require explicit capability.
+- Resumable checkpoints include all training roles/runtime children and validate metadata before
+ Accelerate state mutation. Model-only export scope remains intentional.
+- Distributed checkpoint phases and publication are all-rank symmetric and synchronized.
+
+### Rewards, acceleration, and numerical quality
+
+- Pointwise calls accept tail/source-gated chunks and return one finite value per actual input;
+ groupwise paths preserve complete group order.
+- Per-dataset applicability/weights, async tail flush, and train/eval model deduplication remain
+ correct.
+- Reward-free contracts do not create incidental training reward work.
+- Acceleration entries preserve declared safety/stage and list order; lossy rollout-only plugins do
+ not run on coupled trainers.
+- No hardcoded device bypasses adapter/reward/backend ownership.
+- Each forward has an appropriate autocast boundary when optimizer steps or param swaps occur.
+- Required rank barriers are present without introducing asymmetric collectives or filesystem work.
+
+### Code and documentation quality
+
+- Public functions/methods are typed and have English Google-style docstrings; new source files
+ carry the Apache header.
+- Imports follow project style and sanctioned local-import exceptions only.
+- Errors fail fast with concrete user-facing config values; no silent fallback weakens a contract.
+- README, guidance, examples, code comments/docstrings, and `.agents/` docs describe current owners,
+ supported modes, and paths without claiming unexecuted quality.
+- Published text/log snippets contain no credentials, tokens, personal absolute paths, hostnames, or
+ machine-specific details.
+
+## 4. Verify Proportionally
+
+Run focused tests tied to each changed contract before broad tests. Useful suites include:
+
+- execution: `tests/contracts/test_execution_contract.py`,
+ `tests/trainers/test_execution_kernel.py`;
+- offline: `tests/hparams/test_offline_training_args.py`, offline data/trainer tests;
+- runtime/I/O: component runtime, pipeline contract, output-state lifecycle tests;
+- trajectory: `tests/models/trajectory/` and bridge/reduction tests;
+- variants/optimizer: component variant, role optimization, Muon, and multirole tests;
+- distributed/checkpoint: distributed-plan, checkpoint layout/runtime identity/resume tests;
+- rewards: loader-context and processor/reconstruction tests.
+
+Run Black/isort on changed Python files as the commit gate, then run the documented full-tree checks.
+If the repository has pre-existing full-tree failures, prove the baseline and distinguish them from
+new regressions; do not hide them or expand scope silently. Validate Markdown links and example
+paths for docs changes.
+
+GPU/distributed evidence should cover only affected compositions/backends, but any claimed support
+must have a representative run. Multi-role/Muon or loading/checkpoint changes normally require DDP,
+ZeRO-2, and FSDP2 coverage plus intended early-rejection cases.
+
+## 5. Verdict
+
+- **Safe**: contracts, tests, docs, and evidence agree; proceed only with the user's authorized
+ commit/push scope.
+- **Needs attention**: list each issue with file/line and fix/re-review before commit.
+- **Risky**: halt when behavior, compatibility, data, or distributed correctness remains uncertain
+ and request explicit direction.
+
+After an authorized commit, verify the final diff and formatting. Bug fixes also follow
+`../../knowledge/topics/fix_patterns.md`.
+
+## Frequent Review Findings
+
+- Trainer/argument contract drift or wrong acquisition hook.
+- Offline path invoking rollout/reward or caching supervision state.
+- Adapter capability checked only after weights load.
+- Runtime membership via `hasattr`, alias double movement, or component prepared outside the bundle.
+- Missing frozen-member checkpoint symmetry or lost repeated-block wrap metadata.
+- Frozen reference represented as a variant or one-group-per-role assumption.
+- Training role/runtime child omitted from resume metadata.
+- Muon accepted on an unsupported backend or duplicate activation-checkpoint owners.
+- Public docs naming a pre-refactor function, owner, path, or unverified support status.
From f02b941e9b0097c89cabaa63b5de1aa1b2dadecd Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 08:19:43 +0800
Subject: [PATCH 72/76] fix(examples): move offline fixtures to dataset root
---
.agents/knowledge/topics/fix_patterns.md | 14 ++++
README.md | 4 +
dataset/offline_dpo_sd3_5/train.jsonl | 2 +
.../data => dataset}/sft_sd3_5/train.jsonl | 4 +-
examples/README.md | 4 +-
examples/data/offline_dpo_sd3_5/train.jsonl | 2 -
examples/offline_dpo/lora/sd3_5/default.yaml | 4 +-
examples/sft/lora/sd3_5/default.yaml | 4 +-
guidance/datasets.md | 9 ++-
tests/examples/test_offline_examples.py | 74 +++++++++++++++++++
10 files changed, 108 insertions(+), 13 deletions(-)
create mode 100644 dataset/offline_dpo_sd3_5/train.jsonl
rename {examples/data => dataset}/sft_sd3_5/train.jsonl (51%)
delete mode 100644 examples/data/offline_dpo_sd3_5/train.jsonl
create mode 100644 tests/examples/test_offline_examples.py
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 198877972..a58d3e92e 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -740,6 +740,20 @@ Based on the fix type, write the fix entry to the appropriate document:
must express quality metrics through the evaluation-only reward surface.
- **Related Constraint**: #7
+### Repository dataset fixtures must not live inside the example-config tree
+- **Date**: 2026-08-31
+- **Symptom**: The checked-in SD3.5 SFT and offline-DPO manifests lived under `examples/data`, while
+ every other repository dataset and the public dataset guide used the root `dataset/` hierarchy.
+- **Root Cause**: The initial smoke fixtures were colocated with their configs without preserving
+ the repository boundary between executable example configs and dataset assets.
+- **Fix**: The manifests moved to `dataset/sft_sd3_5` and `dataset/offline_dpo_sd3_5`; their YAML,
+ Markdown links, and directory-depth-sensitive asset paths moved with them. A production-parser
+ regression now loads both configs and manifests and verifies every supervision asset exists.
+- **Lesson**: Treat example configs and their datasets as separate public surfaces. When moving a
+ manifest, recompute every dataset-root-relative media path and test the resolved files rather
+ than checking only the configured directory string.
+- **Related Constraint**: N/A
+
## Cross-refs
- UP: [Hard Constraints](../constraints.md), [Architecture](../architecture.md)
diff --git a/README.md b/README.md
index 4c49a7424..ccdf84ba7 100644
--- a/README.md
+++ b/README.md
@@ -241,6 +241,10 @@ SFT and offline DPO use strict JSONL with `schema_version: 2`. Public media obje
The optional input-only `slot` field binds sparse conditions to adapter-declared semantic
arguments. Unslotted media fills remaining slots positionally; supervision outputs reject slots.
+The checked-in [SFT demonstration fixture](dataset/sft_sd3_5/train.jsonl) and
+[offline-DPO preference fixture](dataset/offline_dpo_sd3_5/train.jsonl) provide minimal examples
+under the repository's canonical `dataset/` root.
+
Prompt and input-condition encodings are cached. Target, chosen, and rejected media are decoded and
encoded on the fly; their VAE latents are never stored in the preprocessing cache. One offline
epoch is one complete dataloader traversal sharded by PyTorch's official `DistributedSampler`. See the
diff --git a/dataset/offline_dpo_sd3_5/train.jsonl b/dataset/offline_dpo_sd3_5/train.jsonl
new file mode 100644
index 000000000..a2185d729
--- /dev/null
+++ b/dataset/offline_dpo_sd3_5/train.jsonl
@@ -0,0 +1,2 @@
+{"schema_version":2,"input":{"prompt":"Render a clean Flow-Factory emblem.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"../../assets/logo-no-bg.png"}]},"rejected":{"media":[{"type":"image","path":"../../assets/wandb_metrics.png"}]}},"metadata":{"fixture":"repository-images"}}
+{"schema_version":2,"input":{"prompt":"Render an experiment dashboard.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"../../assets/wandb_images.png"}]},"rejected":{"media":[{"type":"image","path":"../../assets/logo.png"}]}},"metadata":{"fixture":"repository-images"}}
diff --git a/examples/data/sft_sd3_5/train.jsonl b/dataset/sft_sd3_5/train.jsonl
similarity index 51%
rename from examples/data/sft_sd3_5/train.jsonl
rename to dataset/sft_sd3_5/train.jsonl
index 0e4b77fc1..e3253c3c7 100644
--- a/examples/data/sft_sd3_5/train.jsonl
+++ b/dataset/sft_sd3_5/train.jsonl
@@ -1,2 +1,2 @@
-{"schema_version":2,"input":{"prompt":"Render a clean Flow-Factory emblem.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"../../../assets/logo-no-bg.png"}]}},"metadata":{"fixture":"repository-image"}}
-{"schema_version":2,"input":{"prompt":"Render an experiment dashboard.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"../../../assets/wandb_images.png"}]}},"metadata":{"fixture":"repository-image"}}
+{"schema_version":2,"input":{"prompt":"Render a clean Flow-Factory emblem.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"../../assets/logo-no-bg.png"}]}},"metadata":{"fixture":"repository-image"}}
+{"schema_version":2,"input":{"prompt":"Render an experiment dashboard.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"../../assets/wandb_images.png"}]}},"metadata":{"fixture":"repository-image"}}
diff --git a/examples/README.md b/examples/README.md
index 7d93dbd34..78ff11622 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -27,10 +27,10 @@ ff-train examples/grpo/lora/flux1/default.yaml
## Offline examples
- [`sft` with SD3.5](sft/lora/sd3_5/default.yaml) consumes V2
- `demonstration` records from [`examples/data/sft_sd3_5`](data/sft_sd3_5/train.jsonl).
+ `demonstration` records from [`dataset/sft_sd3_5`](../dataset/sft_sd3_5/train.jsonl).
- [`offline-dpo` with SD3.5](offline_dpo/lora/sd3_5/default.yaml) consumes V2
`preference` records from
- [`examples/data/offline_dpo_sd3_5`](data/offline_dpo_sd3_5/train.jsonl).
+ [`dataset/offline_dpo_sd3_5`](../dataset/offline_dpo_sd3_5/train.jsonl).
The two tiny manifests reuse repository images so their paths resolve without a separate dataset
download. They are configuration and smoke-test fixtures, not quality-training datasets. Offline
diff --git a/examples/data/offline_dpo_sd3_5/train.jsonl b/examples/data/offline_dpo_sd3_5/train.jsonl
deleted file mode 100644
index f8c47b544..000000000
--- a/examples/data/offline_dpo_sd3_5/train.jsonl
+++ /dev/null
@@ -1,2 +0,0 @@
-{"schema_version":2,"input":{"prompt":"Render a clean Flow-Factory emblem.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"../../../assets/logo-no-bg.png"}]},"rejected":{"media":[{"type":"image","path":"../../../assets/wandb_metrics.png"}]}},"metadata":{"fixture":"repository-images"}}
-{"schema_version":2,"input":{"prompt":"Render an experiment dashboard.","media":[]},"supervision":{"type":"preference","chosen":{"media":[{"type":"image","path":"../../../assets/wandb_images.png"}]},"rejected":{"media":[{"type":"image","path":"../../../assets/logo.png"}]}},"metadata":{"fixture":"repository-images"}}
diff --git a/examples/offline_dpo/lora/sd3_5/default.yaml b/examples/offline_dpo/lora/sd3_5/default.yaml
index 1a1698785..de5d1f819 100644
--- a/examples/offline_dpo/lora/sd3_5/default.yaml
+++ b/examples/offline_dpo/lora/sd3_5/default.yaml
@@ -1,4 +1,4 @@
-# Single-process offline-DPO smoke recipe over a V2 preference manifest.
+# Single-process offline-DPO smoke recipe over the checked-in dataset/offline_dpo_sd3_5 fixture.
# Replace the tiny repository-image dataset with real preference pairs for quality runs.
launcher: "accelerate"
config_file: null
@@ -9,7 +9,7 @@ mixed_precision: "bf16"
data:
datasets:
- name: offline_preferences
- dataset_dir: "examples/data/offline_dpo_sd3_5"
+ dataset_dir: "dataset/offline_dpo_sd3_5"
train:
weight: 1 # Offline epochs require unit source weights and full traversal.
max_dataset_size: 2
diff --git a/examples/sft/lora/sd3_5/default.yaml b/examples/sft/lora/sd3_5/default.yaml
index 2d07fa263..3b0648de3 100644
--- a/examples/sft/lora/sd3_5/default.yaml
+++ b/examples/sft/lora/sd3_5/default.yaml
@@ -1,4 +1,4 @@
-# Single-process SFT smoke recipe over a V2 demonstration manifest.
+# Single-process SFT smoke recipe over the checked-in dataset/sft_sd3_5 fixture.
# Replace the tiny repository-image dataset with a real training corpus for quality runs.
launcher: "accelerate"
config_file: null
@@ -9,7 +9,7 @@ mixed_precision: "bf16"
data:
datasets:
- name: offline_demonstrations
- dataset_dir: "examples/data/sft_sd3_5"
+ dataset_dir: "dataset/sft_sd3_5"
train:
weight: 1 # Offline epochs require unit source weights and full traversal.
max_dataset_size: 2
diff --git a/guidance/datasets.md b/guidance/datasets.md
index c4ebe0eea..a2a9acfdd 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -121,9 +121,12 @@ Decoded audio is a detached CPU `float32` waveform shaped `(channels, samples)`.
source-clock truncation, channel conversion, the single model-rate conversion, posterior selection,
and latent packing remain adapter-owned.
-Tiny schema-complete fixtures and configs are available for
-[SFT](../examples/sft/lora/sd3_5/default.yaml) and
-[offline DPO](../examples/offline_dpo/lora/sd3_5/default.yaml).
+Tiny schema-complete fixtures live under the repository `dataset/` root:
+
+- [SFT demonstration manifest](../dataset/sft_sd3_5/train.jsonl) with its
+ [SD3.5 config](../examples/sft/lora/sd3_5/default.yaml).
+- [Offline-DPO preference manifest](../dataset/offline_dpo_sd3_5/train.jsonl) with its
+ [SD3.5 config](../examples/offline_dpo/lora/sd3_5/default.yaml).
### Public offline smoke datasets
diff --git a/tests/examples/test_offline_examples.py b/tests/examples/test_offline_examples.py
new file mode 100644
index 000000000..32a27d93f
--- /dev/null
+++ b/tests/examples/test_offline_examples.py
@@ -0,0 +1,74 @@
+# Copyright 2026 Jayce-Ping
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+
+import pytest
+
+from flow_factory.data_utils.offline_dataset import OfflineSupervisionType, load_offline_manifest
+from flow_factory.data_utils.schema import DemonstrationSupervision
+from flow_factory.hparams import Arguments
+
+ROOT = Path(__file__).resolve().parents[2]
+OFFLINE_EXAMPLES = (
+ (
+ "examples/sft/lora/sd3_5/default.yaml",
+ "dataset/sft_sd3_5",
+ "demonstration",
+ ),
+ (
+ "examples/offline_dpo/lora/sd3_5/default.yaml",
+ "dataset/offline_dpo_sd3_5",
+ "preference",
+ ),
+)
+
+
+@pytest.mark.parametrize(
+ ("config_path", "expected_dataset_dir", "supervision_type"),
+ OFFLINE_EXAMPLES,
+)
+def test_offline_examples_use_repository_dataset_root(
+ config_path: str,
+ expected_dataset_dir: str,
+ supervision_type: OfflineSupervisionType,
+) -> None:
+ """Keep checked-in offline fixtures under the repository dataset root."""
+ config = Arguments.load_from_yaml(str(ROOT / config_path))
+
+ assert len(config.data_args.datasets) == 1
+ configured_dataset_dir = Path(config.data_args.datasets[0].dataset_dir)
+ assert configured_dataset_dir == Path(expected_dataset_dir)
+ assert configured_dataset_dir.parts[0] == "dataset"
+
+ dataset_dir = ROOT / configured_dataset_dir
+ records = load_offline_manifest(
+ dataset_dir / "train.jsonl",
+ supervision_type=supervision_type,
+ )
+ assert len(records) == 2
+
+ assets_dir = (ROOT / "assets").resolve()
+ for record in records:
+ supervision = record.supervision
+ candidates = (
+ (supervision.target,)
+ if isinstance(supervision, DemonstrationSupervision)
+ else (supervision.chosen, supervision.rejected)
+ )
+ for candidate in candidates:
+ for media in candidate.media:
+ media_path = Path(media.path).resolve()
+ assert media_path.is_relative_to(assets_dir)
+ assert media_path.is_file()
From 365380628e2a7cabb0ac81a6796f255b38076ad6 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 08:38:06 +0800
Subject: [PATCH 73/76] fix(data): unify dataset media discriminator as type
---
.agents/knowledge/topics/fix_patterns.md | 17 +++++
README.md | 4 +-
dataset/minimax_h3_ref2va/README.md | 6 +-
dataset/minimax_h3_ref2va/test.jsonl | 4 +-
dataset/minimax_h3_ref2va/train.jsonl | 4 +-
dataset/offline_smoke/build_mini.py | 12 ++--
guidance/datasets.md | 31 +++++----
guidance/new_model.md | 4 +-
src/flow_factory/data_utils/dataset.py | 30 ++++-----
.../data_utils/offline_condition_cache.py | 14 ++---
.../models/minimax_h3/adapters.py | 2 +-
.../models/minimax_h3/workflow.py | 12 ++--
src/flow_factory/samples/references.py | 24 +++----
.../test_offline_condition_cache.py | 14 ++---
tests/data_utils/test_offline_dataset.py | 4 +-
tests/data_utils/test_offline_train_data.py | 5 +-
tests/data_utils/test_ordered_references.py | 63 ++++++++++---------
tests/data_utils/test_schema.py | 4 +-
tests/examples/test_minimax_h3_examples.py | 4 +-
.../models/minimax_h3/test_condition_state.py | 2 +-
tests/models/minimax_h3/test_review_fixes.py | 6 +-
.../minimax_h3/test_workflow_execution.py | 20 +++---
.../test_reward_processor_reconstruction.py | 2 +-
.../samples/test_ordered_reference_samples.py | 24 +++----
tests/trainers/test_collective_packing.py | 2 +-
25 files changed, 172 insertions(+), 142 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index a58d3e92e..8eef41d75 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -754,6 +754,23 @@ Based on the fix type, write the fix entry to the appropriate document:
than checking only the configured directory string.
- **Related Constraint**: N/A
+### Dataset media discriminators must survive projections unchanged
+- **Date**: 2026-08-31
+- **Symptom**: MiniMax H3 Ref2VA examples used a different media discriminator from strict V2
+ records, and offline condition projection translated between the two representations before
+ adapter preprocessing.
+- **Root Cause**: Ordered-reference support introduced a private compatibility representation
+ instead of preserving the public `MediaAsset.type` contract across canonicalization, decoding,
+ cache projection, and adapter dispatch.
+- **Fix**: Online Ref2VA manifests, canonical reference sidecars, decoded entries, offline
+ projection, and MiniMax H3 dispatch now use `type` end to end. The condition-source and H3
+ preprocessing cache versions were advanced so incompatible Arrow caches are rebuilt, and the
+ dataset guide, fixtures, and contract tests follow the same schema.
+- **Lesson**: A projection may change storage shape, such as list-of-struct to canonical JSON, but
+ it should not rename semantic fields. Keep the public discriminator stable until the concrete
+ third-party object-construction boundary and version every cache that stores the old shape.
+- **Related Constraint**: #5
+
## Cross-refs
- UP: [Hard Constraints](../constraints.md), [Architecture](../architecture.md)
diff --git a/README.md b/README.md
index ccdf84ba7..2ce49eac1 100644
--- a/README.md
+++ b/README.md
@@ -229,8 +229,8 @@ The unified structure of dataset is:
## Offline SFT and Preference Data
-SFT and offline DPO use strict JSONL with `schema_version: 2`. Public media objects always use the
-`type` discriminator; `kind` is not accepted in V2:
+SFT and offline DPO use strict JSONL with `schema_version: 2`. Public media objects use `type` as
+their sole discriminator:
```jsonl
{"schema_version":2,"input":{"prompt":"A clean poster.","media":[]},"supervision":{"type":"demonstration","target":{"media":[{"type":"image","path":"targets/poster.png"}]}},"metadata":{}}
diff --git a/dataset/minimax_h3_ref2va/README.md b/dataset/minimax_h3_ref2va/README.md
index f6e8aa3fa..c85df671b 100644
--- a/dataset/minimax_h3_ref2va/README.md
+++ b/dataset/minimax_h3_ref2va/README.md
@@ -9,10 +9,10 @@ identity.
Supported entries:
-- `image`: `kind` and dataset-relative `path`;
-- `video`: `kind`, dataset-relative `path`, optional finite positive `fps`, and optional
+- `image`: `type` and dataset-relative `path`;
+- `video`: `type`, dataset-relative `path`, optional finite positive `fps`, and optional
dataset-relative `audio_path`;
-- `audio`: `kind`, dataset-relative `path`, and optional finite positive `sample_rate`.
+- `audio`: `type`, dataset-relative `path`, and optional finite positive `sample_rate`.
At least one image or video is required; audio-only manifests are invalid. A video `sample_rate`
is valid only when that video also supplies `audio_path`. Manifest `fps` and `sample_rate`
diff --git a/dataset/minimax_h3_ref2va/test.jsonl b/dataset/minimax_h3_ref2va/test.jsonl
index 46a791e4e..eb6d491b8 100644
--- a/dataset/minimax_h3_ref2va/test.jsonl
+++ b/dataset/minimax_h3_ref2va/test.jsonl
@@ -1,2 +1,2 @@
-{"prompt":"Compose a short audiovisual scene while preserving the ordered references.","references":[{"kind":"image","path":"references/style.png"},{"kind":"video","path":"references/motion.mp4","fps":12.0},{"kind":"audio","path":"references/ambience.wav","sample_rate":16000}]}
-{"prompt":"Use the clip motion first, its separate soundtrack, then the style image.","references":[{"kind":"video","path":"references/motion.mp4","fps":12.0,"audio_path":"references/soundtrack.wav","sample_rate":16000},{"kind":"image","path":"references/style.png"}]}
+{"prompt":"Compose a short audiovisual scene while preserving the ordered references.","references":[{"type":"image","path":"references/style.png"},{"type":"video","path":"references/motion.mp4","fps":12.0},{"type":"audio","path":"references/ambience.wav","sample_rate":16000}]}
+{"prompt":"Use the clip motion first, its separate soundtrack, then the style image.","references":[{"type":"video","path":"references/motion.mp4","fps":12.0,"audio_path":"references/soundtrack.wav","sample_rate":16000},{"type":"image","path":"references/style.png"}]}
diff --git a/dataset/minimax_h3_ref2va/train.jsonl b/dataset/minimax_h3_ref2va/train.jsonl
index b3200228c..dcc6c986c 100644
--- a/dataset/minimax_h3_ref2va/train.jsonl
+++ b/dataset/minimax_h3_ref2va/train.jsonl
@@ -1,2 +1,2 @@
-{"prompt":"Create a coherent scene using the references in order.","references":[{"kind":"image","path":"references/style.png"},{"kind":"video","path":"references/motion.mp4","fps":12.0},{"kind":"audio","path":"references/ambience.wav","sample_rate":16000}]}
-{"prompt":"Follow the reference motion and use the separately supplied soundtrack.","references":[{"kind":"video","path":"references/motion.mp4","fps":12.0,"audio_path":"references/soundtrack.wav","sample_rate":16000},{"kind":"image","path":"references/style.png"}]}
+{"prompt":"Create a coherent scene using the references in order.","references":[{"type":"image","path":"references/style.png"},{"type":"video","path":"references/motion.mp4","fps":12.0},{"type":"audio","path":"references/ambience.wav","sample_rate":16000}]}
+{"prompt":"Follow the reference motion and use the separately supplied soundtrack.","references":[{"type":"video","path":"references/motion.mp4","fps":12.0,"audio_path":"references/soundtrack.wav","sample_rate":16000},{"type":"image","path":"references/style.png"}]}
diff --git a/dataset/offline_smoke/build_mini.py b/dataset/offline_smoke/build_mini.py
index 6c5378f39..d9a58e655 100644
--- a/dataset/offline_smoke/build_mini.py
+++ b/dataset/offline_smoke/build_mini.py
@@ -421,14 +421,14 @@ def _candidate_assets(
pool = pools[family]
return {
"chosen": {
- kind: pool[f"chosen_{kind}"]
- for kind in ("video", "audio")
- if f"chosen_{kind}" in pool
+ media_type: pool[f"chosen_{media_type}"]
+ for media_type in ("video", "audio")
+ if f"chosen_{media_type}" in pool
},
"rejected": {
- kind: pool[f"rejected_{kind}"]
- for kind in ("video", "audio")
- if f"rejected_{kind}" in pool
+ media_type: pool[f"rejected_{media_type}"]
+ for media_type in ("video", "audio")
+ if f"rejected_{media_type}" in pool
},
}
pool = pools["image"]
diff --git a/guidance/datasets.md b/guidance/datasets.md
index a2a9acfdd..782fea196 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -84,9 +84,8 @@ remaining slots in declaration order. Duplicate, unknown, and wrong-media-type s
contract validation. Supervision outputs reject `slot` because their order is declared by the
pipeline output contract rather than by condition-argument names.
-Do not write `kind` in a V2 record. Some ordered-reference adapters still consume a validated
-legacy `kind` mapping internally; the V2 condition projection creates that private bridge only at
-the adapter preprocessing boundary. It is not part of the public V2 schema.
+The same `type` discriminator is preserved when V2 media are projected into ordered-reference
+adapter inputs, so public records and adapter preprocessing use one media-entry contract.
### Demonstration supervision
@@ -289,8 +288,8 @@ exact aligned latent duration. For example:
For LTX2 I2AV, bind one image to `first_frame`. H3 FL2VA supports first-only, last-only, and
first-plus-last records; use explicit slots for the last-only form. H3 Ref2VA puts the complete
-ordered image/video/audio reference sequence in `input.media`; the offline projection bridges those
-public `type` objects to the adapter's private legacy reference representation.
+ordered image/video/audio reference sequence in `input.media`; the offline projection preserves
+those ordered `type` entries at the adapter preprocessing boundary.
```jsonl
{"schema_version":2,"input":{"prompt":"Reveal the scene before this ending.","media":[{"type":"image","path":"conditions/end.png","slot":"last_frame"}]},"supervision":{"type":"demonstration","target":{"media":[{"type":"video","path":"targets/story.mp4","fps":24.0},{"type":"audio","path":"targets/story.wav","sample_rate":32000}]}},"metadata":{}}
@@ -457,11 +456,11 @@ directory. See the [FL2VA dataset fixture](../dataset/minimax_h3_fl2va/train.jso
### Ref2VA: `minimax-h3-ref2va`
-The existing online Ref2VA loader uses a legacy non-empty ordered `"references"` array containing
-image, video, and audio entries:
+The online Ref2VA loader uses a non-empty ordered `"references"` array containing image, video,
+and audio entries. Each entry uses the same `type` discriminator as strict V2 media objects:
```jsonl
-{"prompt":"Create a coherent scene using the references in order.","references":[{"kind":"image","path":"references/style.png"},{"kind":"video","path":"references/motion.mp4","fps":12.0},{"kind":"audio","path":"references/ambience.wav","sample_rate":16000}]}
+{"prompt":"Create a coherent scene using the references in order.","references":[{"type":"image","path":"references/style.png"},{"type":"video","path":"references/motion.mp4","fps":12.0},{"type":"audio","path":"references/ambience.wav","sample_rate":16000}]}
```
Array order is semantically significant. It is preserved during validation, encoding, caching, and
@@ -470,19 +469,19 @@ an audio-only array is invalid.
Supported entries:
-| `kind` | Required keys | Optional keys | Decoded value |
+| `type` | Required keys | Optional keys | Decoded value |
|---|---|---|---|
-| `image` | `kind`, `path` | none | RGB image |
-| `video` | `kind`, `path` | `fps`, `audio_path`, `sample_rate` | frames and optional soundtrack |
-| `audio` | `kind`, `path` | `sample_rate` | waveform |
+| `image` | `type`, `path` | none | RGB image |
+| `video` | `type`, `path` | `fps`, `audio_path`, `sample_rate` | frames and optional soundtrack |
+| `audio` | `type`, `path` | `sample_rate` | waveform |
`fps` and `sample_rate` overrides must be finite positive numbers. A video may use its embedded
soundtrack or a separate dataset-relative `audio_path`; a video `sample_rate` override requires
-`audio_path`. Unknown keys and unsupported legacy `kind` values fail before preprocessing.
+`audio_path`. Unknown keys and unsupported `type` values fail before preprocessing.
-This legacy online manifest is distinct from the strict V2 format above. A V2 record always uses
-`input.media[*].type`; offline condition projection performs any required legacy `kind` conversion
-internally.
+The compact online manifest and strict V2 format share the same media-entry discriminator. Offline
+condition projection preserves `input.media[*].type` when it constructs the ordered-reference
+manifest consumed by the adapter.
See the [Ref2VA dataset fixture](../dataset/minimax_h3_ref2va/train.jsonl), its
[local fixture notes](../dataset/minimax_h3_ref2va/README.md), and the
diff --git a/guidance/new_model.md b/guidance/new_model.md
index b3b948edc..4b8f8c0fa 100644
--- a/guidance/new_model.md
+++ b/guidance/new_model.md
@@ -633,8 +633,8 @@ and distillation continue to rely on their established trajectory-wide reduction
SenseNova is an example of an important boundary: its existing condition schema uses grouped
`images` with within-type order. Do not advertise heterogeneous ordered references merely because
-several images are accepted. The public V2 discriminator remains `type`; conversion to a legacy
-adapter-internal `kind` entry, when genuinely required, belongs only in the condition projection.
+several images are accepted. Dataset media and ordered-reference entries use `type` as their sole
+discriminator, including at the adapter preprocessing boundary.
## Advanced: Custom `preprocess_func`
diff --git a/src/flow_factory/data_utils/dataset.py b/src/flow_factory/data_utils/dataset.py
index d24a1dd5a..69330ba96 100644
--- a/src/flow_factory/data_utils/dataset.py
+++ b/src/flow_factory/data_utils/dataset.py
@@ -1245,7 +1245,7 @@ def _supports_ordered_references(preprocess_func: Optional[Callable]) -> bool:
def _canonicalize_ordered_reference_value(value: Any, row_index: int) -> str:
- """Canonicalize either a legacy reference list or an opaque Arrow string."""
+ """Canonicalize either an ordered reference list or an opaque Arrow string."""
if isinstance(value, str):
references = parse_reference_manifest(value, row_index=row_index)
else:
@@ -1319,14 +1319,14 @@ def _load_ordered_reference(
reference_index: int,
) -> Dict[str, Any]:
"""Decode one ordered reference with dataset row/reference context."""
- kind = entry["kind"]
+ reference_type = entry["type"]
resolved_path = _resolve_path(data_root, entry["path"])
failing_path = resolved_path
loaded = dict(entry)
try:
- if kind == "image":
+ if reference_type == "image":
loaded["media"] = Image.open(resolved_path).convert("RGB")
- elif kind == "video":
+ elif reference_type == "video":
frames, fps, audio, sample_rate = _decode_ordered_video(resolved_path)
effective_fps = entry.get("fps", fps)
_require_finite_positive_rate(
@@ -1334,7 +1334,7 @@ def _load_ordered_reference(
"effective fps",
row_index,
reference_index,
- kind,
+ reference_type,
resolved_path,
)
loaded["frames"] = frames
@@ -1349,7 +1349,7 @@ def _load_ordered_reference(
"sample_rate",
row_index,
reference_index,
- kind,
+ reference_type,
failing_path,
)
effective_sample_rate = entry.get("sample_rate", sample_rate)
@@ -1358,19 +1358,19 @@ def _load_ordered_reference(
"effective sample_rate",
row_index,
reference_index,
- kind,
+ reference_type,
failing_path,
)
loaded["audio"] = audio
loaded["sample_rate"] = effective_sample_rate
- elif kind == "audio":
+ elif reference_type == "audio":
audio, sample_rate = _decode_ordered_audio(resolved_path)
_require_finite_positive_rate(
sample_rate,
"sample_rate",
row_index,
reference_index,
- kind,
+ reference_type,
resolved_path,
)
effective_sample_rate = entry.get("sample_rate", sample_rate)
@@ -1379,19 +1379,21 @@ def _load_ordered_reference(
"effective sample_rate",
row_index,
reference_index,
- kind,
+ reference_type,
resolved_path,
)
loaded["media"] = audio
loaded["sample_rate"] = effective_sample_rate
else:
raise ValueError(
- "expected ordered reference kind in ('image', 'video', 'audio'), " f"got {kind!r}"
+ "expected ordered reference type in ('image', 'video', 'audio'), "
+ f"got {reference_type!r}"
)
except (FileNotFoundError, ImportError, OSError, RuntimeError, ValueError) as error:
raise ValueError(
f"failed to decode ordered reference at row {row_index}, "
- f"reference {reference_index}, kind={kind!r}, path={failing_path!r}: {error}"
+ f"reference {reference_index}, type={reference_type!r}, "
+ f"path={failing_path!r}: {error}"
) from error
return loaded
@@ -1401,7 +1403,7 @@ def _require_finite_positive_rate(
rate_name: str,
row_index: int,
reference_index: int,
- kind: str,
+ reference_type: str,
media_path: str,
) -> None:
if (
@@ -1411,7 +1413,7 @@ def _require_finite_positive_rate(
or value <= 0
):
raise ValueError(
- f"at row {row_index}, reference {reference_index}, kind={kind!r}, "
+ f"at row {row_index}, reference {reference_index}, type={reference_type!r}, "
f"path={media_path!r}, expected decoded {rate_name} to be finite positive, "
f"got {value!r}"
)
diff --git a/src/flow_factory/data_utils/offline_condition_cache.py b/src/flow_factory/data_utils/offline_condition_cache.py
index 94b85d4e6..9980435d6 100644
--- a/src/flow_factory/data_utils/offline_condition_cache.py
+++ b/src/flow_factory/data_utils/offline_condition_cache.py
@@ -50,7 +50,7 @@
)
from .schema import MediaAsset, NormalizedDatasetRecord
-_CONDITION_SOURCE_FORMAT = "flow-factory-offline-condition-v3"
+_CONDITION_SOURCE_FORMAT = "flow-factory-offline-condition-v4"
def project_offline_condition_dataset(
@@ -65,9 +65,9 @@ def project_offline_condition_dataset(
Grouped adapters receive ``prompt`` plus per-modality ``images``, ``videos``,
and ``audios`` columns. Ordered-reference adapters receive one canonical JSON
- string per row instead of an Arrow list-of-struct column. The string is
- restored to validated legacy ``kind`` entries only at the adapter preprocess
- boundary, avoiding Arrow's heterogeneous-struct null-key expansion.
+ string per row instead of an Arrow list-of-struct column. The string preserves
+ validated ``type`` entries through the adapter preprocess boundary, avoiding
+ Arrow's heterogeneous-struct null-key expansion.
"""
if not isinstance(ordered_references, bool):
raise TypeError(
@@ -148,7 +148,7 @@ def project_offline_condition_dataset(
if ordered_references:
columns["references"] = [
canonicalize_reference_manifest(
- [_to_legacy_reference(asset) for asset in record.model_input.media],
+ [_to_ordered_reference(asset) for asset in record.model_input.media],
row_index=index,
)
for index, record in enumerate(stable_records)
@@ -339,8 +339,8 @@ def build_offline_condition_cache(
return condition_cache
-def _to_legacy_reference(asset: MediaAsset) -> Dict[str, Any]:
- reference: Dict[str, Any] = {"kind": asset.type, "path": asset.path}
+def _to_ordered_reference(asset: MediaAsset) -> Dict[str, Any]:
+ reference: Dict[str, Any] = {"type": asset.type, "path": asset.path}
if asset.type == "video" and asset.fps is not None:
reference["fps"] = asset.fps
elif asset.type == "audio" and asset.sample_rate is not None:
diff --git a/src/flow_factory/models/minimax_h3/adapters.py b/src/flow_factory/models/minimax_h3/adapters.py
index 719d1c927..dca81020b 100644
--- a/src/flow_factory/models/minimax_h3/adapters.py
+++ b/src/flow_factory/models/minimax_h3/adapters.py
@@ -72,7 +72,7 @@
)
_H3_PREPROCESS_CACHE_FIELDS = frozenset({"height", "width", "num_frames"})
-_H3_PREPROCESS_CACHE_VERSION = "minimax-h3-v2"
+_H3_PREPROCESS_CACHE_VERSION = "minimax-h3-v3"
logger = setup_logger(__name__)
_H3_OPTIONAL_AUDIO_REFERENCE_FORMAT = MediaFormat(
type=MediaType.AUDIO,
diff --git a/src/flow_factory/models/minimax_h3/workflow.py b/src/flow_factory/models/minimax_h3/workflow.py
index b4fc9dd50..25bee6b56 100644
--- a/src/flow_factory/models/minimax_h3/workflow.py
+++ b/src/flow_factory/models/minimax_h3/workflow.py
@@ -653,20 +653,22 @@ def _build_pinned_references(entries: Sequence[Mapping[str, Any]]) -> List[Any]:
symbols = require_minimax_h3_support()
references = []
for entry in entries:
- kind = entry["kind"]
- if kind == "image":
+ reference_type = entry["type"]
+ if reference_type == "image":
references.append(symbols.ImageReference(image=entry["media"]))
- elif kind == "video":
+ elif reference_type == "video":
reference_kwargs = {"frames": entry["frames"], "fps": entry["fps"]}
if entry.get("audio") is not None:
reference_kwargs.update(audio=entry["audio"], sample_rate=entry["sample_rate"])
references.append(symbols.VideoReference(**reference_kwargs))
- elif kind == "audio":
+ elif reference_type == "audio":
references.append(
symbols.AudioReference(audio=entry["media"], sample_rate=entry["sample_rate"])
)
else:
- raise ValueError(f"expected image/video/audio reference kind, received {kind!r}")
+ raise ValueError(
+ f"expected image/video/audio reference type, received {reference_type!r}"
+ )
return references
diff --git a/src/flow_factory/samples/references.py b/src/flow_factory/samples/references.py
index c62612864..cb5dbf9ba 100644
--- a/src/flow_factory/samples/references.py
+++ b/src/flow_factory/samples/references.py
@@ -19,9 +19,9 @@
from typing import Any, Dict, List
_REFERENCE_KEYS = {
- "image": frozenset({"kind", "path"}),
- "video": frozenset({"kind", "path", "fps", "audio_path", "sample_rate"}),
- "audio": frozenset({"kind", "path", "sample_rate"}),
+ "image": frozenset({"type", "path"}),
+ "video": frozenset({"type", "path", "fps", "audio_path", "sample_rate"}),
+ "audio": frozenset({"type", "path", "sample_rate"}),
}
@@ -39,10 +39,10 @@ def canonicalize_reference_manifest(references: Any, row_index: int) -> str:
_validate_reference_entry(entry, row_index, reference_index)
for reference_index, entry in enumerate(references)
]
- if not any(entry["kind"] in ("image", "video") for entry in validated):
+ if not any(entry["type"] in ("image", "video") for entry in validated):
raise ValueError(
f"at row {row_index}, reference 0, expected at least one image or video "
- f"reference, got audio-only kinds={[entry['kind'] for entry in validated]!r}"
+ f"reference, got audio-only types={[entry['type'] for entry in validated]!r}"
)
return json.dumps(validated, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
@@ -69,11 +69,11 @@ def _validate_reference_entry(
f"expected object at row {row_index}, reference {reference_index}, "
f"got {type(entry).__name__}: {entry!r}"
)
- kind = entry.get("kind")
- if not isinstance(kind, str) or kind not in _REFERENCE_KEYS:
+ reference_type = entry.get("type")
+ if not isinstance(reference_type, str) or reference_type not in _REFERENCE_KEYS:
raise ValueError(
- f"at row {row_index}, reference {reference_index}, expected kind in "
- f"{tuple(_REFERENCE_KEYS)}, got {kind!r}"
+ f"at row {row_index}, reference {reference_index}, expected type in "
+ f"{tuple(_REFERENCE_KEYS)}, got {reference_type!r}"
)
path = entry.get("path")
if not isinstance(path, str) or not path:
@@ -82,11 +82,11 @@ def _validate_reference_entry(
f"a non-empty string, got {path!r}"
)
- unknown_keys = set(entry) - _REFERENCE_KEYS[kind]
+ unknown_keys = set(entry) - _REFERENCE_KEYS[reference_type]
if unknown_keys:
raise ValueError(
f"at row {row_index}, reference {reference_index}, unknown keys "
- f"{sorted(unknown_keys)} for kind {kind!r}"
+ f"{sorted(unknown_keys)} for type {reference_type!r}"
)
for rate_name in ("fps", "sample_rate"):
if rate_name in entry:
@@ -101,7 +101,7 @@ def _validate_reference_entry(
f"at row {row_index}, reference {reference_index}, expected "
f"{rate_name} to be finite positive numeric, got {rate!r}"
)
- if kind == "video":
+ if reference_type == "video":
audio_path = entry.get("audio_path")
if audio_path is not None and (not isinstance(audio_path, str) or not audio_path):
raise ValueError(
diff --git a/tests/data_utils/test_offline_condition_cache.py b/tests/data_utils/test_offline_condition_cache.py
index 9f227046c..260e18594 100644
--- a/tests/data_utils/test_offline_condition_cache.py
+++ b/tests/data_utils/test_offline_condition_cache.py
@@ -851,11 +851,11 @@ def test_ordered_heterogeneous_references_cross_arrow_as_canonical_json(
assert isinstance(raw_manifest, str)
assert [set(reference) for reference in raw_references] == [
- {"kind", "path"},
- {"kind", "path", "fps"},
- {"kind", "path", "sample_rate"},
+ {"type", "path"},
+ {"type", "path", "fps"},
+ {"type", "path", "sample_rate"},
]
- assert [reference["kind"] for reference in raw_references] == [
+ assert [reference["type"] for reference in raw_references] == [
"image",
"video",
"audio",
@@ -872,9 +872,9 @@ def test_ordered_heterogeneous_references_cross_arrow_as_canonical_json(
)
loaded = preprocessor.references[0]
- assert set(loaded[0]) == {"kind", "path", "media"}
- assert set(loaded[1]) == {"kind", "path", "fps", "frames"}
- assert set(loaded[2]) == {"kind", "path", "sample_rate", "media"}
+ assert set(loaded[0]) == {"type", "path", "media"}
+ assert set(loaded[1]) == {"type", "path", "fps", "frames"}
+ assert set(loaded[2]) == {"type", "path", "sample_rate", "media"}
assert cache[0][OFFLINE_CONDITION_ID_COLUMN] == compute_offline_condition_id(
record,
index=0,
diff --git a/tests/data_utils/test_offline_dataset.py b/tests/data_utils/test_offline_dataset.py
index 23ce8b16f..849245e4f 100644
--- a/tests/data_utils/test_offline_dataset.py
+++ b/tests/data_utils/test_offline_dataset.py
@@ -232,8 +232,8 @@ def test_manifest_reader_preserves_order_resolves_paths_and_requires_one_type(
{
"schema_version": 2,
"input": {
- "prompt": "legacy key",
- "media": [{"kind": "image", "path": "input.png"}],
+ "prompt": "unknown media key",
+ "media": [{"type": "image", "path": "input.png", "media_type": "image"}],
},
"supervision": {
"type": "demonstration",
diff --git a/tests/data_utils/test_offline_train_data.py b/tests/data_utils/test_offline_train_data.py
index de53c7310..8e0f5c048 100644
--- a/tests/data_utils/test_offline_train_data.py
+++ b/tests/data_utils/test_offline_train_data.py
@@ -488,7 +488,7 @@ def test_builder_supports_homogeneous_offline_preference_sources(tmp_path: Path)
assert "rejected.png" not in repr(dataset._condition_cache[0])
-def test_builder_uses_bridge_ordered_reference_boundary_with_single_row_batches(
+def test_builder_preserves_ordered_reference_type_with_single_row_batches(
tmp_path: Path,
) -> None:
dataset_dir = tmp_path / "ordered"
@@ -510,8 +510,7 @@ def test_builder_uses_bridge_ordered_reference_boundary_with_single_row_batches(
)
assert preprocessor.references is not None
- assert preprocessor.references[0][0]["kind"] == "image"
- assert "type" not in preprocessor.references[0][0]
+ assert preprocessor.references[0][0]["type"] == "image"
(dataset,) = _source_datasets(loader)
assert dataset[0].model_input.media[0].type == "image"
diff --git a/tests/data_utils/test_ordered_references.py b/tests/data_utils/test_ordered_references.py
index 6cc984eac..33f64f21a 100644
--- a/tests/data_utils/test_ordered_references.py
+++ b/tests/data_utils/test_ordered_references.py
@@ -27,9 +27,9 @@
from flow_factory.models.minimax_h3.adapters import MiniMaxH3Ref2VAAdapter
REFERENCES = [
- {"kind": "image", "path": "subject.png"},
- {"kind": "video", "path": "motion.mp4", "fps": 29.97},
- {"kind": "audio", "path": "voice.wav", "sample_rate": 44100},
+ {"type": "image", "path": "subject.png"},
+ {"type": "video", "path": "motion.mp4", "fps": 29.97},
+ {"type": "audio", "path": "voice.wav", "sample_rate": 44100},
]
@@ -125,8 +125,8 @@ def _write_video(path: Path, with_audio: bool) -> None:
def test_ordered_references_round_trip_real_media_and_merged_cache(tmp_path: Path) -> None:
dataset_dir = tmp_path / "dataset"
references = [
- {"kind": "image", "path": "subject.png"},
- {"kind": "audio", "path": "voice.wav", "sample_rate": 22050},
+ {"type": "image", "path": "subject.png"},
+ {"type": "audio", "path": "voice.wav", "sample_rate": 22050},
]
_write_jsonl(dataset_dir, references)
Image.new("RGB", (4, 3), color=(12, 34, 56)).save(dataset_dir / "subject.png")
@@ -142,7 +142,7 @@ def test_ordered_references_round_trip_real_media_and_merged_cache(tmp_path: Pat
)
assert len(preprocessor.received) == 1
- assert [entry["kind"] for entry in preprocessor.received[0]] == ["image", "audio"]
+ assert [entry["type"] for entry in preprocessor.received[0]] == ["image", "audio"]
assert preprocessor.received[0][0]["media"].size == (4, 3)
assert preprocessor.received[0][1]["sample_rate"] == 22050
assert preprocessor.received[0][1]["media"].shape[0] == 1
@@ -162,7 +162,7 @@ def test_ordered_references_round_trip_real_media_and_merged_cache(tmp_path: Pat
def test_video_reference_preserves_frames_fps_embedded_audio_and_rate(tmp_path: Path) -> None:
dataset_dir = tmp_path / "dataset"
- references = [{"kind": "video", "path": "motion.mp4"}]
+ references = [{"type": "video", "path": "motion.mp4"}]
_write_jsonl(dataset_dir, references)
video_path = dataset_dir / "motion.mp4"
_write_video(video_path, with_audio=True)
@@ -188,8 +188,8 @@ def test_reference_decode_error_has_row_reference_and_cause(tmp_path: Path) -> N
dataset_dir = tmp_path / "dataset"
dataset_dir.mkdir()
rows = [
- {"prompt": "valid", "references": [{"kind": "image", "path": "valid.png"}]},
- {"prompt": "missing", "references": [{"kind": "image", "path": "missing.png"}]},
+ {"prompt": "valid", "references": [{"type": "image", "path": "valid.png"}]},
+ {"prompt": "missing", "references": [{"type": "image", "path": "missing.png"}]},
]
(dataset_dir / "train.jsonl").write_text(
"".join(json.dumps(row) + "\n" for row in rows),
@@ -218,7 +218,7 @@ def test_soundtrack_decode_error_reports_soundtrack_path_and_context(tmp_path: P
dataset_dir = tmp_path / "dataset"
references = [
{
- "kind": "video",
+ "type": "video",
"path": "motion.mp4",
"audio_path": "missing-soundtrack.wav",
}
@@ -229,7 +229,7 @@ def test_soundtrack_decode_error_reports_soundtrack_path_and_context(tmp_path: P
with pytest.raises(
ValueError,
- match=r"row 0.*reference 0.*kind='video'.*missing-soundtrack\.wav",
+ match=r"row 0.*reference 0.*type='video'.*missing-soundtrack\.wav",
) as caught:
GeneralDataset(
dataset_dir=str(dataset_dir),
@@ -244,7 +244,7 @@ def test_soundtrack_decode_error_reports_soundtrack_path_and_context(tmp_path: P
@pytest.mark.parametrize(
- ("kind", "decode_target", "decode_result", "rate_name"),
+ ("media_type", "decode_target", "decode_result", "rate_name"),
[
(
"video",
@@ -269,14 +269,17 @@ def test_soundtrack_decode_error_reports_soundtrack_path_and_context(tmp_path: P
def test_decoded_rates_must_be_finite_positive_with_context(
tmp_path: Path,
monkeypatch,
- kind: str,
+ media_type: str,
decode_target: str,
decode_result: Any,
rate_name: str,
) -> None:
- dataset_dir = tmp_path / f"dataset-{kind}-{rate_name}"
- path = "media.mp4" if kind == "video" else "media.wav"
- references = [{"kind": "image", "path": "valid.png"}, {"kind": kind, "path": path}]
+ dataset_dir = tmp_path / f"dataset-{media_type}-{rate_name}"
+ path = "media.mp4" if media_type == "video" else "media.wav"
+ references = [
+ {"type": "image", "path": "valid.png"},
+ {"type": media_type, "path": path},
+ ]
_write_jsonl(dataset_dir, references)
Image.new("RGB", (2, 2)).save(dataset_dir / "valid.png")
monkeypatch.setattr(decode_target, lambda media_path: decode_result)
@@ -284,7 +287,7 @@ def test_decoded_rates_must_be_finite_positive_with_context(
with pytest.raises(
ValueError,
- match=rf"row 0.*reference 1.*kind='{kind}'.*{rate_name}.*finite positive",
+ match=rf"row 0.*reference 1.*type='{media_type}'.*{rate_name}.*finite positive",
):
GeneralDataset(
dataset_dir=str(dataset_dir),
@@ -304,7 +307,7 @@ def test_video_uses_manifest_fps_when_decoder_has_no_rate(tmp_path: Path, monkey
)
loaded = _load_ordered_reference(
- {"kind": "video", "path": "media.mp4", "fps": 24.0},
+ {"type": "video", "path": "media.mp4", "fps": 24.0},
data_root=str(tmp_path),
row_index=3,
reference_index=4,
@@ -323,10 +326,10 @@ def test_video_without_override_rejects_missing_decoded_fps(tmp_path: Path, monk
with pytest.raises(
ValueError,
- match=r"row 3.*reference 4.*kind='video'.*effective fps.*finite positive.*None",
+ match=r"row 3.*reference 4.*type='video'.*effective fps.*finite positive.*None",
):
_load_ordered_reference(
- {"kind": "video", "path": "media.mp4"},
+ {"type": "video", "path": "media.mp4"},
data_root=str(tmp_path),
row_index=3,
reference_index=4,
@@ -335,7 +338,7 @@ def test_video_without_override_rejects_missing_decoded_fps(tmp_path: Path, monk
@pytest.mark.parametrize("rate", [float("nan"), float("inf"), float("-inf"), 0.0, -1.0])
@pytest.mark.parametrize(
- ("kind", "rate_name", "decode_target", "decode_result"),
+ ("media_type", "rate_name", "decode_target", "decode_result"),
[
(
"video",
@@ -355,17 +358,17 @@ def test_effective_override_rates_must_be_finite_positive(
tmp_path: Path,
monkeypatch,
rate: float,
- kind: str,
+ media_type: str,
rate_name: str,
decode_target: str,
decode_result: Any,
) -> None:
monkeypatch.setattr(decode_target, lambda media_path: decode_result)
- entry = {"kind": kind, "path": f"media.{kind}", rate_name: rate}
+ entry = {"type": media_type, "path": f"media.{media_type}", rate_name: rate}
with pytest.raises(
ValueError,
- match=rf"row 3.*reference 4.*kind='{kind}'.*effective {rate_name}.*finite positive",
+ match=rf"row 3.*reference 4.*type='{media_type}'.*effective {rate_name}.*finite positive",
):
_load_ordered_reference(
entry,
@@ -474,14 +477,14 @@ def fingerprint(preprocess_kwargs: Dict[str, Any]) -> str:
"changed_references",
[
[
- {"kind": "image", "path": "subject.png"},
- {"kind": "video", "path": "motion.mp4", "fps": 24.0},
- {"kind": "audio", "path": "voice.wav", "sample_rate": 44100},
+ {"type": "image", "path": "subject.png"},
+ {"type": "video", "path": "motion.mp4", "fps": 24.0},
+ {"type": "audio", "path": "voice.wav", "sample_rate": 44100},
],
[
- {"kind": "image", "path": "subject.png"},
- {"kind": "video", "path": "motion.mp4", "fps": 29.97},
- {"kind": "audio", "path": "voice.wav", "sample_rate": 48000},
+ {"type": "image", "path": "subject.png"},
+ {"type": "video", "path": "motion.mp4", "fps": 29.97},
+ {"type": "audio", "path": "voice.wav", "sample_rate": 48000},
],
list(reversed(REFERENCES)),
],
diff --git a/tests/data_utils/test_schema.py b/tests/data_utils/test_schema.py
index 9c09f5a30..128b936ad 100644
--- a/tests/data_utils/test_schema.py
+++ b/tests/data_utils/test_schema.py
@@ -195,11 +195,11 @@ def test_normalization_expands_dataset_root_and_normalizes_absolute_paths(
@pytest.mark.parametrize(
"media",
[
- {"kind": "image", "path": "image.png"},
+ {"type": "image", "path": "image.png", "media_type": "image"},
{"type": "image", "path": "image.png", "unknown": True},
],
)
-def test_v2_media_rejects_legacy_kind_and_unknown_keys(media: Dict[str, Any]) -> None:
+def test_v2_media_rejects_unknown_keys(media: Dict[str, Any]) -> None:
raw = _demonstration_record(input={"prompt": "strict", "media": [media]})
with pytest.raises(ValidationError):
diff --git a/tests/examples/test_minimax_h3_examples.py b/tests/examples/test_minimax_h3_examples.py
index e56f19d56..c77014b8c 100644
--- a/tests/examples/test_minimax_h3_examples.py
+++ b/tests/examples/test_minimax_h3_examples.py
@@ -165,7 +165,7 @@ def test_ref2va_manifests_are_ordered_valid_and_dataset_relative() -> None:
for row_index, row in enumerate(rows):
references = row["references"]
canonical = json.loads(canonicalize_reference_manifest(references, row_index=row_index))
- assert [entry["kind"] for entry in canonical] == [entry["kind"] for entry in references]
+ assert [entry["type"] for entry in canonical] == [entry["type"] for entry in references]
for reference in references:
path = Path(reference["path"])
assert not path.is_absolute()
@@ -174,7 +174,7 @@ def test_ref2va_manifests_are_ordered_valid_and_dataset_relative() -> None:
audio_path = Path(reference["audio_path"])
assert not audio_path.is_absolute()
assert (dataset_dir / audio_path).is_file()
- assert [entry["kind"] for entry in rows[0]["references"]] == [
+ assert [entry["type"] for entry in rows[0]["references"]] == [
"image",
"video",
"audio",
diff --git a/tests/models/minimax_h3/test_condition_state.py b/tests/models/minimax_h3/test_condition_state.py
index 57d708283..4077ca32b 100644
--- a/tests/models/minimax_h3/test_condition_state.py
+++ b/tests/models/minimax_h3/test_condition_state.py
@@ -59,7 +59,7 @@ def test_conditioned_adapters_declare_preparer_but_t2va_keeps_identity() -> None
MiniMaxH3T2VAAdapter.preprocess_cache_version,
MiniMaxH3FL2VAAdapter.preprocess_cache_version,
MiniMaxH3Ref2VAAdapter.preprocess_cache_version,
- } == {"minimax-h3-v2"}
+ } == {"minimax-h3-v3"}
def test_ref2va_preparer_realizes_prefix_once_and_routes_owned_contexts(
diff --git a/tests/models/minimax_h3/test_review_fixes.py b/tests/models/minimax_h3/test_review_fixes.py
index 41ce8d9ec..27087dc16 100644
--- a/tests/models/minimax_h3/test_review_fixes.py
+++ b/tests/models/minimax_h3/test_review_fixes.py
@@ -119,13 +119,13 @@ def test_sparse_indices_reject_invalid_type_range_and_duplicates(indices: Any) -
("adapter_class", "extra"),
[
(MiniMaxH3T2VAAdapter, {"images": [["frame"]]}),
- (MiniMaxH3T2VAAdapter, {"references": [[{"kind": "image"}]]}),
+ (MiniMaxH3T2VAAdapter, {"references": [[{"type": "image"}]]}),
(MiniMaxH3FL2VAAdapter, {"videos": [["video"]], "images": [["frame"]]}),
(MiniMaxH3FL2VAAdapter, {"images": [[], []]}),
- (MiniMaxH3Ref2VAAdapter, {"images": [["frame"]], "references": [[{"kind": "image"}]]}),
+ (MiniMaxH3Ref2VAAdapter, {"images": [["frame"]], "references": [[{"type": "image"}]]}),
(
MiniMaxH3Ref2VAAdapter,
- {"audios": [[torch.zeros(1)]], "references": [[{"kind": "image"}]]},
+ {"audios": [[torch.zeros(1)]], "references": [[{"type": "image"}]]},
),
],
)
diff --git a/tests/models/minimax_h3/test_workflow_execution.py b/tests/models/minimax_h3/test_workflow_execution.py
index 334937956..3719af628 100644
--- a/tests/models/minimax_h3/test_workflow_execution.py
+++ b/tests/models/minimax_h3/test_workflow_execution.py
@@ -197,8 +197,10 @@ def test_preprocess_adds_outer_batch_to_arrow_cache_fields(monkeypatch) -> None:
def test_ref_preprocess_builds_ordered_pinned_objects_without_returning_them(monkeypatch) -> None:
constructed: List[Any] = []
- def reference_type(kind: str):
- return lambda **kwargs: constructed.append((kind, kwargs)) or SimpleNamespace(kind=kind)
+ def reference_type(media_type: str):
+ return lambda **kwargs: constructed.append((media_type, kwargs)) or SimpleNamespace(
+ media_type=media_type
+ )
monkeypatch.setattr(
"flow_factory.models.minimax_h3.workflow.require_minimax_h3_support",
@@ -217,16 +219,16 @@ def reference_type(kind: str):
)
adapter = _adapter(MiniMaxH3Ref2VAAdapter)
references = [
- {"kind": "image", "path": "i.png", "media": "image"},
+ {"type": "image", "path": "i.png", "media": "image"},
{
- "kind": "video",
+ "type": "video",
"path": "v.mp4",
"frames": "frames",
"fps": 24.0,
"audio": torch.zeros(2, 8),
"sample_rate": 32000,
},
- {"kind": "audio", "path": "a.wav", "media": torch.ones(1, 8), "sample_rate": 16000},
+ {"type": "audio", "path": "a.wav", "media": torch.ones(1, 8), "sample_rate": 16000},
]
result = adapter.preprocess_func(
@@ -238,10 +240,14 @@ def reference_type(kind: str):
num_frames=5,
)
- assert [kind for kind, _ in constructed] == ["image", "video", "audio"]
+ assert [media_type for media_type, _ in constructed] == ["image", "video", "audio"]
assert constructed[1][1]["frames"] == "frames"
assert "video" not in constructed[1][1]
- assert [ref.kind for ref in encoded_inputs["references"]] == ["image", "video", "audio"]
+ assert [ref.media_type for ref in encoded_inputs["references"]] == [
+ "image",
+ "video",
+ "audio",
+ ]
assert result["reference_manifest"] == ["manifest"]
assert "references" not in result
assert all(not isinstance(value, SimpleNamespace) for value in result.values())
diff --git a/tests/rewards/test_reward_processor_reconstruction.py b/tests/rewards/test_reward_processor_reconstruction.py
index f9fd45574..10abf26a8 100644
--- a/tests/rewards/test_reward_processor_reconstruction.py
+++ b/tests/rewards/test_reward_processor_reconstruction.py
@@ -49,7 +49,7 @@ def test_distributed_group_reward_preserves_sample_reconstruction_fields() -> No
)
sample = MiniMaxH3Ref2VASample(
prompt="A reference-conditioned prompt",
- reference_manifest='[{"kind":"image","path":"condition.png"}]',
+ reference_manifest='[{"type":"image","path":"condition.png"}]',
)
rewards = processor.compute_rewards([sample], store_to_samples=False)
diff --git a/tests/samples/test_ordered_reference_samples.py b/tests/samples/test_ordered_reference_samples.py
index 1e6f3cf22..7b0992f8d 100644
--- a/tests/samples/test_ordered_reference_samples.py
+++ b/tests/samples/test_ordered_reference_samples.py
@@ -21,20 +21,20 @@
def test_reference_manifest_preserves_order_and_canonicalizes_keys() -> None:
references = [
- {"path": "style.png", "kind": "image"},
+ {"path": "style.png", "type": "image"},
{
- "kind": "video",
+ "type": "video",
"path": "motion.mp4",
"fps": 24,
"audio_path": "sound.wav",
"sample_rate": 48000,
},
- {"kind": "audio", "path": "ambience.wav", "sample_rate": 32000},
+ {"type": "audio", "path": "ambience.wav", "sample_rate": 32000},
]
manifest = canonicalize_reference_manifest(references, row_index=3)
- assert [entry["kind"] for entry in parse_reference_manifest(manifest, 3)] == [
+ assert [entry["type"] for entry in parse_reference_manifest(manifest, 3)] == [
"image",
"video",
"audio",
@@ -51,14 +51,14 @@ def test_reference_sample_identity_includes_ordered_manifest() -> None:
first = Ref2AVSample(
prompt="animate",
reference_manifest=canonicalize_reference_manifest(
- [{"kind": "image", "path": "first.png"}],
+ [{"type": "image", "path": "first.png"}],
0,
),
)
second = Ref2AVSample(
prompt="animate",
reference_manifest=canonicalize_reference_manifest(
- [{"kind": "image", "path": "second.png"}],
+ [{"type": "image", "path": "second.png"}],
0,
),
)
@@ -70,12 +70,14 @@ def test_reference_sample_identity_includes_ordered_manifest() -> None:
"references,match",
[
([], "non-empty"),
- ([{"kind": "audio", "path": "only.wav"}], "image or video"),
- ([{"kind": "image", "path": ""}], "non-empty string"),
- ([{"kind": "image", "path": "x.png", "fps": 24}], "unknown keys"),
- ([{"kind": "video", "path": "x.mp4", "fps": float("nan")}], "finite positive"),
+ ([{"path": "missing-type.png"}], "expected type"),
+ ([{"type": "document", "path": "unsupported.pdf"}], "expected type"),
+ ([{"type": "audio", "path": "only.wav"}], "image or video"),
+ ([{"type": "image", "path": ""}], "non-empty string"),
+ ([{"type": "image", "path": "x.png", "fps": 24}], "unknown keys"),
+ ([{"type": "video", "path": "x.mp4", "fps": float("nan")}], "finite positive"),
(
- [{"kind": "video", "path": "x.mp4", "sample_rate": 32000}],
+ [{"type": "video", "path": "x.mp4", "sample_rate": 32000}],
"requires audio_path",
),
],
diff --git a/tests/trainers/test_collective_packing.py b/tests/trainers/test_collective_packing.py
index 75bd90bc8..7ff9766b3 100644
--- a/tests/trainers/test_collective_packing.py
+++ b/tests/trainers/test_collective_packing.py
@@ -117,7 +117,7 @@ def test_gather_samples_packs_same_dtype_fields_and_preserves_other_fields():
def test_gather_samples_preserves_concrete_reconstruction_fields() -> None:
accelerator = GatherRecorder()
- manifest = '[{"kind":"image","path":"condition.png"}]'
+ manifest = '[{"path":"condition.png","type":"image"}]'
sample = MiniMaxH3Ref2VASample(
prompt="A reference-conditioned prompt",
reference_manifest=manifest,
From 826cf599de8f4d79ee29c9792e1ff9139c6b1f20 Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 08:41:12 +0800
Subject: [PATCH 74/76] test(docs): align H3 parameter assertion with main
---
.agents/knowledge/topics/fix_patterns.md | 13 +++++++++++++
tests/docs/test_minimax_h3_docs.py | 4 ++--
2 files changed, 15 insertions(+), 2 deletions(-)
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 8eef41d75..66f7899d1 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -771,6 +771,19 @@ Based on the fix type, write the fix entry to the appropriate document:
third-party object-construction boundary and version every cache that stores the old shape.
- **Related Constraint**: #5
+### Rebase conflict resolution includes downstream assertions
+- **Date**: 2026-08-31
+- **Symptom**: After rebasing onto a README correction for the MiniMax H3 parameter count, the
+ merged table was accurate but a child-branch documentation test still required the superseded
+ value.
+- **Root Cause**: The factual conflict was resolved in the edited document, while its regression
+ assertion lived in a cleanly rebased file and therefore received no conflict marker.
+- **Fix**: The documentation assertion now checks the corrected `33B` value inherited from main.
+- **Lesson**: Conflict markers identify overlapping text, not the complete semantic impact of a
+ rebase. After resolving a factual conflict, search for and run downstream assertions that encode
+ the same fact even when Git reports those files as clean.
+- **Related Constraint**: N/A
+
## Cross-refs
- UP: [Hard Constraints](../constraints.md), [Architecture](../architecture.md)
diff --git a/tests/docs/test_minimax_h3_docs.py b/tests/docs/test_minimax_h3_docs.py
index edbe2af12..d166adbcd 100644
--- a/tests/docs/test_minimax_h3_docs.py
+++ b/tests/docs/test_minimax_h3_docs.py
@@ -37,7 +37,7 @@ def test_readme_documents_h3_links_dependency_and_limits() -> None:
assert "all 36 real-weight smoke cells" in text
assert "FL2VA first-plus-last" in text
assert "T2VA is real-weight validated on 1 and 16 GPUs" not in text
- assert text.count("30B | ") == 3
+ assert text.count("33B | ") == 3
for model_type, link in zip(
("minimax-h3-t2va", "minimax-h3-fl2va", "minimax-h3-ref2va"),
@@ -58,7 +58,7 @@ def test_readme_documents_h3_links_dependency_and_limits() -> None:
"data-ward velocity",
"N transitions",
"N + 1 states",
- "30B",
+ "33B",
"do not claim a completed long-run reward trend",
"[Datasets](guidance/datasets.md)",
):
From cc319afae10931a9b09c1c1bbb68415ba9d7860c Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 21:29:27 +0800
Subject: [PATCH 75/76] [deps] chore: require validated Muon runtime stack
---
.agents/knowledge/constraints.md | 4 ++--
.agents/knowledge/dependencies.md | 29 ++++++++++++++++--------
.agents/knowledge/topics/fix_patterns.md | 2 +-
AGENTS.md | 2 +-
README.md | 9 ++++++--
config/accelerate_configs/fsdp2.yaml | 4 ++--
docker/README.md | 2 +-
docker/docker-cuda/Dockerfile | 9 ++++----
pyproject.toml | 18 +++++++--------
scripts/install_geneval_deps.sh | 8 +++----
src/flow_factory/rewards/ocr.py | 16 +++++--------
src/flow_factory/utils/audio.py | 2 +-
12 files changed, 58 insertions(+), 47 deletions(-)
diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md
index 77fcc746b..964c943dc 100644
--- a/.agents/knowledge/constraints.md
+++ b/.agents/knowledge/constraints.md
@@ -203,7 +203,7 @@ The adapter resolves `component_load_dtypes` at native load/materialization time
## Code Quality (21–27)
### 21. Formatting Standards
-- **Black** with `line-length=100`, targeting Python 3.10–3.12
+- **Black** with `line-length=100`, targeting Python 3.11–3.12
- **isort** with `profile="black"`, `line_length=100`
- Comments and docstrings in **English**
@@ -214,7 +214,7 @@ The adapter resolves `component_load_dtypes` at native load/materialization time
- **Top-level imports only**: All `import` / `from ... import ...` statements MUST live at the top of the module, never inside function bodies, methods, `__init__`, or conditional branches. Sanctioned exceptions: (a) optional dependencies wrapped in `try/except ImportError` (e.g., `deepspeed`, `xformers`); (b) backend-gated imports where the target symbol is only resolvable under a specific runtime backend already selected by a preceding feature check (e.g., DeepSpeed/FSDP submodules guarded by `is_deepspeed()` / `is_fsdp2()` in `models/abc.py`); (c) genuine unresolvable circular imports documented inline. Lazy imports added merely for "import speed" or "to keep the module light" are NOT acceptable — every hard dependency already runs through Python's import machinery on a typical import path. Inline imports hide the dependency surface from readers, `isort`, and static-analysis tools, and re-execute on every call in hot loops.
### 23. Type Annotations
-All public methods must have type annotations. Use `typing` module types (`List`, `Dict`, `Optional`, `Tuple`, `Union`) for Python 3.10 compatibility.
+All public methods must have type annotations. Use `typing` module types (`List`, `Dict`, `Optional`, `Tuple`, `Union`) for Python 3.11 compatibility.
### 24. License Header
All source files must include the Apache 2.0 license header with `Copyright 2026 Jayce-Ping`.
diff --git a/.agents/knowledge/dependencies.md b/.agents/knowledge/dependencies.md
index 79198cd3e..a9e25ed11 100644
--- a/.agents/knowledge/dependencies.md
+++ b/.agents/knowledge/dependencies.md
@@ -20,11 +20,11 @@ pip install -e ".[deepspeed]" # Core + DeepSpeed only
pyproject.toml
├── [project.dependencies] Core deps (always installed)
├── [project.optional-dependencies]
-│ ├── deepspeed DeepSpeed >= 0.15.4
+│ ├── deepspeed DeepSpeed >= 0.18.3
│ ├── quantization bitsandbytes >= 0.45.3
│ ├── wandb Weights & Biases tracking
│ ├── swanlab SwanLab tracking
-│ ├── nvidia xformers, nvidia-ml-py
+│ ├── nvidia xformers >= 0.0.35, nvidia-ml-py
│ ├── bagel flash-attn, opencv-python
│ ├── geneval mmcv, mmengine, mmdet, open_clip_torch
│ ├── geneval2 scipy (exact GenEval2 GM parity)
@@ -36,14 +36,15 @@ pyproject.toml
The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ packages). Key ones:
-| Package | Min Version | Purpose |
+| Package | Supported Version | Purpose |
|---------|-------------|---------|
-| `torch` | >= 2.6.0 | PyTorch core and AdamW baseline; Muon needs the optional API noted below |
-| `torchvision` | >= 0.19.0 | Vision utilities |
-| `torchaudio` | >= 2.4.0 | Audio I/O (audio / audio-video models, CLAP) |
+| `torch` | >= 2.10.0 | PyTorch core, AdamW, and the native `torch.optim.Muon` API |
+| `torchvision` | >= 0.25.0 | Vision utilities; paired baseline for PyTorch 2.10 |
+| `torchaudio` | >= 2.10.0 | Audio I/O (audio / audio-video models, CLAP) |
+| `torchcodec` | >= 0.10.0 | Runtime used by TorchAudio 2.10 `load` / `save` APIs |
| `transformers` | >= 4.57.1 | Text encoders, tokenizers |
| `diffusers` | >= 0.40.0 | Diffusion pipelines, schedulers, MiniMax H3 and LTX2 APIs |
-| `accelerate` | >= 1.11.0 | Distributed training, mixed precision |
+| `accelerate` | >= 1.14.0 | Distributed training, mixed precision, and max reduction |
| `peft` | >= 0.17.0 | LoRA, parameter-efficient fine-tuning |
| `datasets` | >= 3.3.2 | Dataset loading |
| `huggingface-hub` | >= 0.35.3 | Model/dataset downloads |
@@ -57,12 +58,18 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
- DeepSpeed is optional — Accelerate alone handles most distributed scenarios.
### Muon
-- The core `torch>=2.6.0` floor does not guarantee `torch.optim.Muon`. Selecting
- `optimizer: muon` requires a build that exposes that API (included in standard releases from
- PyTorch 2.9); runtime capability detection remains authoritative.
+- The `torch>=2.10.0` baseline includes `torch.optim.Muon`. Runtime capability detection remains
+ authoritative for vendor or custom builds that may omit the API.
- Muon is supported with DDP and FSDP2. The pre-load optimizer/backend validator rejects
DeepSpeed and FSDP1 before pretrained weights are loaded.
+### TorchAudio and TorchCodec
+- TorchAudio 2.10 delegates `torchaudio.load` and `torchaudio.save` to TorchCodec. The matching
+ baseline is TorchCodec 0.10, and its native extension also needs FFmpeg 4–8 shared libraries.
+- The CUDA Docker image installs FFmpeg. Non-container installations must provide compatible
+ system FFmpeg libraries; the `imageio[ffmpeg]` executable bundle does not provide those shared
+ libraries.
+
### diffusers
- Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule
may be used for upstream development, but must not silently override the declared runtime dependency.
@@ -84,6 +91,8 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
### accelerate
- Primary distributed backend. `accelerator.prepare()` wraps a single `ModelBundle` (all target components) plus the optimizer as one root (constraint #9).
+- Version 1.14 or newer is required because distributed runtime-identity checks use
+ `Accelerator.reduce(..., reduction="max")`; earlier releases do not implement max reduction.
- Online generation uses framework samplers. Finite SFT/offline-DPO data uses PyTorch's official
`DistributedSampler`; that already-sharded loader is not prepared via Accelerate.
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index 66f7899d1..c4dbfe485 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -711,7 +711,7 @@ Based on the fix type, write the fix entry to the appropriate document:
### Optional optimizer APIs must be validated before model loading
- **Date**: 2026-08-31
-- **Symptom**: A Muon configuration on the supported PyTorch 2.6 baseline could pass backend
+- **Symptom**: A Muon configuration on the then-declared PyTorch 2.6 baseline could pass backend
validation, load pretrained weights, and then fail when optimizer construction accessed the
unavailable `torch.optim.Muon` attribute.
- **Root Cause**: Backend validation assumed that parsing a Muon optimizer configuration implied
diff --git a/AGENTS.md b/AGENTS.md
index 2ebfbf74d..025337c7f 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -7,7 +7,7 @@ Flow-Factory is a unified **online and offline fine-tuning framework** for diffu
- **Algorithms**: SFT, offline DPO, online DPO, GRPO, GRPO-Guard, DPPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1
- **Models**: FLUX.1 (+Kontext), FLUX.2 (+Klein), SD3.5, Qwen-Image (+Edit-Plus), Z-Image, Wan2 (T2V/I2V), LTX2 (T2AV/I2AV), MiniMax H3 (T2VA/FL2VA/Ref2VA), Bagel, SenseNova-U1 (1.0/1.5; T2I + ordered multi-reference I2I)
- **Rewards**: PickScore (+Rank), CLIP, CLAP, ImageBind, OCR, GenEval/GenEval2, HPSv2, VLM-Evaluate, rational-rewards, and custom rewards
-- **Python**: >=3.10 | **PyTorch**: >=2.6.0 | **License**: Apache-2.0
+- **Python**: >=3.11 | **PyTorch**: >=2.10.0 | **License**: Apache-2.0
**Language**: Match user's language.
diff --git a/README.md b/README.md
index 2ce49eac1..df2225bdd 100644
--- a/README.md
+++ b/README.md
@@ -153,6 +153,9 @@ cd Flow-Factory
pip install -e .
```
+Flow-Factory requires Python 3.11 or newer and PyTorch 2.10 or newer. PyTorch 2.10 includes the
+native `torch.optim.Muon` API used by Muon optimizer configs.
+
Optional dependencies, such as `deepspeed`, are also available. Install them with:
```bash
@@ -163,8 +166,10 @@ pip install -e .[deepspeed]
> **Dependency:** MiniMax H3 and LTX2 require the released `diffusers>=0.40.0` API.
> PyAV >=18.0.0 decodes ordered video/audio references and target media.
+> TorchAudio 2.10 delegates audio loading and saving to TorchCodec, which also requires FFmpeg
+> shared libraries. The CUDA image installs those system libraries automatically.
-A CUDA training image (Python 3.12, **uv**-based install, PyTorch 2.8 + `cu129`, `deepspeed`, `wandb`, released `diffusers`) is defined under [`docker/docker-cuda/`](docker/docker-cuda/Dockerfile). See [`docker/README.md`](docker/README.md) for build and run instructions (including `linux/amd64` on Apple Silicon).
+A CUDA training image (Python 3.12, **uv**-based install, PyTorch 2.10 + `cu129`, `deepspeed`, `wandb`, released `diffusers`) is defined under [`docker/docker-cuda/`](docker/docker-cuda/Dockerfile). See [`docker/README.md`](docker/README.md) for build and run instructions (including `linux/amd64` on Apple Silicon).
## Experiment Trackers
@@ -341,7 +346,7 @@ The following reward models are pre-registered and ready to use:
| `rational_rewards_edit` | Pointwise | A reasoning reward model that provides multi-aspect reward for image edit; four aspects → scalar in [0, 1] | [RationalRewards-8B-Edit](https://huggingface.co/TIGER-Lab/RationalRewards-8B-Edit) |
| `qwen_image_bench` | Pointwise | Qwen-Image-Bench "Q-Judger"; hierarchical 5-dim / 56-facet scoring with per-prompt `dims_en` → scalar in [0, 1] | [Qwen-Image-Bench](https://github.com/QwenLM/Qwen-Image-Bench) |
-> **GenEval** requires extra dependencies (mmcv, mmdet, open_clip). Install with: `bash scripts/install_geneval_deps.sh` (Python 3.10 recommended). See [guidance/rewards.md](guidance/rewards.md#dataset-metadata-convention) for dataset format.
+> **GenEval** requires extra dependencies (mmcv, mmdet, open_clip). Install with: `bash scripts/install_geneval_deps.sh` (Python 3.11 or 3.12). See [guidance/rewards.md](guidance/rewards.md#dataset-metadata-convention) for dataset format.
> **VLM-as-Judge** (remote vLLM / OpenAI-style HTTP) is covered in [guidance/rewards.md#vlm-as-judge](guidance/rewards.md#vlm-as-judge) (`vllm_evaluate`, Rational Rewards, `qwen_image_bench`, async tips). For [RationalRewards](https://github.com/TIGER-AI-Lab/RationalRewards) specifically, serve the judge with [`scripts/start_vllm_rational_reward.sh`](scripts/start_vllm_rational_reward.sh) and set YAML `api_base_url` / `vlm_model` to match `--served-model-name` (defaults: `RationalRewards-8B-T2I` / `RationalRewards-8B-Edit`). For [Qwen-Image-Bench](https://github.com/QwenLM/Qwen-Image-Bench), use [`scripts/start_vllm_qwen_image_bench.sh`](scripts/start_vllm_qwen_image_bench.sh) and build the dataset with `python dataset/qwen_image_bench/prepare.py`.
diff --git a/config/accelerate_configs/fsdp2.yaml b/config/accelerate_configs/fsdp2.yaml
index 7a8701ad7..296aebca5 100644
--- a/config/accelerate_configs/fsdp2.yaml
+++ b/config/accelerate_configs/fsdp2.yaml
@@ -4,7 +4,7 @@ downcast_bf16: 'no'
fsdp_config:
fsdp_version: 2
fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
- # Accelerate 1.13 FSDP2 always uses original DTensor parameters and normalizes
+ # Accelerate 1.14+ FSDP2 always uses original DTensor parameters and normalizes
# use_orig_params to None. Do not add the obsolete fsdp_use_orig_params key.
# FSDP2 also rejects any explicit forward_prefetch value, including false.
fsdp_offload_params: false
@@ -23,4 +23,4 @@ same_network: true
tpu_env: []
tpu_use_cluster: false
tpu_use_sudo: false
-use_cpu: false
\ No newline at end of file
+use_cpu: false
diff --git a/docker/README.md b/docker/README.md
index ced5c85de..4df1fbe34 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -1,6 +1,6 @@
# Docker (CUDA) — Flow-Factory Training Image
-Pre-built GPU training image for Flow-Factory: CUDA 12.9, Python 3.12, PyTorch 2.8, DeepSpeed, and W&B — ready to run `ff-train` out of the box.
+Pre-built GPU training image for Flow-Factory: CUDA 12.9, Python 3.12, PyTorch 2.10, DeepSpeed, and W&B — ready to run `ff-train` out of the box.
## Prerequisites
diff --git a/docker/docker-cuda/Dockerfile b/docker/docker-cuda/Dockerfile
index 91d63bdbf..8c61e9b98 100644
--- a/docker/docker-cuda/Dockerfile
+++ b/docker/docker-cuda/Dockerfile
@@ -46,9 +46,10 @@ RUN mkdir -p /app/src/flow_factory && touch /app/src/flow_factory/__init__.py
ARG PYTORCH_INDEX_URL=https://download.pytorch.org/whl/cu129
RUN uv pip install \
- torch==2.8.0 \
- torchvision==0.23.0 \
- torchaudio==2.8.0 \
+ torch==2.10.0 \
+ torchvision==0.25.0 \
+ torchaudio==2.10.0 \
+ torchcodec==0.10.0 \
--index-url "${PYTORCH_INDEX_URL}"
RUN uv pip install ".[deepspeed,wandb]"
@@ -58,7 +59,7 @@ COPY . /app
RUN uv pip install --no-deps -e .
# Fail fast if any layer is broken
-RUN python -c "import deepspeed, diffusers, torch, wandb; assert tuple(map(int, diffusers.__version__.split('.')[:2])) >= (0, 40); print('torch', torch.__version__, 'diffusers', diffusers.__version__)" && ff-train --help >/dev/null
+RUN python -c "import deepspeed, diffusers, torch, torchcodec, wandb; assert hasattr(torch.optim, 'Muon'); assert tuple(map(int, diffusers.__version__.split('.')[:2])) >= (0, 40); print('torch', torch.__version__, 'torchcodec', torchcodec.__version__, 'diffusers', diffusers.__version__)" && ff-train --help >/dev/null
# Intentionally root: DeepSpeed / NCCL may need elevated IPC access.
CMD ["/bin/bash"]
diff --git a/pyproject.toml b/pyproject.toml
index 7957a5024..7d4917844 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ name = "flow-factory"
version = "0.1.0"
description = "Unified RL Fine-tuning Framework for Diffusion/Flow-Matching Models"
readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.11"
license = {text = "Apache-2.0"}
authors = [
{name = "Flow-Factory Team"}
@@ -17,7 +17,6 @@ classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Science/Research",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
- "Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
]
@@ -25,14 +24,15 @@ classifiers = [
# Core
dependencies = [
# Deep Learning Core
- "torch>=2.6.0",
- "torchvision>=0.19.0",
- "torchaudio>=2.4.0",
+ "torch>=2.10.0",
+ "torchvision>=0.25.0",
+ "torchaudio>=2.10.0",
+ "torchcodec>=0.10.0",
# Hugging Face Ecosystem
"transformers>=4.57.1", # Compatible with both v4.x and v5.x
"diffusers>=0.40.0",
- "accelerate>=1.11.0",
+ "accelerate>=1.14.0",
"peft>=0.17.0",
"datasets>=3.3.2",
"tokenizers>=0.22.1",
@@ -58,7 +58,7 @@ dependencies = [
[project.optional-dependencies]
# DeepSpeed (ZeRO optimization)
deepspeed = [
- "deepspeed>=0.15.4",
+ "deepspeed>=0.18.3",
]
# Quantization (8-bit/4-bit)
@@ -80,7 +80,7 @@ wandb = [ "wandb" ]
swanlab = [ "swanlab" ]
nvidia = [
- "xformers>=0.0.27",
+ "xformers>=0.0.35",
"nvidia-ml-py",
]
@@ -115,7 +115,7 @@ where = ["src"]
[tool.black]
line-length = 100
-target-version = ['py310', 'py311', 'py312']
+target-version = ['py311', 'py312']
[tool.isort]
profile = "black"
diff --git a/scripts/install_geneval_deps.sh b/scripts/install_geneval_deps.sh
index 8644172da..b9f2f4c71 100644
--- a/scripts/install_geneval_deps.sh
+++ b/scripts/install_geneval_deps.sh
@@ -4,8 +4,8 @@
# Install GenEval reward model dependencies (mmcv + mmdet + open_clip)
#
# Requirements:
-# - Python 3.10 or 3.12 (tested)
-# - PyTorch >= 2.0 with CUDA
+# - Python 3.11 or 3.12
+# - PyTorch >= 2.10 with CUDA
# - CUDA toolkit (nvcc) for mmcv CUDA ops compilation
# - uv (recommended) or pip
#
@@ -39,8 +39,8 @@ fi
PY_VERSION=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
-if [[ "$PY_VERSION" != "3.10" && "$PY_VERSION" != "3.12" ]]; then
- warn "Python ${PY_VERSION} detected. This script has only been tested with Python 3.10 and 3.12."
+if [[ "$PY_VERSION" != "3.11" && "$PY_VERSION" != "3.12" ]]; then
+ warn "Python ${PY_VERSION} detected. Flow-Factory supports Python 3.11 and 3.12."
warn "Proceeding anyway..."
echo ""
fi
diff --git a/src/flow_factory/rewards/ocr.py b/src/flow_factory/rewards/ocr.py
index 33feea173..853f65e33 100644
--- a/src/flow_factory/rewards/ocr.py
+++ b/src/flow_factory/rewards/ocr.py
@@ -20,8 +20,10 @@
pip install paddlepaddle-gpu==3.3.0 -i https://www.paddlepaddle.org.cn/packages/stable/cu129/
pip install paddleocr
pip install python-Levenshtein
-# Install torch2.8.0 and it will update nvcc toolkits automatically
-pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu129
+# Install the project baseline PyTorch stack and its audio codec runtime
+pip install \
+ torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 torchcodec==0.10.0 \
+ --index-url https://download.pytorch.org/whl/cu129
# Maybe you will need this:
yum install -y mesa-libGL glib2
```
@@ -83,9 +85,7 @@ def _compute_scores_batch(
metadata: Optional[list[str]] = None,
) -> list[float]:
"""Compute mean target-text fidelity for each image."""
- if len(prompt) != len(image) or (
- metadata is not None and len(prompt) != len(metadata)
- ):
+ if len(prompt) != len(image) or (metadata is not None and len(prompt) != len(metadata)):
raise ValueError(
"expected equal OCR batch lengths for prompt and image, with optional "
"metadata matching that length; "
@@ -132,11 +132,7 @@ def _targets_from_prompt(prompt: str, sample_index: int) -> list[str]:
f"expected string prompt for OCR sample {sample_index}, "
f"received {type(prompt).__name__}: {prompt!r}"
)
- targets = [
- target.strip()
- for target in _QUOTED_TEXT.findall(prompt)
- if target.strip()
- ]
+ targets = [target.strip() for target in _QUOTED_TEXT.findall(prompt) if target.strip()]
if not targets:
raise ValueError(
f"expected quoted OCR target in prompt for sample {sample_index}, "
diff --git a/src/flow_factory/utils/audio.py b/src/flow_factory/utils/audio.py
index 924b9d4ce..fc7dcdfa7 100644
--- a/src/flow_factory/utils/audio.py
+++ b/src/flow_factory/utils/audio.py
@@ -243,7 +243,7 @@ def load_audio(
Note:
Backend resolution (see :func:`_load_audio_backend`):
1. ``torchaudio`` — primary backend, handles wav/mp3/flac/ogg/...
- (``torchaudio>=2.4.0`` is a core dependency).
+ (``torchaudio>=2.10.0`` and its TorchCodec runtime are core dependencies).
2. ``soundfile`` — used when ``torchaudio`` is unavailable;
handles wav/flac/ogg.
3. stdlib ``wave`` — last-resort fallback, WAV-only,
From 93e4a935d049ec7701350be14286bdb998ac8a9b Mon Sep 17 00:00:00 2001
From: Jayce-Ping <315229706@qq.com>
Date: Mon, 31 Aug 2026 21:50:18 +0800
Subject: [PATCH 76/76] [deps] fix: restore Python 3.10 compatibility
---
.agents/knowledge/constraints.md | 4 ++--
.agents/knowledge/dependencies.md | 7 +++++++
.agents/knowledge/topics/fix_patterns.md | 13 +++++++++++++
.agents/knowledge/topics/minimax_h3.md | 2 +-
AGENTS.md | 2 +-
README.md | 6 +++---
dataset/minimax_h3_ref2va/README.md | 2 +-
guidance/datasets.md | 2 +-
pyproject.toml | 7 ++++---
scripts/install_geneval_deps.sh | 6 +++---
src/flow_factory/data_utils/dataset.py | 4 ++--
src/flow_factory/data_utils/offline_dataset.py | 4 ++--
tests/docs/test_minimax_h3_docs.py | 2 +-
13 files changed, 41 insertions(+), 20 deletions(-)
diff --git a/.agents/knowledge/constraints.md b/.agents/knowledge/constraints.md
index 964c943dc..77fcc746b 100644
--- a/.agents/knowledge/constraints.md
+++ b/.agents/knowledge/constraints.md
@@ -203,7 +203,7 @@ The adapter resolves `component_load_dtypes` at native load/materialization time
## Code Quality (21–27)
### 21. Formatting Standards
-- **Black** with `line-length=100`, targeting Python 3.11–3.12
+- **Black** with `line-length=100`, targeting Python 3.10–3.12
- **isort** with `profile="black"`, `line_length=100`
- Comments and docstrings in **English**
@@ -214,7 +214,7 @@ The adapter resolves `component_load_dtypes` at native load/materialization time
- **Top-level imports only**: All `import` / `from ... import ...` statements MUST live at the top of the module, never inside function bodies, methods, `__init__`, or conditional branches. Sanctioned exceptions: (a) optional dependencies wrapped in `try/except ImportError` (e.g., `deepspeed`, `xformers`); (b) backend-gated imports where the target symbol is only resolvable under a specific runtime backend already selected by a preceding feature check (e.g., DeepSpeed/FSDP submodules guarded by `is_deepspeed()` / `is_fsdp2()` in `models/abc.py`); (c) genuine unresolvable circular imports documented inline. Lazy imports added merely for "import speed" or "to keep the module light" are NOT acceptable — every hard dependency already runs through Python's import machinery on a typical import path. Inline imports hide the dependency surface from readers, `isort`, and static-analysis tools, and re-execute on every call in hot loops.
### 23. Type Annotations
-All public methods must have type annotations. Use `typing` module types (`List`, `Dict`, `Optional`, `Tuple`, `Union`) for Python 3.11 compatibility.
+All public methods must have type annotations. Use `typing` module types (`List`, `Dict`, `Optional`, `Tuple`, `Union`) for Python 3.10 compatibility.
### 24. License Header
All source files must include the Apache 2.0 license header with `Copyright 2026 Jayce-Ping`.
diff --git a/.agents/knowledge/dependencies.md b/.agents/knowledge/dependencies.md
index a9e25ed11..c84df2886 100644
--- a/.agents/knowledge/dependencies.md
+++ b/.agents/knowledge/dependencies.md
@@ -45,6 +45,7 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
| `transformers` | >= 4.57.1 | Text encoders, tokenizers |
| `diffusers` | >= 0.40.0 | Diffusion pipelines, schedulers, MiniMax H3 and LTX2 APIs |
| `accelerate` | >= 1.14.0 | Distributed training, mixed precision, and max reduction |
+| `av` | >= 17.0.0 | CPU video/audio decoding with Python 3.10 support |
| `peft` | >= 0.17.0 | LoRA, parameter-efficient fine-tuning |
| `datasets` | >= 3.3.2 | Dataset loading |
| `huggingface-hub` | >= 0.35.3 | Model/dataset downloads |
@@ -70,6 +71,12 @@ The authoritative list is `pyproject.toml` `[project.dependencies]` (20+ package
system FFmpeg libraries; the `imageio[ffmpeg]` executable bundle does not provide those shared
libraries.
+### PyAV
+- The tested floor is `av>=17.0.0`. Python 3.10 resolves to a compatible 17.x release because
+ PyAV 18 requires Python 3.11; Python 3.11 and newer may resolve to later PyAV releases.
+- Flow-Factory uses the container, stream, frame, and resampler APIs available in PyAV 17; it does
+ not require a PyAV 18-only API.
+
### diffusers
- Use the released `diffusers>=0.40.0` package as the authoritative API. The repository submodule
may be used for upstream development, but must not silently override the declared runtime dependency.
diff --git a/.agents/knowledge/topics/fix_patterns.md b/.agents/knowledge/topics/fix_patterns.md
index c4dbfe485..bcfab0d20 100644
--- a/.agents/knowledge/topics/fix_patterns.md
+++ b/.agents/knowledge/topics/fix_patterns.md
@@ -784,6 +784,19 @@ Based on the fix type, write the fix entry to the appropriate document:
the same fact even when Git reports those files as clean.
- **Related Constraint**: N/A
+### Dependency floors must preserve declared Python compatibility
+- **Date**: 2026-08-31
+- **Symptom**: Aligning the runtime metadata with `av>=18.0.0` raised Flow-Factory's Python floor
+ from 3.10 to 3.11 even though the framework and its Muon dependency stack still supported 3.10.
+- **Root Cause**: The media-decoder floor followed the latest tested PyAV release without checking
+ whether Flow-Factory used a PyAV 18-only API or whether PyAV 17 covered the same contract.
+- **Fix**: Restore `requires-python>=3.10`, set the tested decoder floor to `av>=17.0.0`, and align
+ classifiers, formatter targets, runtime errors, installation guidance, and agent documentation.
+- **Lesson**: A dependency-induced interpreter-floor increase is not automatically a framework
+ requirement. Verify the used API surface and test the last compatible dependency line before
+ dropping a supported Python version.
+- **Related Constraint**: N/A
+
## Cross-refs
- UP: [Hard Constraints](../constraints.md), [Architecture](../architecture.md)
diff --git a/.agents/knowledge/topics/minimax_h3.md b/.agents/knowledge/topics/minimax_h3.md
index bde33e60c..7da450e11 100644
--- a/.agents/knowledge/topics/minimax_h3.md
+++ b/.agents/knowledge/topics/minimax_h3.md
@@ -53,7 +53,7 @@ and lets that boundary stay strict.
geometry cache fields and a preprocessing cache version.
- Reference paths are dataset-relative. Positive finite `fps` and `sample_rate` overrides follow
`samples/references.py`.
-- PyAV >=18.0.0 decodes video/audio references, including embedded or separate soundtracks.
+- PyAV >=17.0.0 decodes video/audio references, including embedded or separate soundtracks.
## Offline audiovisual output contract
diff --git a/AGENTS.md b/AGENTS.md
index 025337c7f..14aecb73b 100644
--- a/AGENTS.md
+++ b/AGENTS.md
@@ -7,7 +7,7 @@ Flow-Factory is a unified **online and offline fine-tuning framework** for diffu
- **Algorithms**: SFT, offline DPO, online DPO, GRPO, GRPO-Guard, DPPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD, DMD2, TDM, TDM-R1
- **Models**: FLUX.1 (+Kontext), FLUX.2 (+Klein), SD3.5, Qwen-Image (+Edit-Plus), Z-Image, Wan2 (T2V/I2V), LTX2 (T2AV/I2AV), MiniMax H3 (T2VA/FL2VA/Ref2VA), Bagel, SenseNova-U1 (1.0/1.5; T2I + ordered multi-reference I2I)
- **Rewards**: PickScore (+Rank), CLIP, CLAP, ImageBind, OCR, GenEval/GenEval2, HPSv2, VLM-Evaluate, rational-rewards, and custom rewards
-- **Python**: >=3.11 | **PyTorch**: >=2.10.0 | **License**: Apache-2.0
+- **Python**: >=3.10 | **PyTorch**: >=2.10.0 | **License**: Apache-2.0
**Language**: Match user's language.
diff --git a/README.md b/README.md
index df2225bdd..0c29b8227 100644
--- a/README.md
+++ b/README.md
@@ -153,7 +153,7 @@ cd Flow-Factory
pip install -e .
```
-Flow-Factory requires Python 3.11 or newer and PyTorch 2.10 or newer. PyTorch 2.10 includes the
+Flow-Factory requires Python 3.10 or newer and PyTorch 2.10 or newer. PyTorch 2.10 includes the
native `torch.optim.Muon` API used by Muon optimizer configs.
Optional dependencies, such as `deepspeed`, are also available. Install them with:
@@ -165,7 +165,7 @@ pip install -e .[deepspeed]
> **Note**: The Bagel adapter requires `flash-attn` (>= 2.5.8) and `opencv-python`. Install them with `pip install -e .[bagel]` (the `[bagel]` extra is intentionally not part of `[all]` because flash-attn is heavy to build).
> **Dependency:** MiniMax H3 and LTX2 require the released `diffusers>=0.40.0` API.
-> PyAV >=18.0.0 decodes ordered video/audio references and target media.
+> PyAV >=17.0.0 decodes ordered video/audio references and target media.
> TorchAudio 2.10 delegates audio loading and saving to TorchCodec, which also requires FFmpeg
> shared libraries. The CUDA image installs those system libraries automatically.
@@ -346,7 +346,7 @@ The following reward models are pre-registered and ready to use:
| `rational_rewards_edit` | Pointwise | A reasoning reward model that provides multi-aspect reward for image edit; four aspects → scalar in [0, 1] | [RationalRewards-8B-Edit](https://huggingface.co/TIGER-Lab/RationalRewards-8B-Edit) |
| `qwen_image_bench` | Pointwise | Qwen-Image-Bench "Q-Judger"; hierarchical 5-dim / 56-facet scoring with per-prompt `dims_en` → scalar in [0, 1] | [Qwen-Image-Bench](https://github.com/QwenLM/Qwen-Image-Bench) |
-> **GenEval** requires extra dependencies (mmcv, mmdet, open_clip). Install with: `bash scripts/install_geneval_deps.sh` (Python 3.11 or 3.12). See [guidance/rewards.md](guidance/rewards.md#dataset-metadata-convention) for dataset format.
+> **GenEval** requires extra dependencies (mmcv, mmdet, open_clip). Install with: `bash scripts/install_geneval_deps.sh` (Python 3.10 or newer). See [guidance/rewards.md](guidance/rewards.md#dataset-metadata-convention) for dataset format.
> **VLM-as-Judge** (remote vLLM / OpenAI-style HTTP) is covered in [guidance/rewards.md#vlm-as-judge](guidance/rewards.md#vlm-as-judge) (`vllm_evaluate`, Rational Rewards, `qwen_image_bench`, async tips). For [RationalRewards](https://github.com/TIGER-AI-Lab/RationalRewards) specifically, serve the judge with [`scripts/start_vllm_rational_reward.sh`](scripts/start_vllm_rational_reward.sh) and set YAML `api_base_url` / `vlm_model` to match `--served-model-name` (defaults: `RationalRewards-8B-T2I` / `RationalRewards-8B-Edit`). For [Qwen-Image-Bench](https://github.com/QwenLM/Qwen-Image-Bench), use [`scripts/start_vllm_qwen_image_bench.sh`](scripts/start_vllm_qwen_image_bench.sh) and build the dataset with `python dataset/qwen_image_bench/prepare.py`.
diff --git a/dataset/minimax_h3_ref2va/README.md b/dataset/minimax_h3_ref2va/README.md
index c85df671b..d4dfa5684 100644
--- a/dataset/minimax_h3_ref2va/README.md
+++ b/dataset/minimax_h3_ref2va/README.md
@@ -18,7 +18,7 @@ At least one image or video is required; audio-only manifests are invalid. A vid
is valid only when that video also supplies `audio_path`. Manifest `fps` and `sample_rate`
overrides take precedence over decoded metadata where supported.
-PyAV >=18.0.0 is required for reliable video/audio decoding. It preserves video frames and FPS,
+PyAV >=17.0.0 is required for reliable video/audio decoding. It preserves video frames and FPS,
plus embedded or separately referenced audio and its sample rate. Only encoded tensors/layout and
the canonical manifest enter the Arrow cache; upstream reference objects are transient.
diff --git a/guidance/datasets.md b/guidance/datasets.md
index 782fea196..aa3efffed 100644
--- a/guidance/datasets.md
+++ b/guidance/datasets.md
@@ -114,7 +114,7 @@ must carry the supervision type required by its trainer. Prompt-only rows, mixed
preference rows, unknown keys, and non-V2 records fail during manifest loading.
All V2 media paths are resolved against that source's `dataset_dir`; an absolute path is retained.
-Images, videos, and audio have built-in CPU decoders. Video targets require PyAV 18 or newer.
+Images, videos, and audio have built-in CPU decoders. Video targets require PyAV 17 or newer.
Decoded audio is a detached CPU `float32` waveform shaped `(channels, samples)`. A manifest
`sample_rate` is a logical source-clock override and does not pre-resample the decoded samples;
source-clock truncation, channel conversion, the single model-rate conversion, posterior selection,
diff --git a/pyproject.toml b/pyproject.toml
index 7d4917844..a86efc613 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ name = "flow-factory"
version = "0.1.0"
description = "Unified RL Fine-tuning Framework for Diffusion/Flow-Matching Models"
readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.10"
license = {text = "Apache-2.0"}
authors = [
{name = "Flow-Factory Team"}
@@ -17,6 +17,7 @@ classifiers = [
"Development Status :: 3 - Alpha",
"Intended Audience :: Science/Research",
"Topic :: Scientific/Engineering :: Artificial Intelligence",
+ "Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
]
@@ -44,7 +45,7 @@ dependencies = [
"ftfy", # Required by Diffusers Wan prompt normalization
# Basic Utils
- "av>=18.0.0",
+ "av>=17.0.0",
"imageio[ffmpeg]>=2.37.2",
"numpy>=1.26.4",
"pillow>=10.4.0",
@@ -115,7 +116,7 @@ where = ["src"]
[tool.black]
line-length = 100
-target-version = ['py311', 'py312']
+target-version = ['py310', 'py311', 'py312']
[tool.isort]
profile = "black"
diff --git a/scripts/install_geneval_deps.sh b/scripts/install_geneval_deps.sh
index b9f2f4c71..e22d4017c 100644
--- a/scripts/install_geneval_deps.sh
+++ b/scripts/install_geneval_deps.sh
@@ -4,7 +4,7 @@
# Install GenEval reward model dependencies (mmcv + mmdet + open_clip)
#
# Requirements:
-# - Python 3.11 or 3.12
+# - Python >= 3.10
# - PyTorch >= 2.10 with CUDA
# - CUDA toolkit (nvcc) for mmcv CUDA ops compilation
# - uv (recommended) or pip
@@ -39,8 +39,8 @@ fi
PY_VERSION=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')")
-if [[ "$PY_VERSION" != "3.11" && "$PY_VERSION" != "3.12" ]]; then
- warn "Python ${PY_VERSION} detected. Flow-Factory supports Python 3.11 and 3.12."
+if ! python -c "import sys; raise SystemExit(0 if sys.version_info >= (3, 10) else 1)"; then
+ warn "Python ${PY_VERSION} detected. Flow-Factory requires Python 3.10 or newer."
warn "Proceeding anyway..."
echo ""
fi
diff --git a/src/flow_factory/data_utils/dataset.py b/src/flow_factory/data_utils/dataset.py
index 69330ba96..deb46378c 100644
--- a/src/flow_factory/data_utils/dataset.py
+++ b/src/flow_factory/data_utils/dataset.py
@@ -1422,8 +1422,8 @@ def _require_finite_positive_rate(
def _require_pyav() -> Any:
if av is None:
raise ImportError(
- "ordered video/audio references require PyAV>=18.0.0; "
- "install with `pip install 'av>=18.0.0'`"
+ "ordered video/audio references require PyAV>=17.0.0; "
+ "install with `pip install 'av>=17.0.0'`"
)
return av
diff --git a/src/flow_factory/data_utils/offline_dataset.py b/src/flow_factory/data_utils/offline_dataset.py
index 767257ffd..c6f80dcd5 100644
--- a/src/flow_factory/data_utils/offline_dataset.py
+++ b/src/flow_factory/data_utils/offline_dataset.py
@@ -485,8 +485,8 @@ def decode_video(asset: MediaAsset) -> np.ndarray:
"""
if av is None:
raise ImportError(
- "offline target video decoding requires PyAV>=18.0.0; "
- "install with `pip install 'av>=18.0.0'`"
+ "offline target video decoding requires PyAV>=17.0.0; "
+ "install with `pip install 'av>=17.0.0'`"
)
try:
with av.open(asset.path) as container:
diff --git a/tests/docs/test_minimax_h3_docs.py b/tests/docs/test_minimax_h3_docs.py
index d166adbcd..8d0bdf0a9 100644
--- a/tests/docs/test_minimax_h3_docs.py
+++ b/tests/docs/test_minimax_h3_docs.py
@@ -50,7 +50,7 @@ def test_readme_documents_h3_links_dependency_and_limits() -> None:
for required in (
"diffusers>=0.40.0",
"pip install -e .",
- "PyAV >=18.0.0",
+ "PyAV >=17.0.0",
"B=1",
"no CFG",
"shift 12",