Flux Pipeline Tuning with Torch Tweak
This example demonstrates how to use Torch Tweak to tune the Flux text-to-image model from Hugging Face's diffusers library.
Model Details
The Flux model is a text-to-image diffusion model that generates high-quality images from text descriptions. The model is trained on a large dataset of images and text, and can generate realistic images across various domains.
For more information, visit the Flux model page on HuggingFace.
Environment Setup
You can use either of the following options to set up the environment:
Option 1 - virtual environment managed by you
Activate your virtual environment and install the dependencies:
pip install --index-url https://download.pytorch.org/whl/xpu \
--extra-index-url https://pypi.org/simple \
.
Option 2 - virtual environment managed by uv
Install dependencies:
Usage
Tuning the model
To tune the Flux model, run:
tune --model-name black-forest-labs/FLUX.2-klein-4B --prompt "A futuristic cityscape with neon lights"
You can customize the following parameters:
- --model-name: HuggingFace model name or path (default: "black-forest-labs/FLUX.2-klein-4B")
- --prompt: Text prompt for image generation
- --negative-prompt: Negative text prompt (default: "low quality, blurry")
- --height: Height of the generated image (default: 512)
- --width: Width of the generated image (default: 512)
- --steps: Number of inference steps (default: 25)
Generating images with the tuned model
After tuning, generate images with:
The generated image will be saved in the specified output directory.
Dynamic batching
The service uses dynamic batching - requests are grouped and processed together for efficiency. Currently, there is one frontend and one worker. To support multiple workers, move batching to a separate service that handles request grouping.