add rerender lora

animation/rerender-lora/src/config.py

@@ -0,0 +1,198 @@
+import json
+import os
+from typing import Optional, Sequence, Tuple
+
+from src.video_util import get_frame_count
+
+
+class RerenderConfig:
+
+    def __init__(self):
+        ...
+
+    def create_from_parameters(self,
+                               input_path: str,
+                               output_path: str,
+                               prompt: str,
+                               work_dir: Optional[str] = None,
+                               key_subdir: str = 'keys',
+                               frame_count: Optional[int] = None,
+                               interval: int = 10,
+                               crop: Sequence[int] = (0, 0, 0, 0),
+                               sd_model: Optional[str] = None,
+                               a_prompt: str = '',
+                               n_prompt: str = '',
+                               ddim_steps=20,
+                               scale=7.5,
+                               control_type: str = 'HED',
+                               control_strength=1,
+                               seed: int = -1,
+                               image_resolution: int = 512,
+                               use_limit_device_resolution: bool = False,
+                               x0_strength: float = -1,
+                               style_update_freq: int = 10,
+                               cross_period: Tuple[float, float] = (0, 1),
+                               warp_period: Tuple[float, float] = (0, 0.1),
+                               mask_period: Tuple[float, float] = (0.5, 0.8),
+                               ada_period: Tuple[float, float] = (1.0, 1.0),
+                               mask_strength: float = 0.5,
+                               inner_strength: float = 0.9,
+                               smooth_boundary: bool = True,
+                               color_preserve: bool = True,
+                               loose_cfattn: bool = False,
+                               freeu_args: Tuple[int] = (1, 1, 1, 1),
+                               lora_path: Optional[str] = None,
+                               **kwargs):
+        self.input_path = input_path
+        self.output_path = output_path
+        self.prompt = prompt
+        self.work_dir = work_dir
+        if work_dir is None:
+            self.work_dir = os.path.dirname(output_path)
+        self.key_dir = os.path.join(self.work_dir, key_subdir)
+        self.first_dir = os.path.join(self.work_dir, 'first')
+
+        # Split video into frames
+        if not os.path.isfile(input_path):
+            raise FileNotFoundError(f'Cannot find video file {input_path}')
+        self.input_dir = os.path.join(self.work_dir, 'video')
+
+        self.frame_count = frame_count
+        if frame_count is None:
+            self.frame_count = get_frame_count(self.input_path)
+        self.interval = interval
+        self.crop = crop
+        self.sd_model = sd_model
+        self.a_prompt = a_prompt
+        self.n_prompt = n_prompt
+        self.ddim_steps = ddim_steps
+        self.scale = scale
+        self.control_type = control_type
+        if self.control_type == 'canny':
+            self.canny_low = kwargs.get('canny_low', 100)
+            self.canny_high = kwargs.get('canny_high', 200)
+        else:
+            self.canny_low = None
+            self.canny_high = None
+        self.control_strength = control_strength
+        self.seed = seed
+        self.image_resolution = image_resolution
+        self.use_limit_device_resolution = use_limit_device_resolution
+        self.x0_strength = x0_strength
+        self.style_update_freq = style_update_freq
+        self.cross_period = cross_period
+        self.mask_period = mask_period
+        self.warp_period = warp_period
+        self.ada_period = ada_period
+        self.mask_strength = mask_strength
+        self.inner_strength = inner_strength
+        self.smooth_boundary = smooth_boundary
+        self.color_preserve = color_preserve
+        self.loose_cfattn = loose_cfattn
+        self.freeu_args = freeu_args
+        self.lora_path = lora_path
+
+        os.makedirs(self.input_dir, exist_ok=True)
+        os.makedirs(self.work_dir, exist_ok=True)
+        os.makedirs(self.key_dir, exist_ok=True)
+        os.makedirs(self.first_dir, exist_ok=True)
+
+    def create_from_path(self, cfg_path: str):
+        with open(cfg_path, 'r') as fp:
+            cfg = json.load(fp)
+        kwargs = dict()
+
+        def append_if_not_none(key):
+            value = cfg.get(key, None)
+            if value is not None:
+                kwargs[key] = value
+
+        kwargs['input_path'] = cfg['input']
+        kwargs['output_path'] = cfg['output']
+        kwargs['prompt'] = cfg['prompt']
+        append_if_not_none('work_dir')
+        append_if_not_none('key_subdir')
+        append_if_not_none('frame_count')
+        append_if_not_none('interval')
+        append_if_not_none('crop')
+        append_if_not_none('sd_model')
+        append_if_not_none('a_prompt')
+        append_if_not_none('n_prompt')
+        append_if_not_none('ddim_steps')
+        append_if_not_none('scale')
+        append_if_not_none('control_type')
+        if kwargs.get('control_type', '') == 'canny':
+            append_if_not_none('canny_low')
+            append_if_not_none('canny_high')
+        append_if_not_none('control_strength')
+        append_if_not_none('seed')
+        append_if_not_none('image_resolution')
+        append_if_not_none('use_limit_device_resolution')
+        append_if_not_none('x0_strength')
+        append_if_not_none('style_update_freq')
+        append_if_not_none('cross_period')
+        append_if_not_none('warp_period')
+        append_if_not_none('mask_period')
+        append_if_not_none('ada_period')
+        append_if_not_none('mask_strength')
+        append_if_not_none('inner_strength')
+        append_if_not_none('smooth_boundary')
+        append_if_not_none('color_perserve')
+        append_if_not_none('loose_cfattn')
+        append_if_not_none('freeu_args')
+        append_if_not_none('lora_path')
+        self.create_from_parameters(**kwargs)
+
+    def create_from_json(self, json_dict):
+        kwargs = dict()
+        def append_if_not_none(key):
+            value = json_dict.get(key, None)
+            if value is not None:
+                kwargs[key] = value
+
+        kwargs['input_path'] = json_dict['input']
+        kwargs['output_path'] = json_dict['output']
+        kwargs['prompt'] = json_dict['prompt']
+        append_if_not_none('work_dir')
+        append_if_not_none('key_subdir')
+        append_if_not_none('frame_count')
+        append_if_not_none('interval')
+        append_if_not_none('crop')
+        append_if_not_none('sd_model')
+        append_if_not_none('a_prompt')
+        append_if_not_none('n_prompt')
+        append_if_not_none('ddim_steps')
+        append_if_not_none('scale')
+        append_if_not_none('control_type')
+        if kwargs.get('control_type', '') == 'canny':
+            append_if_not_none('canny_low')
+            append_if_not_none('canny_high')
+        append_if_not_none('control_strength')
+        append_if_not_none('seed')
+        append_if_not_none('image_resolution')
+        append_if_not_none('use_limit_device_resolution')
+        append_if_not_none('x0_strength')
+        append_if_not_none('style_update_freq')
+        append_if_not_none('cross_period')
+        append_if_not_none('warp_period')
+        append_if_not_none('mask_period')
+        append_if_not_none('ada_period')
+        append_if_not_none('mask_strength')
+        append_if_not_none('inner_strength')
+        append_if_not_none('smooth_boundary')
+        append_if_not_none('color_perserve')
+        append_if_not_none('loose_cfattn')
+        append_if_not_none('freeu_args')
+        self.create_from_parameters(**kwargs)
+
+    @property
+    def use_warp(self):
+        return self.warp_period[0] <= self.warp_period[1]
+
+    @property
+    def use_mask(self):
+        return self.mask_period[0] <= self.mask_period[1]
+
+    @property
+    def use_ada(self):
+        return self.ada_period[0] <= self.ada_period[1]

