Image-Edit Fine-tuning Framework

Overview

This repository provides a comprehensive framework for fine-tuning image editing tasks. The framework supports FLUX Kontext, DreamOmni2, Qwen-Image-Edit, and Qwen-Image-Edit-2509 model architectures. Our implementation focuses on efficient training through LoRA (Low-Rank Adaptation) and features an optimized embedding cache system that achieves 2-3x training acceleration.

New

🚀 DreamOmni2 Trainer Support (v3.3.0): Full support for DreamOmni2 LoRA fine-tuning with cumulative offset positioning for multi-image conditioning. Features include VLM-based prompt optimization and proper spatial relationship encoding through cumulative offsets. See CHANGELOG:v3.3.0 for details.
- Cumulative offset support for multi-image position encoding (RoPE)
- Optional VLM prompt optimization using Qwen2.5-VL model
- Compatible with existing Flux Kontext training pipeline
- Example config: tests/test_configs/test_dreamomni2_fp16.yaml
🔧 Multi-Logger Support (v3.2.0): Added unified logging interface supporting TensorBoard, Weights & Biases, and SwanLab. Switch between logging backends with simple config changes. See Logging Guide for details.
- Unified LoggerManager API for all logging operations
- Support for TensorBoard (local), wandb (cloud), and SwanLab (cloud)
- Automatic authentication via .env files
- Distributed training support with automatic process coordination

TensorBoard

Wandb

SwanLab

✨ Validation Visualization Improvements (v3.1.0): Enhanced validation sampling visualization with step-by-step progress display. Added support for viewing model improvement during training at regular intervals (every 100 steps).

Step 100

Step 200

Step 300

Step 400

Step 500

Step 600

Validation sampling images showing results after every 100 training steps (steps 100-600) Tested with tests/test_configs/test_example_fluxkontext_fp16.yaml --cache

🔧 FSDP LoRA Checkpoint Fix (v3.0.2): Fixed FSDP training issues with LoRA checkpoint saving and documented performance benchmarks across different training strategies. BF16 FSDP shows optimal memory-performance balance (10GB memory, 1.7 FPS) compared to FP4 DDP (25GB memory, 0.4 FPS). See CHANGELOG:v3.0.2 for details.
🧪 Test Infrastructure Improvements (v3.0.1): Enhanced test resources management with HuggingFace Hub integration. Test resources are now automatically downloaded on-demand, reducing repository size and improving maintainability. Added comprehensive E2E tests for Qwen-Image-Edit and Qwen-Image-Edit-Plus sampling workflows. See CHANGELOG:v3.0.1 for details.
🎯 Multi-Resolution Mixed Training (v3.0.0): Revolutionary multi-resolution training support that allows training with multiple resolution candidates in a single session. The system intelligently selects the optimal resolution based on each image's aspect ratio, supporting both simple shared-resolution mode and advanced per-image-type configuration. Compatible with all model architectures and training modes. See CHANGELOG:v3.0.0 for complete details.
- Simple mode: multi_resolutions: [320*320, 512*512, 640*640]
- Advanced mode with per-type control: multi_resolutions: {target: [...], controls: [[...], [...]]}
- Example config: tests/test_configs/test_example_fluxkontext_multiresolution.yaml
📚 Documentation Improvements (v2.4.1): Comprehensive documentation updates including MIT license badge, enhanced data preparation guide (Folder/HuggingFace/CSV sources), and English language standardization. See CHANGELOG:v2.4.1 for details.
🔥 Dynamic Shape Support (v2.4.0): For Qwen-Image-Edit or Plus, we introduce the fixed number of pixels condition for batch process such that it support multiple shapes.
- data.init_args.processor.init_args.target_pixels: 512*512
- data.init_args.processor.init_args.controls_pixels: [512*512] But this still got limitations for the randomness of shapes used in training. Next we may add H/W buckets to support real dynamic shapes training.
Qwen-Image-Edit-Plus (2509) Support (v2.3.0): Complete support for the enhanced Qwen-Image-Edit-Plus model architecture with native multi-image composition capabilities. Read here for changes of the Qwen-Image-Edit-Plus version. Refer predict notebook for the predict example notebook. Pretrained model provided in TsienDragon/qwen-image-edit-plus-lora-face-seg

