Colorizenet

🚀 彩色化网络（ColorizeNet）模型卡片

本模型是一个基于ControlNet训练的图像彩色化模型，可将黑白图像转换为彩色图像，为黑白影像赋予生动色彩，具有较高的实用价值。

✨ 主要特性

• 基于ControlNet架构，结合预训练的Stable Diffusion 2.1模型进行微调，在图像彩色化任务上表现出色。
• 利用COCO数据集进行训练，数据丰富，模型泛化能力强。

📦 安装指南

文档中未提及具体安装步骤，暂无法提供。

💻 使用示例

基础用法

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)

高级用法

# 读取待彩色化的图像
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()

# 准备模型的输入和参数
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)

# 采样并对结果进行后处理
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]

📚 详细文档

模型详情

模型描述

彩色化网络（ColorizeNet）是一个基于ControlNet的图像彩色化模型，使用了Stability AI提出的预训练Stable Diffusion 2.1模型进行训练。

• 微调基础模型：[https://huggingface.co/stabilityai/stable-diffusion-2-1]

模型来源

• 仓库地址：[https://github.com/rensortino/ColorizeNet]

训练数据

该模型在COCO数据集上进行训练，使用数据集中的所有图像，并将它们转换为灰度图像，以用于控制ControlNet。 [https://huggingface.co/datasets/detection-datasets/coco]

模型运行

实例化模型并加载其配置和权重，然后读取待彩色化的图像，准备模型的输入和参数，最后进行采样和结果后处理。具体代码见使用示例部分。

模型结果

黑白输入	彩色输出

📄 许可证

本模型使用Apache 2.0许可证。