animation/rerender-lora/src/controller.py

@@ -0,0 +1,143 @@
+import gc
+
+import torch
+import torch.nn.functional as F
+
+from flow.flow_utils import flow_warp
+
+# AdaIn
+
+
+def calc_mean_std(feat, eps=1e-5):
+    # eps is a small value added to the variance to avoid divide-by-zero.
+    size = feat.size()
+    assert (len(size) == 4)
+    N, C = size[:2]
+    feat_var = feat.view(N, C, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(N, C, 1, 1)
+    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
+    return feat_mean, feat_std
+
+
+class AttentionControl():
+
+    def __init__(self,
+                 inner_strength,
+                 mask_period,
+                 cross_period,
+                 ada_period,
+                 warp_period,
+                 loose_cfatnn=False):
+        self.step_store = self.get_empty_store()
+        self.cur_step = 0
+        self.total_step = 0
+        self.cur_index = 0
+        self.init_store = False
+        self.restore = False
+        self.update = False
+        self.flow = None
+        self.mask = None
+        self.restorex0 = False
+        self.updatex0 = False
+        self.inner_strength = inner_strength
+        self.cross_period = cross_period
+        self.mask_period = mask_period
+        self.ada_period = ada_period
+        self.warp_period = warp_period
+        self.up_resolution = 1280 if loose_cfatnn else 1281
+
+    @staticmethod
+    def get_empty_store():
+        return {
+            'first': [],
+            'previous': [],
+            'x0_previous': [],
+            'first_ada': []
+        }
+
+    def forward(self, context, is_cross: bool, place_in_unet: str):
+        cross_period = (self.total_step * self.cross_period[0],
+                        self.total_step * self.cross_period[1])
+        if not is_cross and place_in_unet == 'up' and context.shape[
+                2] < self.up_resolution:
+            if self.init_store:
+                self.step_store['first'].append(context.detach())
+                self.step_store['previous'].append(context.detach())
+            if self.update:
+                tmp = context.clone().detach()
+            if self.restore and self.cur_step >= cross_period[0] and \
+                    self.cur_step <= cross_period[1]:
+                context = torch.cat(
+                    (self.step_store['first'][self.cur_index],
+                     self.step_store['previous'][self.cur_index]),
+                    dim=1).clone()
+            if self.update:
+                self.step_store['previous'][self.cur_index] = tmp
+            self.cur_index += 1
+        return context
+
+    def update_x0(self, x0):
+        if self.init_store:
+            self.step_store['x0_previous'].append(x0.detach())
+            style_mean, style_std = calc_mean_std(x0.detach())
+            self.step_store['first_ada'].append(style_mean.detach())
+            self.step_store['first_ada'].append(style_std.detach())
+        if self.updatex0:
+            tmp = x0.clone().detach()
+        if self.restorex0:
+            if self.cur_step >= self.total_step * self.ada_period[
+                    0] and self.cur_step <= self.total_step * self.ada_period[
+                        1]:
+                x0 = F.instance_norm(x0) * self.step_store['first_ada'][
+                    2 * self.cur_step +
+                    1] + self.step_store['first_ada'][2 * self.cur_step]
+            if self.cur_step >= self.total_step * self.warp_period[
+                    0] and self.cur_step <= self.total_step * self.warp_period[
+                        1]:
+                pre = self.step_store['x0_previous'][self.cur_step]
+                x0 = flow_warp(pre, self.flow, mode='nearest') * self.mask + (
+                    1 - self.mask) * x0
+        if self.updatex0:
+            self.step_store['x0_previous'][self.cur_step] = tmp
+        return x0
+
+    def set_warp(self, flow, mask):
+        self.flow = flow.clone()
+        self.mask = mask.clone()
+
+    def __call__(self, context, is_cross: bool, place_in_unet: str):
+        context = self.forward(context, is_cross, place_in_unet)
+        return context
+
+    def set_step(self, step):
+        self.cur_step = step
+
+    def set_total_step(self, total_step):
+        self.total_step = total_step
+        self.cur_index = 0
+
+    def clear_store(self):
+        del self.step_store
+        torch.cuda.empty_cache()
+        gc.collect()
+        self.step_store = self.get_empty_store()
+
+    def set_task(self, task, restore_step=1.0):
+        self.init_store = False
+        self.restore = False
+        self.update = False
+        self.cur_index = 0
+        self.restore_step = restore_step
+        self.updatex0 = False
+        self.restorex0 = False
+        if 'initfirst' in task:
+            self.init_store = True
+            self.clear_store()
+        if 'updatestyle' in task:
+            self.update = True
+        if 'keepstyle' in task:
+            self.restore = True
+        if 'updatex0' in task:
+            self.updatex0 = True
+        if 'keepx0' in task:
+            self.restorex0 = True

animation/rerender-lora/src/ddim_v_hacked.py

@@ -0,0 +1,587 @@
+"""SAMPLING ONLY."""
+
+# CrossAttn precision handling
+import os
+
+import einops
+import numpy as np
+import torch
+from tqdm import tqdm
+
+from deps.ControlNet.ldm.modules.diffusionmodules.util import (
+    extract_into_tensor, make_ddim_sampling_parameters, make_ddim_timesteps,
+    noise_like)
+
+_ATTN_PRECISION = os.environ.get('ATTN_PRECISION', 'fp32')
+
+
+def register_attention_control(model, controller=None):
+
+    def ca_forward(self, place_in_unet):
+
+        def forward(x, context=None, mask=None):
+            h = self.heads
+
+            q = self.to_q(x)
+            is_cross = context is not None
+            context = context if is_cross else x
+            context = controller(context, is_cross, place_in_unet)
+
+            k = self.to_k(context)
+            v = self.to_v(context)
+
+            q, k, v = map(
+                lambda t: einops.rearrange(t, 'b n (h d) -> (b h) n d', h=h),
+                (q, k, v))
+
+            # force cast to fp32 to avoid overflowing
+            if _ATTN_PRECISION == 'fp32':
+                with torch.autocast(enabled=False, device_type='cuda'):
+                    q, k = q.float(), k.float()
+                    sim = torch.einsum('b i d, b j d -> b i j', q,
+                                       k) * self.scale
+            else:
+                sim = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale
+
+            del q, k
+
+            if mask is not None:
+                mask = einops.rearrange(mask, 'b ... -> b (...)')
+                max_neg_value = -torch.finfo(sim.dtype).max
+                mask = einops.repeat(mask, 'b j -> (b h) () j', h=h)
+                sim.masked_fill_(~mask, max_neg_value)
+
+            # attention, what we cannot get enough of
+            sim = sim.softmax(dim=-1)
+
+            out = torch.einsum('b i j, b j d -> b i d', sim, v)
+            out = einops.rearrange(out, '(b h) n d -> b n (h d)', h=h)
+            return self.to_out(out)
+
+        return forward
+
+    class DummyController:
+
+        def __call__(self, *args):
+            return args[0]
+
+        def __init__(self):
+            self.cur_step = 0
+
+    if controller is None:
+        controller = DummyController()
+
+    def register_recr(net_, place_in_unet):
+        if net_.__class__.__name__ == 'CrossAttention':
+            net_.forward = ca_forward(net_, place_in_unet)
+        elif hasattr(net_, 'children'):
+            for net__ in net_.children():
+                register_recr(net__, place_in_unet)
+
+    sub_nets = model.named_children()
+    for net in sub_nets:
+        if 'input_blocks' in net[0]:
+            register_recr(net[1], 'down')
+        elif 'output_blocks' in net[0]:
+            register_recr(net[1], 'up')
+        elif 'middle_block' in net[0]:
+            register_recr(net[1], 'mid')
+
+
+class DDIMVSampler(object):
+
+    def __init__(self, model, schedule='linear', **kwargs):
+        super().__init__()
+        self.model = model
+        self.ddpm_num_timesteps = model.num_timesteps
+        self.schedule = schedule
+
+    def register_buffer(self, name, attr):
+        if type(attr) == torch.Tensor:
+            if attr.device != torch.device('cuda'):
+                attr = attr.to(torch.device('cuda'))
+        setattr(self, name, attr)
+
+    def make_schedule(self,
+                      ddim_num_steps,
+                      ddim_discretize='uniform',
+                      ddim_eta=0.,
+                      verbose=True):
+        self.ddim_timesteps = make_ddim_timesteps(
+            ddim_discr_method=ddim_discretize,
+            num_ddim_timesteps=ddim_num_steps,
+            num_ddpm_timesteps=self.ddpm_num_timesteps,
+            verbose=verbose)
+        alphas_cumprod = self.model.alphas_cumprod
+        assert alphas_cumprod.shape[0] == self.ddpm_num_timesteps, \
+            'alphas have to be defined for each timestep'
+
+        def to_torch(x):
+            return x.clone().detach().to(torch.float32).to(self.model.device)
+
+        self.register_buffer('betas', to_torch(self.model.betas))
+        self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
+        self.register_buffer('alphas_cumprod_prev',
+                             to_torch(self.model.alphas_cumprod_prev))
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.register_buffer('sqrt_alphas_cumprod',
+                             to_torch(np.sqrt(alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod',
+                             to_torch(np.sqrt(1. - alphas_cumprod.cpu())))
+        self.register_buffer('log_one_minus_alphas_cumprod',
+                             to_torch(np.log(1. - alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recip_alphas_cumprod',
+                             to_torch(np.sqrt(1. / alphas_cumprod.cpu())))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod',
+                             to_torch(np.sqrt(1. / alphas_cumprod.cpu() - 1)))
+
+        # ddim sampling parameters
+        ddim_sigmas, ddim_alphas, ddim_alphas_prev = \
+            make_ddim_sampling_parameters(
+                alphacums=alphas_cumprod.cpu(),
+                ddim_timesteps=self.ddim_timesteps,
+                eta=ddim_eta,
+                verbose=verbose)
+        self.register_buffer('ddim_sigmas', ddim_sigmas)
+        self.register_buffer('ddim_alphas', ddim_alphas)
+        self.register_buffer('ddim_alphas_prev', ddim_alphas_prev)
+        self.register_buffer('ddim_sqrt_one_minus_alphas',
+                             np.sqrt(1. - ddim_alphas))
+        sigmas_for_original_sampling_steps = ddim_eta * torch.sqrt(
+            (1 - self.alphas_cumprod_prev) / (1 - self.alphas_cumprod) *
+            (1 - self.alphas_cumprod / self.alphas_cumprod_prev))
+        self.register_buffer('ddim_sigmas_for_original_num_steps',
+                             sigmas_for_original_sampling_steps)
+
+    @torch.no_grad()
+    def sample(self,
+               S,
+               batch_size,
+               shape,
+               conditioning=None,
+               callback=None,
+               img_callback=None,
+               quantize_x0=False,
+               eta=0.,
+               mask=None,
+               x0=None,
+               xtrg=None,
+               noise_rescale=None,
+               temperature=1.,
+               noise_dropout=0.,
+               score_corrector=None,
+               corrector_kwargs=None,
+               verbose=True,
+               x_T=None,
+               log_every_t=100,
+               unconditional_guidance_scale=1.,
+               unconditional_conditioning=None,
+               dynamic_threshold=None,
+               ucg_schedule=None,
+               controller=None,
+               strength=0.0,
+               **kwargs):
+        if conditioning is not None:
+            if isinstance(conditioning, dict):
+                ctmp = conditioning[list(conditioning.keys())[0]]
+                while isinstance(ctmp, list):
+                    ctmp = ctmp[0]
+                cbs = ctmp.shape[0]
+                if cbs != batch_size:
+                    print(f'Warning: Got {cbs} conditionings'
+                          f'but batch-size is {batch_size}')
+
+            elif isinstance(conditioning, list):
+                for ctmp in conditioning:
+                    if ctmp.shape[0] != batch_size:
+                        print(f'Warning: Got {cbs} conditionings'
+                              f'but batch-size is {batch_size}')
+
+            else:
+                if conditioning.shape[0] != batch_size:
+                    print(f'Warning: Got {conditioning.shape[0]}'
+                          f'conditionings but batch-size is {batch_size}')
+
+        self.make_schedule(ddim_num_steps=S, ddim_eta=eta, verbose=verbose)
+        # sampling
+        C, H, W = shape
+        size = (batch_size, C, H, W)
+        print(f'Data shape for DDIM sampling is {size}, eta {eta}')
+
+        samples, intermediates = self.ddim_sampling(
+            conditioning,
+            size,
+            callback=callback,
+            img_callback=img_callback,
+            quantize_denoised=quantize_x0,
+            mask=mask,
+            x0=x0,
+            xtrg=xtrg,
+            noise_rescale=noise_rescale,
+            ddim_use_original_steps=False,
+            noise_dropout=noise_dropout,
+            temperature=temperature,
+            score_corrector=score_corrector,
+            corrector_kwargs=corrector_kwargs,
+            x_T=x_T,
+            log_every_t=log_every_t,
+            unconditional_guidance_scale=unconditional_guidance_scale,
+            unconditional_conditioning=unconditional_conditioning,
+            dynamic_threshold=dynamic_threshold,
+            ucg_schedule=ucg_schedule,
+            controller=controller,
+            strength=strength,
+        )
+        return samples, intermediates
+
+    @torch.no_grad()
+    def ddim_sampling(self,
+                      cond,
+                      shape,
+                      x_T=None,
+                      ddim_use_original_steps=False,
+                      callback=None,
+                      timesteps=None,
+                      quantize_denoised=False,
+                      mask=None,
+                      x0=None,
+                      xtrg=None,
+                      noise_rescale=None,
+                      img_callback=None,
+                      log_every_t=100,
+                      temperature=1.,
+                      noise_dropout=0.,
+                      score_corrector=None,
+                      corrector_kwargs=None,
+                      unconditional_guidance_scale=1.,
+                      unconditional_conditioning=None,
+                      dynamic_threshold=None,
+                      ucg_schedule=None,
+                      controller=None,
+                      strength=0.0):
+
+        if strength == 1 and x0 is not None:
+            return x0, None
+
+        register_attention_control(self.model.model.diffusion_model,
+                                   controller)
+
+        device = self.model.betas.device
+        b = shape[0]
+        if x_T is None:
+            img = torch.randn(shape, device=device)
+        else:
+            img = x_T
+
+        if timesteps is None:
+            timesteps = self.ddpm_num_timesteps if ddim_use_original_steps \
+                else self.ddim_timesteps
+        elif timesteps is not None and not ddim_use_original_steps:
+            subset_end = int(
+                min(timesteps / self.ddim_timesteps.shape[0], 1) *
+                self.ddim_timesteps.shape[0]) - 1
+            timesteps = self.ddim_timesteps[:subset_end]
+
+        intermediates = {'x_inter': [img], 'pred_x0': [img]}
+        time_range = reversed(range(
+            0, timesteps)) if ddim_use_original_steps else np.flip(timesteps)
+        total_steps = timesteps if ddim_use_original_steps \
+            else timesteps.shape[0]
+        print(f'Running DDIM Sampling with {total_steps} timesteps')
+
+        iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps)
+        if controller is not None:
+            controller.set_total_step(total_steps)
+        if mask is None:
+            mask = [None] * total_steps
+
+        dir_xt = 0
+        for i, step in enumerate(iterator):
+            if controller is not None:
+                controller.set_step(i)
+            index = total_steps - i - 1
+            ts = torch.full((b, ), step, device=device, dtype=torch.long)
+
+            if strength >= 0 and i == int(
+                    total_steps * strength) and x0 is not None:
+                img = self.model.q_sample(x0, ts)
+            if mask is not None and xtrg is not None:
+                # TODO: deterministic forward pass?
+                if type(mask) == list:
+                    weight = mask[i]
+                else:
+                    weight = mask
+                if weight is not None:
+                    rescale = torch.maximum(1. - weight, (1 - weight**2)**0.5 *
+                                            controller.inner_strength)
+                    if noise_rescale is not None:
+                        rescale = (1. - weight) * (
+                            1 - noise_rescale) + rescale * noise_rescale
+                    img_ref = self.model.q_sample(xtrg, ts)
+                    img = img_ref * weight + (1. - weight) * (
+                        img - dir_xt) + rescale * dir_xt
+
+            if ucg_schedule is not None:
+                assert len(ucg_schedule) == len(time_range)
+                unconditional_guidance_scale = ucg_schedule[i]
+
+            outs = self.p_sample_ddim(
+                img,
+                cond,
+                ts,
+                index=index,
+                use_original_steps=ddim_use_original_steps,
+                quantize_denoised=quantize_denoised,
+                temperature=temperature,
+                noise_dropout=noise_dropout,
+                score_corrector=score_corrector,
+                corrector_kwargs=corrector_kwargs,
+                unconditional_guidance_scale=unconditional_guidance_scale,
+                unconditional_conditioning=unconditional_conditioning,
+                dynamic_threshold=dynamic_threshold,
+                controller=controller,
+                return_dir=True)
+            img, pred_x0, dir_xt = outs
+            if callback:
+                callback(i)
+            if img_callback:
+                img_callback(pred_x0, i)
+
+            if index % log_every_t == 0 or index == total_steps - 1:
+                intermediates['x_inter'].append(img)
+                intermediates['pred_x0'].append(pred_x0)
+
+        return img, intermediates
+
+    @torch.no_grad()
+    def p_sample_ddim(self,
+                      x,
+                      c,
+                      t,
+                      index,
+                      repeat_noise=False,
+                      use_original_steps=False,
+                      quantize_denoised=False,
+                      temperature=1.,
+                      noise_dropout=0.,
+                      score_corrector=None,
+                      corrector_kwargs=None,
+                      unconditional_guidance_scale=1.,
+                      unconditional_conditioning=None,
+                      dynamic_threshold=None,
+                      controller=None,
+                      return_dir=False):
+        b, *_, device = *x.shape, x.device
+
+        if unconditional_conditioning is None or \
+                unconditional_guidance_scale == 1.:
+            model_output = self.model.apply_model(x, t, c)
+        else:
+            model_t = self.model.apply_model(x, t, c)
+            model_uncond = self.model.apply_model(x, t,
+                                                  unconditional_conditioning)
+            model_output = model_uncond + unconditional_guidance_scale * (
+                model_t - model_uncond)
+
+        if self.model.parameterization == 'v':
+            e_t = self.model.predict_eps_from_z_and_v(x, t, model_output)
+        else:
+            e_t = model_output
+
+        if score_corrector is not None:
+            assert self.model.parameterization == 'eps', 'not implemented'
+            e_t = score_corrector.modify_score(self.model, e_t, x, t, c,
+                                               **corrector_kwargs)
+
+        if use_original_steps:
+            alphas = self.model.alphas_cumprod
+            alphas_prev = self.model.alphas_cumprod_prev
+            sqrt_one_minus_alphas = self.model.sqrt_one_minus_alphas_cumprod
+            sigmas = self.model.ddim_sigmas_for_original_num_steps
+        else:
+            alphas = self.ddim_alphas
+            alphas_prev = self.ddim_alphas_prev
+            sqrt_one_minus_alphas = self.ddim_sqrt_one_minus_alphas
+            sigmas = self.ddim_sigmas
+
+        # select parameters corresponding to the currently considered timestep
+        a_t = torch.full((b, 1, 1, 1), alphas[index], device=device)
+        a_prev = torch.full((b, 1, 1, 1), alphas_prev[index], device=device)
+        sigma_t = torch.full((b, 1, 1, 1), sigmas[index], device=device)
+        sqrt_one_minus_at = torch.full((b, 1, 1, 1),
+                                       sqrt_one_minus_alphas[index],
+                                       device=device)
+
+        # current prediction for x_0
+        if self.model.parameterization != 'v':
+            pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt()
+        else:
+            pred_x0 = self.model.predict_start_from_z_and_v(x, t, model_output)
+
+        if quantize_denoised:
+            pred_x0, _, *_ = self.model.first_stage_model.quantize(pred_x0)
+
+        if dynamic_threshold is not None:
+            raise NotImplementedError()
+        '''
+        if mask is not None and xtrg is not None:
+            pred_x0 = xtrg * mask + (1. - mask) * pred_x0
+        '''
+
+        if controller is not None:
+            pred_x0 = controller.update_x0(pred_x0)
+
+        # direction pointing to x_t
+        dir_xt = (1. - a_prev - sigma_t**2).sqrt() * e_t
+        noise = sigma_t * noise_like(x.shape, device,
+                                     repeat_noise) * temperature
+        if noise_dropout > 0.:
+            noise = torch.nn.functional.dropout(noise, p=noise_dropout)
+        x_prev = a_prev.sqrt() * pred_x0 + dir_xt + noise
+
+        if return_dir:
+            return x_prev, pred_x0, dir_xt
+        return x_prev, pred_x0
+
+    @torch.no_grad()
+    def encode(self,
+               x0,
+               c,
+               t_enc,
+               use_original_steps=False,
+               return_intermediates=None,
+               unconditional_guidance_scale=1.0,
+               unconditional_conditioning=None,
+               callback=None):
+        timesteps = np.arange(self.ddpm_num_timesteps
+                              ) if use_original_steps else self.ddim_timesteps
+        num_reference_steps = timesteps.shape[0]
+
+        assert t_enc <= num_reference_steps
+        num_steps = t_enc
+
+        if use_original_steps:
+            alphas_next = self.alphas_cumprod[:num_steps]
+            alphas = self.alphas_cumprod_prev[:num_steps]
+        else:
+            alphas_next = self.ddim_alphas[:num_steps]
+            alphas = torch.tensor(self.ddim_alphas_prev[:num_steps])
+
+        x_next = x0
+        intermediates = []
+        inter_steps = []
+        for i in tqdm(range(num_steps), desc='Encoding Image'):
+            t = torch.full((x0.shape[0], ),
+                           timesteps[i],
+                           device=self.model.device,
+                           dtype=torch.long)
+            if unconditional_guidance_scale == 1.:
+                noise_pred = self.model.apply_model(x_next, t, c)
+            else:
+                assert unconditional_conditioning is not None
+                e_t_uncond, noise_pred = torch.chunk(
+                    self.model.apply_model(
+                        torch.cat((x_next, x_next)), torch.cat((t, t)),
+                        torch.cat((unconditional_conditioning, c))), 2)
+                noise_pred = e_t_uncond + unconditional_guidance_scale * (
+                    noise_pred - e_t_uncond)
+            xt_weighted = (alphas_next[i] / alphas[i]).sqrt() * x_next
+            weighted_noise_pred = alphas_next[i].sqrt() * (
+                (1 / alphas_next[i] - 1).sqrt() -
+                (1 / alphas[i] - 1).sqrt()) * noise_pred
+            x_next = xt_weighted + weighted_noise_pred
+            if return_intermediates and i % (num_steps // return_intermediates
+                                             ) == 0 and i < num_steps - 1:
+                intermediates.append(x_next)
+                inter_steps.append(i)
+            elif return_intermediates and i >= num_steps - 2:
+                intermediates.append(x_next)
+                inter_steps.append(i)
+            if callback:
+                callback(i)
+
+        out = {'x_encoded': x_next, 'intermediate_steps': inter_steps}
+        if return_intermediates:
+            out.update({'intermediates': intermediates})
+        return x_next, out
+
+    @torch.no_grad()
+    def stochastic_encode(self, x0, t, use_original_steps=False, noise=None):
+        # fast, but does not allow for exact reconstruction
+        # t serves as an index to gather the correct alphas
+        if use_original_steps:
+            sqrt_alphas_cumprod = self.sqrt_alphas_cumprod
+            sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod
+        else:
+            sqrt_alphas_cumprod = torch.sqrt(self.ddim_alphas)
+            sqrt_one_minus_alphas_cumprod = self.ddim_sqrt_one_minus_alphas
+
+        if noise is None:
+            noise = torch.randn_like(x0)
+        if t >= len(sqrt_alphas_cumprod):
+            return noise
+        return (
+            extract_into_tensor(sqrt_alphas_cumprod, t, x0.shape) * x0 +
+            extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x0.shape) *
+            noise)
+
+    @torch.no_grad()
+    def decode(self,
+               x_latent,
+               cond,
+               t_start,
+               unconditional_guidance_scale=1.0,
+               unconditional_conditioning=None,
+               use_original_steps=False,
+               callback=None):
+
+        timesteps = np.arange(self.ddpm_num_timesteps
+                              ) if use_original_steps else self.ddim_timesteps
+        timesteps = timesteps[:t_start]
+
+        time_range = np.flip(timesteps)
+        total_steps = timesteps.shape[0]
+        print(f'Running DDIM Sampling with {total_steps} timesteps')
+
+        iterator = tqdm(time_range, desc='Decoding image', total=total_steps)
+        x_dec = x_latent
+        for i, step in enumerate(iterator):
+            index = total_steps - i - 1
+            ts = torch.full((x_latent.shape[0], ),
+                            step,
+                            device=x_latent.device,
+                            dtype=torch.long)
+            x_dec, _ = self.p_sample_ddim(
+                x_dec,
+                cond,
+                ts,
+                index=index,
+                use_original_steps=use_original_steps,
+                unconditional_guidance_scale=unconditional_guidance_scale,
+                unconditional_conditioning=unconditional_conditioning)
+            if callback:
+                callback(i)
+        return x_dec
+
+
+def calc_mean_std(feat, eps=1e-5):
+    # eps is a small value added to the variance to avoid divide-by-zero.
+    size = feat.size()
+    assert (len(size) == 4)
+    N, C = size[:2]
+    feat_var = feat.view(N, C, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(N, C, 1, 1)
+    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
+    return feat_mean, feat_std
+
+
+def adaptive_instance_normalization(content_feat, style_feat):
+    assert (content_feat.size()[:2] == style_feat.size()[:2])
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+
+    normalized_feat = (content_feat -
+                       content_mean.expand(size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)

animation/rerender-lora/src/freeu.py

@@ -0,0 +1,94 @@
+import torch
+import torch.fft as fft
+
+
+def Fourier_filter(x, threshold, scale):
+
+    x_freq = fft.fftn(x, dim=(-2, -1))
+    x_freq = fft.fftshift(x_freq, dim=(-2, -1))
+
+    B, C, H, W = x_freq.shape
+    mask = torch.ones((B, C, H, W)).cuda()
+
+    crow, ccol = H // 2, W // 2
+    mask[..., crow - threshold:crow + threshold,
+         ccol - threshold:ccol + threshold] = scale
+    x_freq = x_freq * mask
+
+    x_freq = fft.ifftshift(x_freq, dim=(-2, -1))
+
+    x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real
+
+    return x_filtered
+
+from deps.ControlNet.ldm.modules.diffusionmodules.util import \
+    timestep_embedding  # noqa：E501
+
+
+# backbone_scale1=1.1, backbone_scale2=1.2, skip_scale1=1.0, skip_scale2=0.2
+def freeu_forward(self,
+                  backbone_scale1=1.,
+                  backbone_scale2=1.,
+                  skip_scale1=1.,
+                  skip_scale2=1.):
+
+    def forward(x,
+                timesteps=None,
+                context=None,
+                control=None,
+                only_mid_control=False,
+                **kwargs):
+        hs = []
+        with torch.no_grad():
+            t_emb = timestep_embedding(timesteps,
+                                       self.model_channels,
+                                       repeat_only=False)
+            emb = self.time_embed(t_emb)
+            h = x.type(self.dtype)
+            for module in self.input_blocks:
+                h = module(h, emb, context)
+                hs.append(h)
+            h = self.middle_block(h, emb, context)
+
+        if control is not None:
+            h += control.pop()
+        '''
+        for i, module in enumerate(self.output_blocks):
+            if only_mid_control or control is None:
+                h = torch.cat([h, hs.pop()], dim=1)
+            else:
+                h = torch.cat([h, hs.pop() + control.pop()], dim=1)
+            h = module(h, emb, context)
+        '''
+        for i, module in enumerate(self.output_blocks):
+            hs_ = hs.pop()
+
+            if h.shape[1] == 1280:
+                hidden_mean = h.mean(1).unsqueeze(1)
+                B = hidden_mean.shape[0]
+                hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True) 
+                hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True)
+                hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3)
+                h[:, :640] = h[:, :640] * ((backbone_scale1 - 1) * hidden_mean + 1)        
+                # h[:, :640] = h[:, :640] * backbone_scale1
+                hs_ = Fourier_filter(hs_, threshold=1, scale=skip_scale1)
+            if h.shape[1] == 640:
+                hidden_mean = h.mean(1).unsqueeze(1)
+                B = hidden_mean.shape[0]
+                hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True) 
+                hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True)
+                hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3) 
+                h[:, :320] = h[:, :320] * ((backbone_scale2 - 1) * hidden_mean + 1)
+                # h[:, :320] = h[:, :320] * backbone_scale2
+                hs_ = Fourier_filter(hs_, threshold=1, scale=skip_scale2)
+
+            if only_mid_control or control is None:
+                h = torch.cat([h, hs_], dim=1)
+            else:
+                h = torch.cat([h, hs_ + control.pop()], dim=1)
+            h = module(h, emb, context)
+
+        h = h.type(x.dtype)
+        return self.out(h)
+
+    return forward