⚠️ Important: Qwen-Image-Edit-Plus (2509) requires the latest version of diffusers. Please install it using:

pip install --upgrade "git+https://github.com/huggingface/diffusers.git"

If you don't need Qwen-Image-Edit-Plus (2509) support, you can use the older version from requirements.txt.

Original Image

LoRA for Face Segmentation

Original with Character Composition (Plus version support multiple image composition natively)

LoRA with Character Composition

- **CSV dataset support**: 2025 Sep 24 - support csv dataset path - **Multi Control**: 2025 Sep 16

control 0

control 1

generated results

Support Multi Controls. The process logic is concat the latent of all control latents. And use different latent_id to identify them.

Pretrain Model is provided in Huggingface TsienDragon/character-compositing

TODO

Resolve validation DDP hang when validation samples use mixed resolutions (current workaround: keep validation shapes aligned). See docs/TODO.md.

Key Features

Multi-Model Support: Complete support for Qwen-Image-Edit, FLUX Kontext, DreamOmni2, and Qwen-Image-Edit-2509 model architectures
Multi-Resolution Mixed Training: Train with multiple resolution candidates (e.g., 320x320, 512x512, 640x640) in a single session with intelligent aspect-ratio-aware selection
Multi-Precision Training: FP16, FP8, and FP4 quantization levels for different hardware requirements
Efficient Fine-tuning: LoRA-based parameter-efficient fine-tuning with minimal memory footprint
Edit Mask Loss feature documentation in docs/spec/losses/image_edit_mask_loss.md Advanced mask-weighted loss function for focused training on edit regions
Speed Optimization including quantilizationand flash attention in docs/guide/speed_optimization.md
Embedding Cache System: Proprietary caching mechanism for 2-3x training acceleration
Validation Sampling: Real-time training progress monitoring with TensorBoard visualization
Resume Training: Seamless training resumption from checkpoints with full state recovery
HuggingFace Integration: Full compatibility with HuggingFace ecosystem for LoRA weights sharing and deployment
Auto-Upload to HuggingFace: One-click upload of trained LoRA weights to HuggingFace Hub
Multi-GPU Support: Distributed training capabilities with gradient accumulation
FSDP Training: Fully Sharded Data Parallel training with optimized memory usage (10GB vs 25GB)
Quantization Support: FP4/FP8/FP16 quantization for reduced memory usage and performance optimization
Flexible Architecture: Modular design supporting various vision-language tasks
Production Ready: Comprehensive testing suite and deployment configurations
Multi Control: Support Multiple Controls for Image-Edit model that can support images compositing tasks.

Documentation

Comprehensive documentation is organized in the docs/ directory:

Getting Started Guide — Setup, training, inference, and configuration guides
- Setup — Environment setup and installation
- Training — Training workflows and best practices
- Inference — Model inference and deployment
- Data Preparation — Dataset preparation and formats
- Configuration — Configuration system reference
- LoRA Fine-tuning — LoRA training strategies
- Validation Sampling — Training progress visualization
- Speed Optimization — Performance tuning and quantization
- HuggingFace Integration — Dataset and model sharing
- Testing Reference — Testing guidelines
- Debug Guide — Troubleshooting and debugging
Technical Specifications — Detailed module specifications
- Models — Architecture specs for FLUX Kontext, DreamOmni2, Qwen-Image-Edit, and Qwen-Image-Edit-Plus
- Losses — Loss function implementations and design rationale
- Data — Data loading and caching system specifications
- Trainer — Training loop and trainer architecture
Changelog — Version history and release notes
- Changelog Overview — Summary of all releases
- Latest: v3.3.0 — DreamOmni2 trainer support
Development Plans — Roadmap and feature planning
- Plan Index — Active and completed initiatives
TODO List — Current development tasks and progress

Dataset

Here we provided two toy datasets in the huggingface that user can efficiently use to train the model.

