Skip to content

feat(flux2): add Flux2KleinImg2ImgPipeline#14044

Draft
Sainava wants to merge 1 commit into
huggingface:mainfrom
Sainava:feature/flux2-klein-img2img
Draft

feat(flux2): add Flux2KleinImg2ImgPipeline#14044
Sainava wants to merge 1 commit into
huggingface:mainfrom
Sainava:feature/flux2-klein-img2img

Conversation

@Sainava

@Sainava Sainava commented Jun 22, 2026

Copy link
Copy Markdown

What does this PR do?

Fixes #14042

Description and Motivation

The diffusers library currently provides a dedicated Flux2KleinPipeline for text-to-image generation, but lacks a native Image-to-Image (Img2Img) pipeline for the FLUX.2 Klein architecture.

This PR introduces a dedicated Flux2KleinImg2ImgPipeline by adapting the existing Diffusers img2img workflow to the FLUX.2 Klein architecture, carefully preserving its specific latent processing and normalization behavior.

Implementation Details

  • Standard Img2Img Workflow: Reuses the standard Diffusers workflow, including proper image preprocessing, VAE encoding, timestep slicing based on strength, latent noise initialization, and denoising via FlowMatchEulerDiscreteScheduler (using scale_noise).
  • Klein Architecture Alignment: Integrates seamlessly with Klein's specific _encode_vae_image path, preserving crucial patchification and batch-normalization behavior.
  • Safety & Hardware Optimizations: Maintains existing Klein image area safeguards (capping at 1024x1024) to avoid excessive memory usage.
  • Input Validation: Adds robust defensive checks, including strength boundary validation (0, 1] and batch-size divisibility checking before tensor duplication.

Testing

  • Verified import and pipeline registration through the Diffusers lazy-loading system.
    • Verified local end-to-end tensor flow, latent preparation, sequence length calculation, and dimension packing/unpacking using CPU mock implementations of the Klein Transformer and VAE.

(Note for Reviewers: I am opening this as a Draft PR to get feedback on the overall approach and architecture. Once the implementation direction is approved, I will add the standard fast/slow pipeline tests, import tests, and API documentation updates!)

Before submitting

Who can review?

@yiyixuxu @sayakpaul

@github-actions github-actions Bot added fixes-issue pipelines size/L PR with diff > 200 LOC labels Jun 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add Flux2KleinImg2ImgPipeline for FLUX.2 Klein

1 participant