animation/rerender-lora/src/img_util.py

@@ -0,0 +1,23 @@
+import einops
+import torch
+import torch.nn.functional as F
+
+
+@torch.no_grad()
+def find_flat_region(mask):
+    device = mask.device
+    kernel_x = torch.Tensor([[-1, 0, 1], [-1, 0, 1],
+                             [-1, 0, 1]]).unsqueeze(0).unsqueeze(0).to(device)
+    kernel_y = torch.Tensor([[-1, -1, -1], [0, 0, 0],
+                             [1, 1, 1]]).unsqueeze(0).unsqueeze(0).to(device)
+    mask_ = F.pad(mask.unsqueeze(0), (1, 1, 1, 1), mode='replicate')
+
+    grad_x = torch.nn.functional.conv2d(mask_, kernel_x)
+    grad_y = torch.nn.functional.conv2d(mask_, kernel_y)
+    return ((abs(grad_x) + abs(grad_y)) == 0).float()[0]
+
+
+def numpy2tensor(img):
+    x0 = torch.from_numpy(img.copy()).float().cuda() / 255.0 * 2.0 - 1.
+    x0 = torch.stack([x0], dim=0)
+    return einops.rearrange(x0, 'b h w c -> b c h w').clone()

animation/rerender-lora/src/import_util.py