Face segmentation dataset: TsienDragon/face_segmentation_20
Character composition dataset: TsienDragon/character-composition

Quick usage:

from qflux.utils.huggingface import load_editing_dataset

dd = load_editing_dataset("TsienDragon/face_segmentation_20")
sample = dd["train"][0]

Dataset structure reference and upload/download instructions are in docs/guide/huggingface-related.md.

Added CSV data format support (v2.2.0): Dataset management based on CSV metadata files is now supported, providing a more flexible dataset structure. For datasets that use a CSV metadata file, use the upload_editing_dataset_from_csv() function, which supports mixed image formats and flexible directory structures. The CSV format allows custom column name mappings to accommodate different dataset structure requirements.

⚠️ Important: Before using this framework, you must prepare your dataset. See the Data Preparation guide for step-by-step instructions.

Quick Start

Prerequisites

Python 3.12+
CUDA 12.0+ (for GPU training)
18GB+ VRAM recommended Other environment may works as well but did not test yet.

Requirement Installation

# Clone repository
git clone https://github.com/yourusername/qwen-image-finetune.git
cd qwen-image-finetune

# Automated setup
./setup.sh

# Or with custom path and HF token
./setup.sh /your/path hf_your_token_here

⚠️ Important for Qwen-Image-Edit-Plus (2509) Users: If you plan to use Qwen-Image-Edit-Plus (2509) model, you need to install the latest version of diffusers from GitHub:

pip install --upgrade "git+https://github.com/huggingface/diffusers.git"

If you don't need Qwen-Image-Edit-Plus (2509) support, you can use the older version specified in requirements.txt (diffusers>=0.36.0).

Refer docs/guide/speed_optimization.md to install flash-attn to accelerate training. It provides the greatest benefit with long prompts or large sequence lengths; for short prompts, the speedup may be limited.

Train with Toy Dataset

⚠️ Important: All training commands must be executed from the src/ directory. You need to cd src/ before running any python3 -m qflux.main commands. Otherwise, you will encounter import errors.

prepare the datasets or use Hugging Face dataset (recommended). Refer tests/test_configs/test_example_fluxkontext_fp16.yaml
prepare your config. Now suppose you have the config Chose your model, optimizer, etc.
(Optional) build cache first to speed up training (recommended) It save the GPU memory since in the training, you dont need image encoder and prompt encoder anymore if you have the cache.

⚠️ Important: You must run the command from the src/ directory. Otherwise, you will encounter import errors.

cd src/
python3 -m qflux.main --config ../configs/my_config.yaml --cache

The GPU devices used in cache are specified in the config as well. For example

cache:
  devices:
    vae: cuda:1
    text_encoder: cuda:0
    text_encoder_2: cuda:2
  cache_dir: ${logging.output_dir}/${logging.tracker_project_name}/cache
  use_cache: true
  prompt_empty_drop_keys:
    - prompt_embeds
    - pooled_prompt_embeds

Here vae encoder and text encoders could use different GPU ids if your GPU memory is not enough.

start training

⚠️ Important: You must run the command from the src/ directory. Otherwise, you will encounter import errors.

Prepaare a accelerate_config to specify single gpu training or multi-gpu training

# three gpu training using accelerate
cd src/
CUDA_VISIBLE_DEVICES=1,2,4 accelerate launch --config_file ../accelerate_config.yaml -m qflux.main --config ../configs/my_config.yaml

Or do not use accelerate_config.yaml and specify the accelerate parameters in the bash script directly Looks like

cd src/
CUDA_VISIBLE_DEVICES=0,1 \
accelerate launch \
  --num_processes 2 \
  --mixed_precision bf16 \
  -m qflux.main --config $config_file

resume training In the config file add the resumed checkpoint folder

...
resume: <path_to_checkpoint_folder>
...

Then run the script same as above

Configuration Guide

The framework provides various pre-configured training setups for different models and hardware requirements:

