| 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] |
| [cv2.imwrite(f"colorized_{i}.jpg", cv2.cvtColor(result, cv2.COLOR_RGB2BGR)) for i, result in enumerate(colored_results)] |