@@ -0,0 +1,10 @@
+import os
+import sys
+
+cur_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+gmflow_dir = os.path.join(cur_dir, 'deps/gmflow')
+controlnet_dir = os.path.join(cur_dir, 'deps/ControlNet')
+sys.path.insert(0, gmflow_dir)
+sys.path.insert(0, controlnet_dir)
+
+import deps.ControlNet.share  # noqa: F401 E402

animation/rerender-lora/src/lora.py

@@ -0,0 +1,182 @@
+import json
+import pathlib
+
+import torch
+from torch import nn
+from safetensors import safe_open
+
+unet_mapping = json.load(open('unet_lora_mapping.json'))
+
+
+class LoraLinear(torch.nn.Module):
+    def __init__(self, old_linear, up=None, down=None):
+        super().__init__()
+        self.old_linear = old_linear
+        self.up = up
+        self.down = down
+        self.alpha = 1.0
+        self.new_linear = None
+
+    def forward(self, x):
+        if self.new_linear is None:
+            self.new_linear = nn.Linear(self.down.shape[1], self.up.shape[0])
+            self.new_linear.weight.data = self.up @ self.down
+            self.new_linear = self.new_linear.to('cuda').to(torch.float32)
+        x1 = self.new_linear(x) * (self.alpha / float(self.up[1].size()[0])) if self.alpha is not None else 1
+        x2 = self.old_linear(x)
+        return x1 + x2
+
+
+class LoraConv2d(torch.nn.Module):
+    def __init__(self, old_conv2d, up=None, down=None):
+        super().__init__()
+        self.old_conv2d = old_conv2d
+        self.up = up
+        self.down = down
+        self.alpha = None
+        self.new_conv2d = None
+
+    def forward(self, x):
+        if self.new_conv2d is None:
+            self.new_conv2d = nn.Conv2d(self.down.shape[1], self.up.shape[0], 1)
+            self.new_conv2d.weight.data = (self.up.squeeze() @ self.down.squeeze()).unsqueeze(-1).unsqueeze(-1)
+            self.new_conv2d = self.new_conv2d.to('cuda').to(torch.float32)
+        x1 = self.new_conv2d(x) * (self.alpha / float(self.up[1].size()[0])) if self.alpha is not None else 1
+        x2 = self.old_conv2d(x)
+        return x1 + x2
+
+
+def _get_nested_attr(obj, attr_path):
+    elements = attr_path.split('.')
+    for elem in elements:
+        if elem.isdigit():  # If the element is a digit, access it as an index
+            obj = obj[int(elem)]
+        else:  # Otherwise, access it as an attribute
+            obj = getattr(obj, elem)
+    return obj
+
+
+def _set_nested_attr(obj, attr_path, value):
+    elements = attr_path.split('.')
+    for elem in elements[:-1]:
+        if elem.isdigit():
+            obj = obj[int(elem)]
+        else:
+            obj = getattr(obj, elem)
+    setattr(obj, elements[-1], value)
+
+
+def _replace_underscores(input_string, exceptions):
+    # Split the input string by underscores
+    parts = input_string.split('_')
+
+    # Reconstruct the string, replacing underscores with dots
+    # except in the specified exceptions
+    output_parts = []
+    skip_next = False
+    for i, part in enumerate(parts):
+        if skip_next:
+            skip_next = False
+            continue
+
+        # Check if this part combined with the next part is in exceptions
+        if i < len(parts) - 1 and '_'.join([part, parts[i + 1]]) in exceptions:
+            output_parts.append('_'.join([part, parts[i + 1]]))
+            skip_next = True
+        else:
+            output_parts.append(part)
+
+    return '.'.join(output_parts)
+
+
+def apply_lora(unet, text_encoder, lora_weights):
+    for param_name, param_value in lora_weights.items():
+        if 'unet' in param_name:
+            # Determine layer type
+            if 'lora_down' in param_name:
+                layer_type = 'down'
+            elif 'lora_up' in param_name:
+                layer_type = 'up'
+            elif 'alpha' in param_name:
+                layer_type = 'alpha'
+            else:
+                raise ValueError(f'Unknown layer type for {param_name}')
+
+            base_param_name = param_name.split('.')[0]
+            replaced_name = _replace_underscores(base_param_name, exceptions=[
+                'up_blocks', 'mid_block', 'down_blocks', 'transformer_blocks',
+                'to_k', 'to_q', 'to_v', 'to_out', 'to_in', 'proj_in', 'proj_out',
+            ])
+
+            for diffusers_name, sd_name in unet_mapping.items():
+                base_diffusers = diffusers_name.rsplit('.', 1)[0]
+                base_sd = sd_name.rsplit('.', 1)[0]
+                if base_diffusers in replaced_name:
+                    replaced_name = replaced_name.replace(base_diffusers, base_sd)
+
+            replaced_name = replaced_name.replace('lora.unet', 'diffusion_model')
+            param_value = param_value.to('cuda').to(torch.float32)
+            curModule = _get_nested_attr(unet, replaced_name)
+
+            # Update module type if needed
+            if isinstance(curModule, torch.nn.Linear):
+                _set_nested_attr(unet, replaced_name, LoraLinear(curModule))
+            elif isinstance(curModule, torch.nn.Conv2d):
+                _set_nested_attr(unet, replaced_name, LoraConv2d(curModule))
+            elif not isinstance(curModule, (LoraLinear, LoraConv2d)):
+                raise ValueError(f'Unknown layer type: {type(curModule)}')
+
+            if layer_type == 'down':
+                _set_nested_attr(unet, replaced_name + '.down', param_value)
+            elif layer_type == 'up':
+                _set_nested_attr(unet, replaced_name + '.up', param_value)
+            elif layer_type == 'alpha':
+                _set_nested_attr(unet, replaced_name + '.alpha', param_value)
+        elif 'te_text_model' in param_name:
+            # Determine layer type
+            if 'lora_down' in param_name:
+                layer_type = 'down'
+            elif 'lora_up' in param_name:
+                layer_type = 'up'
+            elif 'alpha' in param_name:
+                layer_type = 'alpha'
+            else:
+                raise ValueError(f'Unknown layer type for {param_name}')
+
+            base_param_name = param_name.split('.')[0]
+            replaced_name = base_param_name.replace('lora_te', 'transformer')
+            replaced_name = _replace_underscores(replaced_name, exceptions=[
+                'text_model', 'self_attn', 'k_proj', 'v_proj', 'q_proj', 'out_proj',
+            ])
+            curModule = _get_nested_attr(text_encoder, replaced_name)
+            if isinstance(curModule, torch.nn.Linear):
+                _set_nested_attr(text_encoder, replaced_name, LoraLinear(curModule))
+            elif not isinstance(curModule, LoraLinear):
+                raise ValueError(f'Unknown layer type: {type(curModule)}')
+            param_value = param_value.to('cuda').to(torch.float32)
+            if layer_type == 'down':
+                _set_nested_attr(text_encoder, replaced_name + '.down', param_value)
+            elif layer_type == 'up':
+                _set_nested_attr(text_encoder, replaced_name + '.up', param_value)
+            elif layer_type == 'alpha':
+                _set_nested_attr(text_encoder, replaced_name + '.alpha', param_value)
+
+
+def load_lora(filepath):
+    path = pathlib.Path(filepath)
+    if path.suffix == '.safetensors':
+        try:
+            print(f'Loading LORA weights using safetensors from {filepath}')
+            lora_weights = {}
+            with safe_open(filepath, framework="pt", device='cuda') as f:
+                for key in f.keys():
+                    lora_weights[key] = f.get_tensor(key)
+        except Exception as e:
+            raise Exception(f'Error loading LORA weights: {str(e)}')
+    else:
+        try:
+            print(f'Loading LORA weights using torch from {filepath}')
+            lora_weights = torch.load(filepath, map_location='cuda')
+        except Exception as e:
+            raise Exception(f'Error loading LORA weights from {filepath}: {str(e)}')
+    return lora_weights