Config File	Model	Precision	Key Features	GPU Memory	Recommended GPU	fps (second/batch)
fluxkontext fp16 character composition	Flux-Kontext	BF16	Multi Control Image Lora Training	A100	26G	2.9
fluxkontext fp16 face segmentation	Flux-Kontext	FP16	Standard Lora Training	A100	27G	3.4
qwen-image-edit fp16 character composition	Qwen-Image-Edit	FP16	Multi Control Image Lora Training	A100	42G	2.8
qwen-image-edit fp16 face segmentation	Qwen-Image-Edit	FP16	Standard Lora Training	A100	43G	3.8
qwen-iamge-edit-plus character composition	Qwen-Image-Edit-Plus	fp4	fp4 lora training	A100	33	3.8
qwen-iamge-edit-plus fp4 face segmentation	Qwen-Image-Edit-Plus	fp4	fp4 lora training	A100	27.9	3.6

GPU recommended with the following settings:

batchsize: 2
gradient-checkpoint: True
Adam8bit
image shape
- character_composition: [[384, 672], [512,512]] 或使用像素约束 controls_pixels: [512*512]
- face-segmentation: [[832, 576]] 或使用像素约束 controls_pixels: [512*512]

Usage Example:

⚠️ Important: You must run the command from the src/ directory. Otherwise, you will encounter import errors.

# For FLUX Kontext FP4 training on RTX 4090
cd src/
CUDA_VISIBLE_DEVICES=0 accelerate launch --config_file ../accelerate_config.yaml -m qflux.main --config $config_file

See doc docs/configuration.md for more details about the configs

FSDP training

Setup In the accelerate config, choose proper distributed_type and choose proper num_processes (the number of gpus you want to use)

distributed_type: FSDP #NO # MULTI_GPU, FSDP
num_processes: 2

Memory cost compare

config	model	dtype	desc	machine	GPU memory	speed (second / batch)
qwen-config	Qwen-Image-Edit	FP16	Standard Lora Training	A100x1	43G	3.8
qwen-config	Qwen-Image-Edit	FP16	Standard Lora Training	A100x2(DDP)	55G	3.1
qwen-config	Qwen-Image-Edit	FP16	Standard Lora Training	A100x3(DDP)	55G	3.41
qwen-config	Qwen-Image-Edit	FP16	Standard Lora Training	A100x2(FSDP)	38G	4.66
qwen-config	Qwen-Image-Edit	FP16	Standard Lora Training	A100x3(FSDP)	22.2G	4.8

Parameter Summary Information over different local ranks

Rank-0	Rank-1	Rank-2

Training with RTX4090

⚠️ Important: You must run the command from the src/ directory. Otherwise, you will encounter import errors.

# Config Example
# configs/face_seg_fp4_4090.yaml
cd src/
NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1 \
CUDA_VISIBLE_DEVICES=0,1 accelerate launch --config_file ../accelerate_config.yaml -m qflux.main --config $config_file

For multi-gpu training, need to set

distributed_type: MULTI_GPU  #for multi-gpu training

in the accelerate_config.yaml

Validation Sampling During Training

The framework provides built-in validation sampling during training, allowing you to monitor your model's progress in real-time. This feature is configured in your YAML file:

validation:
  enabled: true        # Enable validation sampling
  steps: 100           # Run validation every 100 steps
  max_samples: 2       # Use up to 2 samples from validation dataset
  seed: 42             # Fixed seed for reproducible results
  dataset:             # Validation dataset configuration
    # Dataset configuration similar to training dataset

See the Validation Sampling Guide for detailed configuration options.

Launch TensorBoard to view the validation results:

tensorboard --logdir=/path/to/output/logs

Finetune Examples

Single Control

Qwen-Image-Edit LoRA Fine-tuning Results

This project demonstrates fine-tuning the Qwen-VL model for face segmentation tasks. Below shows the comparison between pre and post fine-tuning results:

Input Image	Base Model	LoRA Fine-tuned Result
Original input	Base Qwen-Image-Edit model (bfloat16)	LoRA fine-tuned model in bf16
Original input	Base Qwen-Image-Edit model（fp4)	LoRA fine-tuned model (Base fp4,Lora bf16)

Experiment Details:

