feat([flux2-multi-reference-python-runner]): wire diffusers KV multi-reference editing for FLUX.2#524
Merged
Conversation
…reference editing in flux2 runner scripts/flux2_macos.py now branches into Flux2KleinKVPipeline when one or more --reference-images are supplied, builds the parallel PIL list at the output resolution, and forwards it as the pipeline's image= kwarg so the K/V-cached reference-token attention path actually runs. The three load helpers (sdnq, int8, bf16) take an explicit pipeline_cls so either pipeline class can drive each quantization branch; component init signatures are identical so the same flux2-klein repo loads either way. flux2-klein-9b's tokenizerRepo moves to black-forest-labs/FLUX.2-klein-9B-kv in the seed registry (data.reference/media-models.json) and the inline DEFAULT_REGISTRY (server/lib/mediaModels.js). The runner's existing probe_hf_auth then surfaces a clean USER_ERROR:gated_repo message pointing the user at the kv license URL on first use — accepting it enables both single-image and multi-reference renders on this model. Migration 045 swaps the persisted value on existing installs only when it still matches the pre-change shipped string, preserving any user-customized fork pin. The Python runner warns when --reference-strengths carries non-default values (the KV pipeline doesn't yet expose per-reference attention weighting; route still emits 1.0 for unset slots so the warning only fires on explicit non-1.0 weights) and when --image-path is supplied alongside refs (init image gets ignored — i2i strength and KV-cache reference conditioning are incompatible in the current pipeline).
- Refuse multi-reference editing on int8 quantization with a clear error instead of crashing mid-inference; Flux2KleinKVPipeline's reference-attention hooks don't compose with the quanto-rehydrated transformer attached in the int8 load path. - Rename sidecar field referenceStrengths -> referenceStrengthsRequested so it doesn't pretend per-ref weights were applied (the KV pipeline doesn't expose per-reference attention weighting today; the runner only warns on non-1.0 entries).
- Lazy-import the chosen pipeline class so a venv whose diffusers pin predates Flux2KleinKVPipeline still serves single-image renders; the KV branch fails loudly with an INSTALL_FLUX2 hint instead of breaking every FLUX.2 render with an opaque ImportError. The existing FLUX.2 venv health probe only verifies Flux2KleinPipeline; before this change, importing the KV class at module load would mark such an install healthy then crash on every single-image render too. - Rename the JS sidecar field referenceImageStrengths -> referenceImageStrengthsRequested to mirror the Python rename and match the contract that per-ref weights are NOT yet honored. - Update the Image Gen reference-images warning to reflect the new reality (multi-ref is wired; per-ref strength sliders are advisory until diffusers exposes a per-ref attention knob).
- Extend the int8+kv refuse gate to bf16 (quantization='none') too: the bf16 model's repo points at the base 9B (not the kv variant) whose transformer wasn't tuned for the reference-editing task, so the KV pipeline would run but produce off-task output. Lifting the gate per quantization once the corresponding KV path is validated is tracked as [flux2-bf16-kv-multi-reference-support]. - Pass formats=['PNG','JPEG','WEBP'] to PIL Image.open for the init and reference inputs so a magic-byte-spoofed file can't route into PIL's EPS/PDF/Ghostscript handler. Defense in depth — the route layer already enforces the mime allow-list, but PIL's Image.open dispatches on magic bytes, not extensions. - Add PLAN follow-ups for the bf16+kv gate lift and the deferred sidecar-meta test coverage gap.
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
1 participant
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
scripts/flux2_macos.pynow branches intoFlux2KleinKVPipelinewhen one or more--reference-imagesare provided, loads each ref as a PIL at the output resolution, and forwards them asimage=[…]so the K/V-cached reference-token attention path actually runs. Eachload_pipeline_*helper takes an explicitpipeline_clsso any quantization branch (sdnq / int8 / bf16) can drive either pipeline class — component init signatures are identical so the same flux2-klein repo loads either way.flux2-klein-9b'stokenizerRepomoves toblack-forest-labs/FLUX.2-klein-9B-kvin both the seed registry (data.reference/media-models.json) and the inlineDEFAULT_REGISTRY(server/lib/mediaModels.js). The runner's existingprobe_hf_authsurfaces a cleanUSER_ERROR:gated_repomessage pointing the user at the kv license URL on first use — accepting that license enables both single-image and multi-reference renders.045-flux2-klein-9b-kv-tokenizer.jspatches the persistedtokenizerRepoon existing installs only when it still matches the pre-change shipped value, preserving any user-customized fork pin.--reference-strengthscarries non-default values (the KV pipeline doesn't yet expose per-reference attention weighting) and when--image-pathis supplied alongside refs (the init image gets ignored — i2i strength and KV-cache reference conditioning are incompatible in the current pipeline).PLAN slug:
[flux2-multi-reference-python-runner].Test plan
vitest run lib/mediaModels routes/imageGen.multipart services/imageGen/local scripts/migrations/045— 61 pass.scripts/flux2_macos.pywith two reference images: argparse accepts the new args, the runner reportsFlux2KleinKVPipelineas the chosen class, warns on a non-default 0.5 strength, then hitsprobe_hf_authand surfaces aUSER_ERROR:gated_repo:black-forest-labs/FLUX.2-klein-9B-kvwith the license URL.