animation/rerender-lora/src/video_util.py

@@ -0,0 +1,127 @@
+import os
+
+import cv2
+import torch
+import imageio
+import numpy as np
+
+
+def video_to_frame(video_path: str,
+                   frame_dir: str,
+                   filename_pattern: str = 'frame%03d.jpg',
+                   log: bool = True,
+                   frame_edit_func=None):
+    os.makedirs(frame_dir, exist_ok=True)
+
+    vidcap = cv2.VideoCapture(video_path)
+    success, image = vidcap.read()
+
+    if log:
+        print('img shape: ', image.shape[0:2])
+
+    count = 0
+    while success:
+        if frame_edit_func is not None:
+            image = frame_edit_func(image)
+
+        cv2.imwrite(os.path.join(frame_dir, filename_pattern % count), image)
+        success, image = vidcap.read()
+        if log:
+            print('Read a new frame: ', success, count)
+        count += 1
+
+    vidcap.release()
+
+
+def frame_to_video(video_path: str, frame_dir: str, fps=30, log=True):
+
+    first_img = True
+    writer = imageio.get_writer(video_path, fps=fps)
+
+    file_list = sorted(os.listdir(frame_dir))
+    for file_name in file_list:
+        if not (file_name.endswith('jpg') or file_name.endswith('png')):
+            continue
+
+        fn = os.path.join(frame_dir, file_name)
+        curImg = imageio.imread(fn)
+
+        if first_img:
+            H, W = curImg.shape[0:2]
+            if log:
+                print('img shape', (H, W))
+            first_img = False
+
+        writer.append_data(curImg)
+
+    writer.close()
+
+
+def get_fps(video_path: str):
+    video = cv2.VideoCapture(video_path)
+    fps = video.get(cv2.CAP_PROP_FPS)
+    video.release()
+    return fps
+
+
+def get_frame_count(video_path: str):
+    video = cv2.VideoCapture(video_path)
+    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
+    video.release()
+    return frame_count
+
+
+def resize_image(input_image, resolution):
+    H, W, C = input_image.shape
+    H = float(H)
+    W = float(W)
+    aspect_ratio = W / H
+    k = float(resolution) / min(H, W)
+    H *= k
+    W *= k
+    if H < W:
+        W = resolution
+        H = int(resolution / aspect_ratio)
+    else:
+        H = resolution
+        W = int(aspect_ratio * resolution)
+    H = int(np.round(H / 64.0)) * 64
+    W = int(np.round(W / 64.0)) * 64
+    img = cv2.resize(
+        input_image, (W, H),
+        interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
+    return img
+
+
+def prepare_frames(input_path: str, output_dir: str, resolution: int, crop, use_limit_device_resolution=False):
+    l, r, t, b = crop
+
+    if use_limit_device_resolution:
+        resolution = vram_limit_device_resolution(resolution)
+
+    def crop_func(frame):
+        H, W, C = frame.shape
+        left = np.clip(l, 0, W)
+        right = np.clip(W - r, left, W)
+        top = np.clip(t, 0, H)
+        bottom = np.clip(H - b, top, H)
+        frame = frame[top:bottom, left:right]
+        return resize_image(frame, resolution)
+
+    video_to_frame(input_path, output_dir, '%04d.png', False, crop_func)
+
+
+def vram_limit_device_resolution(resolution, device="cuda"):
+    # get max limit target size
+    gpu_vram = torch.cuda.get_device_properties(device).total_memory / (1024 ** 3)
+    # table of gpu memory limit
+    gpu_table = {24: 1280, 18: 1024, 14: 768, 10: 640, 8: 576, 7: 512, 6: 448, 5: 320, 4: 192, 0: 0}
+    # get user resize for gpu
+    device_resolution = max(val for key, val in gpu_table.items() if key <= gpu_vram)
+    print(f"Limit VRAM is {gpu_vram} Gb and size {device_resolution}.")
+    if gpu_vram < 4:
+        print(f"Small VRAM to use GPU. Configuration resolution will be used.")
+    if resolution < device_resolution:
+        print(f"Video will not resize")
+        return resolution
+    return device_resolution