Prompt: "change the image from the face to the face segmentation mask"
Row 1 - BF16 LoRA: Base model (BF16) + LoRA adapters (BF16) - Checkpoint 900 steps on 20 samples
Row 2 - FP4 LoRA: Base model (BF16) + LoRA adapters (FP4 quantized) - Checkpoint 1000 steps on 20 samples
Inference Steps: 20, CFG Scale: 2.5

Key Observations:

Both LoRA variants significantly outperform the base model
BF16 LoRA shows slightly better detail preservation
FP4 quantized LoRA maintains competitive quality while being more memory efficient
Base model uses BF16 precision in both experiments; only the LoRA adapters differ in quantization

🔥 Flux Kontext Compare：FP16, FP8, FP4 LoRA Fine-tuning Results

Input Image	精度类型	Base Model (无LoRA)	LoRA Fine-tuned Model
Original input image	FP16	Base Flux Kontext model (FP16)	FP16 LoRA fine-tuned model
	FP8	FP8 base model (无LoRA)	FP8 LoRA fine-tuned model
	FP4	FP4 base model (无LoRA)	FP4 LoRA fine-tuned model

Multi Control

Qwen Image Edit with multi-controls

Prompt Image 1	Prompt Image 2	Generated Image

Multi Control Examples from TsienDragon/qwen-image-edit-character-composition

Flux Kontext with multi controls

Prompt Image 1	Prompt Image 2	Generated Image

Multi Control Examples from TsienDragon/flux-kontext-character-composition

Speed

cache	Batch Size	Quantization	Gradient Checkpoint	Flash Attention	Device	GPU Used	Training Speed	Num of Process	config example
cache	2	bf16	True	False	A100	48.6 G	18.3 s/it	1	QwenEdit-bf16
cache	2	fp4	True	False	A100	22.47	10.6 s/it	1	QwenEdit-fp4
cache	2	bf16	True	True	A100	50.2 G	10.34 s/it	1	QwenEdit-bf16
cache	2	fp4	True	True	A100	23.7 G	10.8 s/it	1	QwenEdit-bf16
non-cache	2	fp4	True	True	A100	54.8/53.9G	20.1 s/it	2	QwenEdit-fp4-non-cache
cache	2	fp4	True	True	rtx4090	23.3/22.8G	12.8 s/it	2	Qwenedit-fp4
cache	1	fp4	True	True	rtx4090	18.8/17.9G	6.34 s/it	2	Qwenedit-fp4
cache	2	bf16	True	True	A100	31.3 G	6.65 s/ it	1	FLuxKontext-bf16
cache	2	bf16	True	True	A100	31.32/31.32G	6.69 s/ it	2	FLuxKontext-bf16
cache	2	bf16	True	True	A100	31.9/31.9G	6.78 s/ it	2	FLuxKontext-bf16-prodigy-optimizer
cache	2	bf16	True	True	A100	31.8/31.8G	6.77 s/ it	2	FLuxKontext-fp8
cache	2	bf16	True	True	A100	16.3G	8.24 s/ it	1	FLuxKontext-fp4
cache	2	bf16	False	True	A100	OOM	-	1	FLuxKontext-fp4

prodigy-optimizer: parameter free optimizer No need to tune lr any more
4090: train on 4090 need to set NCCL_P2P_DISABLE=1 NCCL_IB_DISABLE=1. Other setting are same

Check this docs for more training guides docs/training.md

Inference

Single Control

Use trainer in this repo

# Inference with trained LoRA model
from qflux.trainer.qwen_image_edit_trainer import QwenImageEditTrainer
from qflux.data.config import load_config_from_yaml
from PIL import Image

# Load configuration
config = load_config_from_yaml("configs/face_seg_config.yaml")
config.model.lora.pretrained_weight = "/path/to/your/lora/weights.safetensors"

# Initialize trainer (LoRA will be loaded automatically in setup_predict)
trainer = QwenImageEditTrainer(config)

# Setup for inference
trainer.setup_predict()

# Load input image
input_image = Image.open("data/face_seg/control_images/060002_4_028450_FEMALE_30.jpg")

