🚀 彩色化網絡(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許可證。