Repository navigation
Add img2img_flux2: image + prompt editing with FLUX.2 klein - #72
Merged
Merged
Conversation
Port the AutoencoderKLFlux2 encoder and the klein pipeline's reference-image path from diffusers: the reference is encoded once, its tokens (position ids T = 10) ride after the noisy latents in every forward, and only the latents' predictions drive the scheduler. flux2_load_pipeline(encoder = TRUE) loads the encoder; txt2img_flux2() shares the generation body unchanged. Matches diffusers on the real klein-4B weights at fp32 on the CPU (256x256 edit: PSNR 123 dB; encoder relative error 2e-6).
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.
Adds
img2img_flux2(prompt, image = ): image + prompt editing with FLUX.2 klein-4B, using the model's native reference-image path (same weights, no new download). One reference image for now.What's new
flux2_vae_encoder()/load_flux2_vae_encoder(): the AutoencoderKLFlux2 encode path (encoder body, quant_conv, BN stats), reusing the decoder's resnet and mid-block modules. Returns the posterior mean (sample_mode = "argmax").flux2_encode_image(): encode, patchify 2x2, BN-normalize, pack to tokens.flux2_prepare_image_ids(): reference position ids with T = 10 per image (list-shaped for a later multi-image change).flux2_preprocess_image(): diffusers' preprocessing (min side 64, aspect <= 8, Lanczos area cap at 1024^2, floor to /16, center crop, [-1, 1]). The Lanczos resize is a port of Pillow's coefficient algorithm as a separable matmul.flux2_load_pipeline(encoder = TRUE)loads the encoder; the default leaves txt2img memory unchanged.txt2img_flux2()andimg2img_flux2()share one internal generator. The reference tokens are appended to every transformer forward and the prediction is sliced back to the target tokens; mu uses the target length only.Validation (real weights, fp32 CPU, vs diffusers)
tools/gen_fixtures_flux2_vae_encoder.py). Full suite: 1487 passed, 0 failed.Not done
resident_load, gpu.ctl entry, xtx.api) is a follow-up PR.Version bump to 0.2.2.15 is in its own commit.