Skip to content

Add img2img_flux2: image + prompt editing with FLUX.2 klein - #72

Merged
TroyHernandez merged 2 commits into
mainfrom
feat/flux2-img2img
Oct 28, 2026
Merged

TroyHernandez merged 2 commits into
mainfrom
feat/flux2-img2img

Conversation

@TroyHernandez

Copy link
Copy Markdown
Contributor

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() and img2img_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)

  • Encoder, 512x512: relative error ~2e-6 at every block.
  • Preprocessing: exact at 1024x1024; the 1920x1080 resize path is off by at most 1.1/255 (mean 0.25/255), from PIL's 8-bit rounding.
  • End-to-end 256x256 edit, seed 7, shared prompt embeds: PSNR 123 dB.
  • 27 new tinytests, including an encoder parity fixture (tiny random VAE, generated by tools/gen_fixtures_flux2_vae_encoder.py). Full suite: 1487 passed, 0 failed.

Not done

  • GPU memory/time at 1024x1024 with a reference: not measured (the card was held by the fleet worker). One reference doubles the image tokens, so expect a peak above txt2img's 12.5 GB.
  • Fleet support (resident_load, gpu.ctl entry, xtx.api) is a follow-up PR.

Version bump to 0.2.2.15 is in its own commit.

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).
@TroyHernandez
TroyHernandez merged commit 1f5bf61 into main Oct 28, 2026
2 checks passed
@TroyHernandez
TroyHernandez deleted the feat/flux2-img2img branch October 28, 2026 02:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant