Instruction-tuning Stable Diffusion with InstructPix2Pix
Hugging Face has demonstrated that Stable Diffusion can be instruction-tuned to perform specific image translation and low-level processing tasks, such as cartoonization and deraining, by adapting the InstructPix2Pix training strategy. This approach allows the model to follow precise natural language instructions to modify input images more faithfully than general-purpose pre-trained models.
Instruction-Tuning for Image Editing
Instruction-tuning is a supervised method used to teach models to solve tasks based on natural language prompts. While originally developed for language models (such as FLAN), Hugging Face applied this philosophy to Stable Diffusion using the InstructPix2Pix framework.
While pre-trained InstructPix2Pix models are capable of following general editing instructions, they often struggle with specific transformations. To address this, the researchers leveraged paired datasets for specialized tasks to create instruction-prompted datasets, enabling the model to learn specific mapping between an instruction and a visual transformation.
Dataset Preparation and Methodology
To train the model, Hugging Face created specialized datasets by pairing input images with target outputs and corresponding natural language instructions.
Cartoonization Dataset
To teach the model to "cartoonize" images, the team implemented a three-step pipeline:
- Instruction Generation: ChatGPT was used to generate 50 synonymous sentences for the instruction "Cartoonize the image."
- Label Generation: A pre-trained Whitebox CartoonGAN model was used to process 5,000 samples from the Imagenette dataset to create the target cartoonized images.
- Exemplar Creation: These components were combined into exemplars consisting of the original image, the instruction, and the cartoonized output.
Low-Level Image Processing Dataset
Following the FLAN multi-task approach, the team constructed a single dataset comprising several low-level tasks. This differs from traditional image processing where models are typically trained independently on individual datasets.
| Task | Prompt | Dataset | Number of Samples |
|---|---|---|---|
| Deblurring | "deblur the blurry image" | REDS | 1,200 |
| Deraining | "derain the image" | Rain13k | 686 |
| Denoising | "denoise the noisy image" | SIDD | 8 |
| Low-light enhancement | "enhance the low-light image" | LOL | 23 |
Training Experiments and Results
Training was conducted using two primary starting points: fine-tuning from an existing InstructPix2Pix checkpoint and fine-tuning from a Stable Diffusion v1-5 checkpoint using the InstructPix2Pix methodology. The researchers found that starting from an InstructPix2Pix checkpoint led to faster adaptation and higher generation quality.
Performance Analysis
- Cartoonization: The instruction-tuned model more faithfully matched the outputs of the CartoonGAN model compared to the pre-trained InstructPix2Pix model.
- Deraining: The model produced compelling results that aligned closely with the ground-truth images.
- Failures: The model struggled with low-light image enhancement and deblurring, which the researchers attribute to an insufficient number of training exemplars for those specific tasks.
- Generalization: The model failed to produce expected outputs for ImageNette classes it had not seen sufficiently during training, suggesting a need for larger-scale datasets.
Limitations and Future Directions
Despite the progress, the researchers identified several critical challenges for realistic image editing applications:
- Resolution: The systems must be scaled to work with large, high-resolution original images.
- Hallucinations: Diffusion models sometimes "invent" or re-interpret instructions, modifying the image space in ways that are unacceptable for professional editing.
Open Questions for the Community
- How does scaling the dataset size (e.g., moving toward the 30,000+ samples used in original InstructPix2Pix) impact quality?
- Does increasing the task mixture or training for longer durations improve generalization to unseen tasks or combined instructions (e.g., "Deblur the image and denoise it")?
- Would using synonymous instructions during training, rather than just during dataset creation, improve performance?
- Could a ControlNet training setup provide better results for these specific image-to-image tasks?