SIGN IN SIGN UP

feat(flux2klein): Onboard multi-image editing support in pure Flax NNX (4B & 9B) with E2E parity and smoke tests

- Add multi-image editing pipeline (FlaxFlux2KleinPipeline) supporting arbitrary reference image conditioning
- Port VAE encoder and decoder to Flax NNX (AutoencoderKLFlux2NNX)
- Implement efficient image preprocessing, tiling, aspect-ratio-aware padding, and packing utilities
- Support fast sharded safetensors loading and TP sharding for Klein 4B & 9B
- Use standard HF_HOME resolution without machine-specific hardcoded paths
- Ensure clean compatibility across Transformers 4.x and 5.x via lazy module loading and dynamic FlaxPreTrainedModel lookup
- Add end-to-end multi-image editing parity test and preprocessing unit test suite
A
Amey Pasarkar committed
e86be47477241a826ad1f0754d61b43eaff82306
Parent: 7cfb880
Committed by amepas <ameypasarkar@gmail.com> on 8/28/2026, 1:09:46 AM