Conversation
Reference images travel as a slot-major input Part chain. The driver loads N per prompt, transposes prompt-major to slot-major, and chains one input Part per slot, so every Part stays row-aligned to the prompt batch -- Part.__post_init__ requires len(primitive) == len(sample_ids), and primitives are keyed by canonical modality, so N references cannot share one Part. FLUX consumes them as a longer token sequence. Each reference is already resized to the generation canvas, so the condition tensors stay dense and Flux2KleinConditions / predict_noise need no change -- they concatenate tokens and ids blindly. The only model-side change is the RoPE time offset, now scale * (slot + 1), replicating upstream's t_coords = [scale + scale * k for k in range(N)] so the trainside (recomputed) and SGLang (captured) ids stay identical. They agreed at N=1 only because both produced t=10. On the SGLang path one reference keeps the flat per-prompt condition_image shape; two or more nest one list per prompt, which is what upstream's preprocess_condition_image and ImageVAEEncodingStage already iterate. The per-prompt stash now keys on nesting rather than length, because a single prompt with N references is a length-1 outer list. Batches must be uniform in reference count: the lineage walk resolves every ancestor id by name, so a Part holding only some samples is invalid. Ragged batches are rejected with a message pointing at bucketing. Models that do not support multiple references (Qwen-Image-Edit-Plus, HunyuanImage 3.0, BAGEL) keep their own fail-fast guards, which now surface on N>1 instead of the loader truncating first.
Keep ordered reference slots for reward requests so EditReward scores against the primary source, and reject extra references in single-image model paths instead of silently ignoring them.
Close condition-image handles deterministically and describe both supported per-prompt wire shapes accurately.
Carry every image through the VLM encoding and emit the SGLang image_data shape required for both single-candidate and parallel-sampling requests.
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add the DyRef two-stage training path on top of #417's ordered
ImageSet/ImageSetsfoundation.This is a stacked draft. It depends on #417 and should be rebased onto
mainafter #417 lands.Related Issue
Refs #470
Depends on #417.
Test Plan
SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure— passed.python3 -m compileall -q unirl datasets/dyref— passed.img_shapes, SGLang capture splitting, replay reconstruction, and legacy N=1 wire compatibility — passed.{1,2,3}SDE scheduling — passed.Compatibility / Risk
role="target"image that is projected only into reward metadata; it never becomes a model input.LoraConfig.initial_adapter_pathis optional. It accepts standard PEFT adapter directories and intentionally rejects rank/alpha mismatch and meta-initialized loading.stretchas the default condition resize mode; only the DyRef recipes opt into aspect-preserving center crop.dar_min_weight=0.5collapses typical published similarity values to uniform weights. These recipes use0.01so DAR remains active; GPU validation must report the resulting weight distribution.Reviewer Notes
fe40529; the earlier commits belong to the prerequisite PR.unirl/data/data_source.py), then Qwen rollout/replay parity, reward math, and adapter warm-start.Checklist