# Generate face segmentation
result = trainer.predict(
    prompt_image=input_image,
    prompt="change the image from the face to the face segmentation mask",
    num_inference_steps=20,
    true_cfg_scale=4.0
)
# show the image
result[0]
# Save result
result[0].save("output_segmentation.png")
print("Generated face segmentation saved as output_segmentation.png")

Use diffusers pipeline

import torch
from PIL import Image
from diffusers import QwenImageEditPipeline
from diffusers.utils import load_image

pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=torch.bfloat16,height=512, width=512,
    output_type='np')
pipe.to("cuda:0")

pipe.load_lora_weights("TsienDragon/qwen-image-edit-lora-face-segmentation")
images_out = pipe(prompt_image, prompt,negative_prompt="", num_inference_steps=20, output_type='pil', true_cfg_scale=1.0).images

Multi Control

Notebooks Tutorials

Debug

Record of bugs encountered in docs/debug.md

🤝 Contributing to Documentation

We welcome contributions to improve this documentation:

Found an Error? Open an issue or submit a PR
Missing Information? Suggest additions or improvements
Want to Help? Contact the maintainers for contribution guidelines

Documentation Standards

Use clear, concise language
Include practical examples
Provide complete code snippets
Add troubleshooting sections
Keep content up to date

Star History

Getting Help

Documentation Issues

Missing Information: Check if it's covered in another section
Outdated Content: Open an issue to report outdated information
Unclear Instructions: Suggest improvements via issues or PRs

Technical Support

Training Issues: See Training Guide troubleshooting
Data Preparation: See Data Preparation Guide for dataset setup
HuggingFace Model & Dataset: See HuggingFace Related Guide for cloud datasets and LoRA model management
Optimizer Selection: See Training Guide for available optimizers
FLUX Kontext Training: See Training Guide for multi-precision training
Setup Problems: Check Setup Guide common issues
Performance: Review Cache System optimization
General Questions: Open a GitHub issue with detailed description

External Resources

Related Projects

Community

GitHub Discussions - General discussions and Q&A
Issues - Bug reports and feature requests
Pull Requests - Code contributions

📝 Note: This documentation is continuously updated. Last updated: 2025/09/26

⭐ Tip: Use the navigation links above to jump to specific topics, or browse sequentially for a complete understanding of the framework.

License

This project is licensed under the MIT License. See the LICENSE file for details.

Name		Name	Last commit message	Last commit date
.claude		.claude
.cursor/rules		.cursor/rules
.vscode		.vscode
configs		configs
docs		docs
scripts		scripts
skills		skills
src/qflux		src/qflux
tests		tests
.claudeignore		.claudeignore
.gitignore		.gitignore
.pre-commit-config.yaml		.pre-commit-config.yaml
AGENTS.md		AGENTS.md
CLAUDE.md		CLAUDE.md
LICENSE		LICENSE
README.md		README.md
VERSION		VERSION
accelerate_config.yaml		accelerate_config.yaml
pyproject.toml		pyproject.toml
pytest.ini		pytest.ini
requirements.txt		requirements.txt
run.sh		run.sh
setup.sh		setup.sh

Latest commit

History

Folders and files

Repository files navigation

Image-Edit Fine-tuning Framework

Overview

New

TODO

Key Features

Table of Contents

Documentation

Dataset

Quick Start

Prerequisites

Requirement Installation

Train with Toy Dataset

Configuration Guide

FSDP training

Training with RTX4090

Validation Sampling During Training

Finetune Examples

Single Control

Qwen-Image-Edit LoRA Fine-tuning Results

🔥 Flux Kontext Compare：FP16, FP8, FP4 LoRA Fine-tuning Results

Multi Control

Qwen Image Edit with multi-controls

Flux Kontext with multi controls

Speed

Inference

Single Control

Multi Control

Notebooks Tutorials

Debug

🤝 Contributing to Documentation

Documentation Standards

Star History

Getting Help

Documentation Issues

Technical Support

External Resources

Related Projects

Community

License

About

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages