Colorizenet

🚀 ColorizeNet Model Card

ColorizeNet is a model designed for image colorization. It leverages ControlNet technology to transform black - and - white images into colorful ones, offering a practical solution for enhancing old or monochrome images.

✨ Features

• Based on ControlNet for effective image colorization.
• Trained using the pre - trained Stable Diffusion model version 2.1.
• Capable of processing images from the COCO dataset.

📦 Installation

There is no explicit installation information provided in the original document. So this section is skipped.

💻 Usage Examples

Basic Usage

import random

import cv2
import einops
import numpy as np
import torch
from pytorch_lightning import seed_everything

from utils.data import HWC3, apply_color, resize_image
from utils.ddim import DDIMSampler
from utils.model import create_model, load_state_dict

model = create_model('./models/cldm_v21.yaml').cpu()
model.load_state_dict(load_state_dict(
    'lightning_logs/version_6/checkpoints/colorizenet-sd21.ckpt', location='cuda'))
model = model.cuda()
ddim_sampler = DDIMSampler(model)

Advanced Usage

# Read the image to be colorized
input_image = cv2.imread("sample_data/sample1_bw.jpg")
input_image = HWC3(input_image)
img = resize_image(input_image, resolution=512)
H, W, C = img.shape

num_samples = 1
control = torch.from_numpy(img.copy()).float().cuda() / 255.0
control = torch.stack([control for _ in range(num_samples)], dim=0)
control = einops.rearrange(control, 'b h w c -> b c h w').clone()

# Prepare the input and parameters of the model
seed = 1294574436
seed_everything(seed)
prompt = "Colorize this image"
n_prompt = ""
guess_mode = False
strength = 1.0
eta = 0.0
ddim_steps = 20
scale = 9.0

cond = {"c_concat": [control], "c_crossattn": [
    model.get_learned_conditioning([prompt] * num_samples)]}
un_cond = {"c_concat": None if guess_mode else [control], "c_crossattn": [
    model.get_learned_conditioning([n_prompt] * num_samples)]}
shape = (4, H // 8, W // 8)

model.control_scales = [strength * (0.825 ** float(12 - i)) for i in range(13)] if guess_mode else (
    [strength] * 13)

# Sample and post - process the results
samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
                                             shape, cond, verbose=False, eta=eta,
                                             unconditional_guidance_scale=scale,
                                             unconditional_conditioning=un_cond)

x_samples = model.decode_first_stage(samples)
x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c')
             * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)

results = [x_samples[i] for i in range(num_samples)]
colored_results = [apply_color(img, result) for result in results]

📚 Documentation

Model Details

Model Description

ColorizeNet is an image colorization model based on ControlNet, trained using the pre - trained Stable Diffusion model version 2.1 proposed by Stability AI.

• Finetuned from model : [https://huggingface.co/stabilityai/stable-diffusion-2-1]

Model Sources

• Repository: [https://github.com/rensortino/ColorizeNet]

Training Data

The model has been trained on COCO, using all the images in the dataset and converting them to grayscale to use them to condition the ControlNet. [https://huggingface.co/datasets/detection-datasets/coco]

🔧 Technical Details

The model is a ControlNet - based image colorization model. It uses the pre - trained Stable Diffusion model version 2.1. The training data is from the COCO dataset, where all images are converted to grayscale to condition the ControlNet. The model is instantiated using the create_model function and its weights are loaded from a checkpoint file. The sampling process uses the DDIMSampler to generate colorized images from black - and - white inputs.

📄 License

This model is licensed under the Apache - 2.0 license.

Results

BW Input	Colorized