mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-18 07:41:03 +00:00
add rerender lora
This commit is contained in:
salt
@@ -0,0 +1,474 @@
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import einops
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchvision.transforms as T
|
||||
from blendmodes.blend import BlendType, blendLayers
|
||||
from PIL import Image
|
||||
from pytorch_lightning import seed_everything
|
||||
from safetensors.torch import load_file
|
||||
from skimage import exposure
|
||||
|
||||
import src.import_util # noqa: F401
|
||||
from deps.ControlNet.annotator.canny import CannyDetector
|
||||
from deps.ControlNet.annotator.hed import HEDdetector
|
||||
from deps.ControlNet.annotator.util import HWC3
|
||||
from deps.ControlNet.cldm.cldm import ControlLDM
|
||||
from deps.ControlNet.cldm.model import create_model, load_state_dict
|
||||
from deps.gmflow.gmflow.gmflow import GMFlow
|
||||
from flow.flow_utils import get_warped_and_mask
|
||||
from src.config import RerenderConfig
|
||||
from src.controller import AttentionControl
|
||||
from src.ddim_v_hacked import DDIMVSampler
|
||||
from src.freeu import freeu_forward
|
||||
from src.img_util import find_flat_region, numpy2tensor
|
||||
from src.video_util import frame_to_video, get_fps, prepare_frames
|
||||
from src.lora import load_lora, apply_lora
|
||||
|
||||
blur = T.GaussianBlur(kernel_size=(9, 9), sigma=(18, 18))
|
||||
totensor = T.PILToTensor()
|
||||
|
||||
|
||||
def setup_color_correction(image):
|
||||
correction_target = cv2.cvtColor(np.asarray(image.copy()),
|
||||
cv2.COLOR_RGB2LAB)
|
||||
return correction_target
|
||||
|
||||
|
||||
def apply_color_correction(correction, original_image):
|
||||
image = Image.fromarray(
|
||||
cv2.cvtColor(
|
||||
exposure.match_histograms(cv2.cvtColor(np.asarray(original_image),
|
||||
cv2.COLOR_RGB2LAB),
|
||||
correction,
|
||||
channel_axis=2),
|
||||
cv2.COLOR_LAB2RGB).astype('uint8'))
|
||||
|
||||
image = blendLayers(image, original_image, BlendType.LUMINOSITY)
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def rerender(cfg: RerenderConfig, first_img_only: bool, key_video_path: str):
|
||||
# Preprocess input
|
||||
prepare_frames(cfg.input_path, cfg.input_dir, cfg.image_resolution, cfg.crop, cfg.use_limit_device_resolution)
|
||||
|
||||
# Load models
|
||||
if cfg.control_type == 'HED':
|
||||
detector = HEDdetector()
|
||||
elif cfg.control_type == 'canny':
|
||||
canny_detector = CannyDetector()
|
||||
low_threshold = cfg.canny_low
|
||||
high_threshold = cfg.canny_high
|
||||
|
||||
def apply_canny(x):
|
||||
return canny_detector(x, low_threshold, high_threshold)
|
||||
|
||||
detector = apply_canny
|
||||
|
||||
model: ControlLDM = create_model(
|
||||
'./deps/ControlNet/models/cldm_v15.yaml').cpu()
|
||||
if cfg.control_type == 'HED':
|
||||
model.load_state_dict(
|
||||
load_state_dict('./models/control_sd15_hed.pth', location='cuda'))
|
||||
elif cfg.control_type == 'canny':
|
||||
model.load_state_dict(
|
||||
load_state_dict('./models/control_sd15_canny.pth',
|
||||
location='cuda'))
|
||||
model = model.cuda()
|
||||
model.control_scales = [cfg.control_strength] * 13
|
||||
|
||||
if cfg.sd_model is not None:
|
||||
model_ext = os.path.splitext(cfg.sd_model)[1]
|
||||
if model_ext == '.safetensors':
|
||||
model.load_state_dict(load_file(cfg.sd_model), strict=False)
|
||||
elif model_ext == '.ckpt' or model_ext == '.pth':
|
||||
model.load_state_dict(torch.load(cfg.sd_model)['state_dict'],
|
||||
strict=False)
|
||||
|
||||
# apply lora if exists
|
||||
if cfg.lora_path:
|
||||
lora_weights = load_lora(cfg.lora_path)
|
||||
apply_lora(model.model, model.cond_stage_model, lora_weights)
|
||||
|
||||
try:
|
||||
model.first_stage_model.load_state_dict(torch.load(
|
||||
'./models/vae-ft-mse-840000-ema-pruned.ckpt')['state_dict'],
|
||||
strict=False)
|
||||
except Exception:
|
||||
print('Warning: We suggest you download the fine-tuned VAE',
|
||||
'otherwise the generation quality will be degraded')
|
||||
|
||||
model.model.diffusion_model.forward = \
|
||||
freeu_forward(model.model.diffusion_model, *cfg.freeu_args)
|
||||
ddim_v_sampler = DDIMVSampler(model)
|
||||
|
||||
flow_model = GMFlow(
|
||||
feature_channels=128,
|
||||
num_scales=1,
|
||||
upsample_factor=8,
|
||||
num_head=1,
|
||||
attention_type='swin',
|
||||
ffn_dim_expansion=4,
|
||||
num_transformer_layers=6,
|
||||
).to('cuda')
|
||||
|
||||
checkpoint = torch.load('models/gmflow_sintel-0c07dcb3.pth',
|
||||
map_location=lambda storage, loc: storage)
|
||||
weights = checkpoint['model'] if 'model' in checkpoint else checkpoint
|
||||
flow_model.load_state_dict(weights, strict=False)
|
||||
flow_model.eval()
|
||||
|
||||
num_samples = 1
|
||||
ddim_steps = 20
|
||||
scale = 7.5
|
||||
|
||||
seed = cfg.seed
|
||||
if seed == -1:
|
||||
seed = random.randint(0, 65535)
|
||||
eta = 0.0
|
||||
|
||||
prompt = cfg.prompt
|
||||
a_prompt = cfg.a_prompt
|
||||
n_prompt = cfg.n_prompt
|
||||
prompt = prompt + ', ' + a_prompt
|
||||
|
||||
style_update_freq = cfg.style_update_freq
|
||||
pixelfusion = True
|
||||
color_preserve = cfg.color_preserve
|
||||
|
||||
x0_strength = 1 - cfg.x0_strength
|
||||
mask_period = cfg.mask_period
|
||||
firstx0 = True
|
||||
controller = AttentionControl(cfg.inner_strength, cfg.mask_period,
|
||||
cfg.cross_period, cfg.ada_period,
|
||||
cfg.warp_period, cfg.loose_cfattn)
|
||||
|
||||
imgs = sorted(os.listdir(cfg.input_dir))
|
||||
imgs = [os.path.join(cfg.input_dir, img) for img in imgs]
|
||||
if cfg.frame_count >= 0:
|
||||
imgs = imgs[:cfg.frame_count]
|
||||
|
||||
with torch.no_grad():
|
||||
frame = cv2.imread(imgs[0])
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
img = HWC3(frame)
|
||||
H, W, C = img.shape
|
||||
|
||||
img_ = numpy2tensor(img)
|
||||
# if color_preserve:
|
||||
# img_ = numpy2tensor(img)
|
||||
# else:
|
||||
# img_ = apply_color_correction(color_corrections,
|
||||
# Image.fromarray(img))
|
||||
# img_ = totensor(img_).unsqueeze(0)[:, :3] / 127.5 - 1
|
||||
encoder_posterior = model.encode_first_stage(img_.cuda())
|
||||
x0 = model.get_first_stage_encoding(encoder_posterior).detach()
|
||||
|
||||
detected_map = detector(img)
|
||||
detected_map = HWC3(detected_map)
|
||||
# For visualization
|
||||
detected_img = 255 - detected_map
|
||||
|
||||
control = torch.from_numpy(detected_map.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()
|
||||
cond = {
|
||||
'c_concat': [control],
|
||||
'c_crossattn':
|
||||
[model.get_learned_conditioning([prompt] * num_samples)]
|
||||
}
|
||||
un_cond = {
|
||||
'c_concat': [control],
|
||||
'c_crossattn':
|
||||
[model.get_learned_conditioning([n_prompt] * num_samples)]
|
||||
}
|
||||
shape = (4, H // 8, W // 8)
|
||||
|
||||
controller.set_task('initfirst')
|
||||
seed_everything(seed)
|
||||
samples, _ = ddim_v_sampler.sample(ddim_steps,
|
||||
num_samples,
|
||||
shape,
|
||||
cond,
|
||||
verbose=False,
|
||||
eta=eta,
|
||||
unconditional_guidance_scale=scale,
|
||||
unconditional_conditioning=un_cond,
|
||||
controller=controller,
|
||||
x0=x0,
|
||||
strength=x0_strength)
|
||||
x_samples = model.decode_first_stage(samples)
|
||||
pre_result = x_samples
|
||||
pre_img = img
|
||||
first_result = pre_result
|
||||
first_img = pre_img
|
||||
|
||||
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)
|
||||
color_corrections = setup_color_correction(Image.fromarray(x_samples[0]))
|
||||
Image.fromarray(x_samples[0]).save(os.path.join(cfg.first_dir,
|
||||
'first.jpg'))
|
||||
cv2.imwrite(os.path.join(cfg.first_dir, 'first_edge.jpg'), detected_img)
|
||||
|
||||
if first_img_only:
|
||||
exit(0)
|
||||
|
||||
for i in range(0, min(len(imgs), cfg.frame_count) - 1, cfg.interval):
|
||||
cid = i + 1
|
||||
print(cid)
|
||||
if cid <= (len(imgs) - 1):
|
||||
frame = cv2.imread(imgs[cid])
|
||||
else:
|
||||
frame = cv2.imread(imgs[len(imgs) - 1])
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
img = HWC3(frame)
|
||||
|
||||
if color_preserve:
|
||||
img_ = numpy2tensor(img)
|
||||
else:
|
||||
img_ = apply_color_correction(color_corrections,
|
||||
Image.fromarray(img))
|
||||
img_ = totensor(img_).unsqueeze(0)[:, :3] / 127.5 - 1
|
||||
encoder_posterior = model.encode_first_stage(img_.cuda())
|
||||
x0 = model.get_first_stage_encoding(encoder_posterior).detach()
|
||||
|
||||
detected_map = detector(img)
|
||||
detected_map = HWC3(detected_map)
|
||||
|
||||
control = torch.from_numpy(detected_map.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()
|
||||
cond['c_concat'] = [control]
|
||||
un_cond['c_concat'] = [control]
|
||||
|
||||
image1 = torch.from_numpy(pre_img).permute(2, 0, 1).float()
|
||||
image2 = torch.from_numpy(img).permute(2, 0, 1).float()
|
||||
warped_pre, bwd_occ_pre, bwd_flow_pre = get_warped_and_mask(
|
||||
flow_model, image1, image2, pre_result, False)
|
||||
blend_mask_pre = blur(
|
||||
F.max_pool2d(bwd_occ_pre, kernel_size=9, stride=1, padding=4))
|
||||
blend_mask_pre = torch.clamp(blend_mask_pre + bwd_occ_pre, 0, 1)
|
||||
|
||||
image1 = torch.from_numpy(first_img).permute(2, 0, 1).float()
|
||||
warped_0, bwd_occ_0, bwd_flow_0 = get_warped_and_mask(
|
||||
flow_model, image1, image2, first_result, False)
|
||||
blend_mask_0 = blur(
|
||||
F.max_pool2d(bwd_occ_0, kernel_size=9, stride=1, padding=4))
|
||||
blend_mask_0 = torch.clamp(blend_mask_0 + bwd_occ_0, 0, 1)
|
||||
|
||||
if firstx0:
|
||||
mask = 1 - F.max_pool2d(blend_mask_0, kernel_size=8)
|
||||
controller.set_warp(
|
||||
F.interpolate(bwd_flow_0 / 8.0,
|
||||
scale_factor=1. / 8,
|
||||
mode='bilinear'), mask)
|
||||
else:
|
||||
mask = 1 - F.max_pool2d(blend_mask_pre, kernel_size=8)
|
||||
controller.set_warp(
|
||||
F.interpolate(bwd_flow_pre / 8.0,
|
||||
scale_factor=1. / 8,
|
||||
mode='bilinear'), mask)
|
||||
|
||||
controller.set_task('keepx0, keepstyle')
|
||||
seed_everything(seed)
|
||||
samples, intermediates = ddim_v_sampler.sample(
|
||||
ddim_steps,
|
||||
num_samples,
|
||||
shape,
|
||||
cond,
|
||||
verbose=False,
|
||||
eta=eta,
|
||||
unconditional_guidance_scale=scale,
|
||||
unconditional_conditioning=un_cond,
|
||||
controller=controller,
|
||||
x0=x0,
|
||||
strength=x0_strength)
|
||||
direct_result = model.decode_first_stage(samples)
|
||||
|
||||
if not pixelfusion:
|
||||
pre_result = direct_result
|
||||
pre_img = img
|
||||
viz = (
|
||||
einops.rearrange(direct_result, 'b c h w -> b h w c') * 127.5 +
|
||||
127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
|
||||
|
||||
else:
|
||||
|
||||
blend_results = (1 - blend_mask_pre
|
||||
) * warped_pre + blend_mask_pre * direct_result
|
||||
blend_results = (
|
||||
1 - blend_mask_0) * warped_0 + blend_mask_0 * blend_results
|
||||
|
||||
bwd_occ = 1 - torch.clamp(1 - bwd_occ_pre + 1 - bwd_occ_0, 0, 1)
|
||||
blend_mask = blur(
|
||||
F.max_pool2d(bwd_occ, kernel_size=9, stride=1, padding=4))
|
||||
blend_mask = 1 - torch.clamp(blend_mask + bwd_occ, 0, 1)
|
||||
|
||||
encoder_posterior = model.encode_first_stage(blend_results)
|
||||
xtrg = model.get_first_stage_encoding(
|
||||
encoder_posterior).detach() # * mask
|
||||
blend_results_rec = model.decode_first_stage(xtrg)
|
||||
encoder_posterior = model.encode_first_stage(blend_results_rec)
|
||||
xtrg_rec = model.get_first_stage_encoding(
|
||||
encoder_posterior).detach()
|
||||
xtrg_ = (xtrg + 1 * (xtrg - xtrg_rec)) # * mask
|
||||
blend_results_rec_new = model.decode_first_stage(xtrg_)
|
||||
tmp = (abs(blend_results_rec_new - blend_results).mean(
|
||||
dim=1, keepdims=True) > 0.25).float()
|
||||
mask_x = F.max_pool2d((F.interpolate(
|
||||
tmp, scale_factor=1 / 8., mode='bilinear') > 0).float(),
|
||||
kernel_size=3,
|
||||
stride=1,
|
||||
padding=1)
|
||||
|
||||
mask = (1 - F.max_pool2d(1 - blend_mask, kernel_size=8)
|
||||
) # * (1-mask_x)
|
||||
|
||||
if cfg.smooth_boundary:
|
||||
noise_rescale = find_flat_region(mask)
|
||||
else:
|
||||
noise_rescale = torch.ones_like(mask)
|
||||
masks = []
|
||||
for j in range(ddim_steps):
|
||||
if j <= ddim_steps * mask_period[
|
||||
0] or j >= ddim_steps * mask_period[1]:
|
||||
masks += [None]
|
||||
else:
|
||||
masks += [mask * cfg.mask_strength]
|
||||
|
||||
# mask 3
|
||||
# xtrg = ((1-mask_x) *
|
||||
# (xtrg + xtrg - xtrg_rec) + mask_x * samples) * mask
|
||||
# mask 2
|
||||
# xtrg = (xtrg + 1 * (xtrg - xtrg_rec)) * mask
|
||||
xtrg = (xtrg + (1 - mask_x) * (xtrg - xtrg_rec)) * mask # mask 1
|
||||
|
||||
tasks = 'keepstyle, keepx0'
|
||||
if not firstx0:
|
||||
tasks += ', updatex0'
|
||||
if i % style_update_freq == 0:
|
||||
tasks += ', updatestyle'
|
||||
controller.set_task(tasks, 1.0)
|
||||
|
||||
seed_everything(seed)
|
||||
samples, _ = ddim_v_sampler.sample(
|
||||
ddim_steps,
|
||||
num_samples,
|
||||
shape,
|
||||
cond,
|
||||
verbose=False,
|
||||
eta=eta,
|
||||
unconditional_guidance_scale=scale,
|
||||
unconditional_conditioning=un_cond,
|
||||
controller=controller,
|
||||
x0=x0,
|
||||
strength=x0_strength,
|
||||
xtrg=xtrg,
|
||||
mask=masks,
|
||||
noise_rescale=noise_rescale)
|
||||
x_samples = model.decode_first_stage(samples)
|
||||
pre_result = x_samples
|
||||
pre_img = img
|
||||
|
||||
viz = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 +
|
||||
127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
|
||||
|
||||
Image.fromarray(viz[0]).save(
|
||||
os.path.join(cfg.key_dir, f'{cid:04d}.png'))
|
||||
if key_video_path is not None:
|
||||
fps = get_fps(cfg.input_path)
|
||||
fps //= cfg.interval
|
||||
frame_to_video(key_video_path, cfg.key_dir, fps, False)
|
||||
|
||||
|
||||
def postprocess(cfg: RerenderConfig, ne: bool, max_process: int, tmp: bool,
|
||||
ps: bool):
|
||||
video_base_dir = cfg.work_dir
|
||||
o_video = cfg.output_path
|
||||
fps = get_fps(cfg.input_path)
|
||||
|
||||
end_frame = cfg.frame_count - 1
|
||||
interval = cfg.interval
|
||||
key_dir = os.path.split(cfg.key_dir)[-1]
|
||||
use_e = '-ne' if ne else ''
|
||||
use_tmp = '-tmp' if tmp else ''
|
||||
use_ps = '-ps' if ps else ''
|
||||
o_video_cmd = f'--output {o_video}'
|
||||
|
||||
cmd = (
|
||||
f'python video_blend.py {video_base_dir} --beg 1 --end {end_frame} '
|
||||
f'--itv {interval} --key {key_dir} {use_e} {o_video_cmd} --fps {fps} '
|
||||
f'--n_proc {max_process} {use_tmp} {use_ps}')
|
||||
print(cmd)
|
||||
os.system(cmd)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--cfg', type=str, default=None)
|
||||
parser.add_argument('--input',
|
||||
type=str,
|
||||
default=None,
|
||||
help='The input path to video.')
|
||||
parser.add_argument('--output', type=str, default=None)
|
||||
parser.add_argument('--prompt', type=str, default=None)
|
||||
parser.add_argument('--key_video_path', type=str, default=None)
|
||||
parser.add_argument('-one',
|
||||
action='store_true',
|
||||
help='Run the first frame with ControlNet only')
|
||||
parser.add_argument('-nr',
|
||||
action='store_true',
|
||||
help='Do not run rerender and do postprocessing only')
|
||||
parser.add_argument('-nb',
|
||||
action='store_true',
|
||||
help='Do not run postprocessing and run rerender only')
|
||||
parser.add_argument(
|
||||
'-ne',
|
||||
action='store_true',
|
||||
help='Do not run ebsynth (use previous ebsynth temporary output)')
|
||||
parser.add_argument('-nps',
|
||||
action='store_true',
|
||||
help='Do not run poisson gradient blending')
|
||||
parser.add_argument('--n_proc',
|
||||
type=int,
|
||||
default=4,
|
||||
help='The max process count')
|
||||
parser.add_argument('--tmp',
|
||||
action='store_true',
|
||||
help='Keep ebsynth temporary output')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
cfg = RerenderConfig()
|
||||
if args.cfg is not None:
|
||||
cfg.create_from_path(args.cfg)
|
||||
if args.input is not None:
|
||||
print('Config has been loaded. --input is ignored.')
|
||||
if args.output is not None:
|
||||
print('Config has been loaded. --output is ignored.')
|
||||
if args.prompt is not None:
|
||||
print('Config has been loaded. --prompt is ignored.')
|
||||
else:
|
||||
if args.input is None:
|
||||
print('Config not found. --input is required.')
|
||||
exit(0)
|
||||
if args.output is None:
|
||||
print('Config not found. --output is required.')
|
||||
exit(0)
|
||||
if args.prompt is None:
|
||||
print('Config not found. --prompt is required.')
|
||||
exit(0)
|
||||
cfg.create_from_parameters(args.input, args.output, args.prompt)
|
||||
|
||||
if not args.nr:
|
||||
rerender(cfg, args.one, args.key_video_path)
|
||||
torch.cuda.empty_cache()
|
||||
if not args.nb:
|
||||
postprocess(cfg, args.ne, args.n_proc, args.tmp, not args.nps)
|
||||
Reference in New Issue
Block a user