Spaces:
Running
on
Zero
Running
on
Zero
File size: 7,010 Bytes
691af46 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
# Text-Guided-Image-Colorization
This project utilizes the power of **Stable Diffusion (SDXL/SDXL-Light)** and the **BLIP (Bootstrapping Language-Image Pre-training)** captioning model to provide an interactive image colorization experience. Users can influence the generated colors of objects within images, making the colorization process more personalized and creative.
## Table of Contents
- [Features](#features)
- [Installation](#installation)
- [Quick Start](#quick-start)
- [Dataset Usage](#dataset-usage)
- [Training](#training)
- [Evaluation](#evaluation)
- [Results](#results)
- [License](#license)
## Features
- **Interactive Colorization**: Users can specify desired colors for different objects in the image.
- **ControlNet Approach**: Enhanced colorization capabilities through retraining with ControlNet, allowing SDXL to better adapt to the image colorization task.
- **High-Quality Outputs**: Leverage the latest advancements in diffusion models to generate vibrant and realistic colorizations.
- **User-Friendly Interface**: Easy-to-use interface for seamless interaction with the model.
## Installation
To set up the project locally, follow these steps:
1. **Clone the Repository**:
```bash
git clone https://github.com/nick8592/text-guided-image-colorization.git
cd text-guided-image-colorization
```
2. **Install Dependencies**:
Make sure you have Python 3.7 or higher installed. Then, install the required packages:
```bash
pip install -r requirements.txt
```
Install `torch` and `torchvision` matching your CUDA version:
```bash
pip install torch torchvision --index-url https://download.pytorch.org/whl/cuXXX
```
Replace `XXX` with your CUDA version (e.g., `118` for CUDA 11.8). For more info, see [PyTorch Get Started](https://pytorch.org/get-started/locally/).
3. **Download Pre-trained Models**:
| Models | Hugging Face (Recommand) | Other |
|:---:|:---:|:---:|
|SDXL-Lightning Caption|[link](https://huggingface.co/nickpai/sdxl_light_caption_output)|[link](https://gofile.me/7uE8s/FlEhfpWPw) (2kNJfV)|
|SDXL-Lightning Custom Caption (Recommand)|[link](https://huggingface.co/nickpai/sdxl_light_custom_caption_output)|[link](https://gofile.me/7uE8s/AKmRq5sLR) (KW7Fpi)|
```bash
text-guided-image-colorization/sdxl_light_caption_output
└── checkpoint-30000
├── controlnet
│ ├── diffusion_pytorch_model.safetensors
│ └── config.json
├── optimizer.bin
├── random_states_0.pkl
├── scaler.pt
└── scheduler.bin
```
## Quick Start
1. Run the `gradio_ui.py` script:
```bash
python gradio_ui.py
```
2. Open the provided URL in your web browser to access the Gradio-based user interface.
3. Upload an image and use the interface to control the colors of specific objects in the image. But still the model can generate images without a specific prompt.
4. The model will generate a colorized version of the image based on your input (or automatic). See the [demo video](https://x.com/weichenpai/status/1829513077588631987).
## Dataset Usage
You can find more details about the dataset usage in the [Dataset-for-Image-Colorization](https://github.com/nick8592/Dataset-for-Image-Colorization).
## Training
For training, you can use one of the following scripts:
- `train_controlnet.sh`: Trains a model using [Stable Diffusion v2](https://huggingface.co/stabilityai/stable-diffusion-2-1)
- `train_controlnet_sdxl.sh`: Trains a model using [SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0)
- `train_controlnet_sdxl_light.sh`: Trains a model using [SDXL-Lightning](https://huggingface.co/ByteDance/SDXL-Lightning)
Although the training code for SDXL is provided, due to a lack of GPU resources, I wasn't able to train the model by myself. Therefore, there might be some errors when you try to train the model.
## Evaluation
For evaluation, you can use one of the following scripts:
- `eval_controlnet.sh`: Evaluates the model using [Stable Diffusion v2](https://huggingface.co/stabilityai/stable-diffusion-2-1) for a folder of images.
- `eval_controlnet_sdxl_light.sh`: Evaluates the model using [SDXL-Lightning](https://huggingface.co/ByteDance/SDXL-Lightning) for a folder of images.
- `eval_controlnet_sdxl_light_single.sh`: Evaluates the model using [SDXL-Lightning](https://huggingface.co/ByteDance/SDXL-Lightning) for a single image.
## Results
### Prompt-Guided
| Caption | Condition 1 | Condition 2 | Condition 3 |
|:---:|:---:|:---:|:---:|
|  |  |  | |
| a photography of a woman in a soccer uniform kicking a soccer ball | + "green shirt"| + "purple shirt" | + "red shirt" |
|  |  |  | |
| a photography of a photo of a truck | + "bright red car"| + "dark blue car" | + "black car" |
|  |  |  | |
| a photography of a cat wearing a hat on his head | + "orange hat"| + "pink hat" | + "yellow hat" |
### Prompt-Free
Ground truth images are provided solely for reference purpose in the image colorization task.
| Grayscale Image | Colorized Result | Ground Truth |
|:---:|:---:|:---:|
|  |  |  |
|  |  |  |
|  |  |  |
|  |  |  |
|  |  |  |
## License
This project is licensed under the MIT License. See the [LICENSE](LICENSE) file for more details.
|