c2dbfba56f79c968b0d59e1af1493bfa9e6cfbd19338ec4a352438df2c986640

SD-CN-Animation/scripts/core/flow_utils.py

@@ -0,0 +1,156 @@
+import sys, os
+
+import numpy as np
+import cv2
+
+from collections import namedtuple
+import torch
+import argparse
+from RAFT.raft import RAFT
+from RAFT.utils.utils import InputPadder
+
+import modules.paths as ph
+import gc
+
+RAFT_model = None
+fgbg = cv2.createBackgroundSubtractorMOG2(history=500, varThreshold=16, detectShadows=True)
+
+def background_subtractor(frame, fgbg):
+  fgmask = fgbg.apply(frame)
+  return cv2.bitwise_and(frame, frame, mask=fgmask)
+
+def RAFT_clear_memory():
+  global RAFT_model
+  del RAFT_model
+  gc.collect()
+  torch.cuda.empty_cache()
+  RAFT_model = None
+
+def RAFT_estimate_flow(frame1, frame2, device='cuda'):
+  global RAFT_model
+
+  org_size = frame1.shape[1], frame1.shape[0]
+  size = frame1.shape[1] // 16 * 16, frame1.shape[0] // 16 * 16
+  frame1 = cv2.resize(frame1, size)
+  frame2 = cv2.resize(frame2, size)
+
+  model_path = ph.models_path + '/RAFT/raft-things.pth'
+  remote_model_path = 'https://drive.google.com/uc?id=1MqDajR89k-xLV0HIrmJ0k-n8ZpG6_suM'
+
+  if not os.path.isfile(model_path):
+    from basicsr.utils.download_util import load_file_from_url
+    os.makedirs(os.path.dirname(model_path), exist_ok=True)
+    load_file_from_url(remote_model_path, file_name=model_path)
+
+  if RAFT_model is None:
+    args = argparse.Namespace(**{
+        'model': ph.models_path + '/RAFT/raft-things.pth',
+        'mixed_precision': True,
+        'small': False,
+        'alternate_corr': False,
+        'path': ""
+    })
+
+    RAFT_model = torch.nn.DataParallel(RAFT(args))
+    RAFT_model.load_state_dict(torch.load(args.model))
+
+    RAFT_model = RAFT_model.module
+    RAFT_model.to(device)
+    RAFT_model.eval()
+
+  with torch.no_grad():
+    frame1_torch = torch.from_numpy(frame1).permute(2, 0, 1).float()[None].to(device)
+    frame2_torch = torch.from_numpy(frame2).permute(2, 0, 1).float()[None].to(device)
+
+    padder = InputPadder(frame1_torch.shape)
+    image1, image2 = padder.pad(frame1_torch, frame2_torch)
+
+    # estimate optical flow
+    _, next_flow = RAFT_model(image1, image2, iters=20, test_mode=True)
+    _, prev_flow = RAFT_model(image2, image1, iters=20, test_mode=True)
+
+    next_flow = next_flow[0].permute(1, 2, 0).cpu().numpy()
+    prev_flow = prev_flow[0].permute(1, 2, 0).cpu().numpy()
+
+    fb_flow = next_flow + prev_flow
+    fb_norm = np.linalg.norm(fb_flow, axis=2)
+
+    occlusion_mask = fb_norm[..., None].repeat(3, axis=-1)
+
+  next_flow = cv2.resize(next_flow, org_size)
+  prev_flow = cv2.resize(prev_flow, org_size)
+
+  return next_flow, prev_flow, occlusion_mask
+
+def compute_diff_map(next_flow, prev_flow, prev_frame, cur_frame, prev_frame_styled, args_dict):
+  h, w = cur_frame.shape[:2]
+  fl_w, fl_h = next_flow.shape[:2]
+
+  # normalize flow
+  next_flow = next_flow / np.array([fl_h,fl_w]) 
+  prev_flow = prev_flow / np.array([fl_h,fl_w])
+
+  # compute occlusion mask
+  fb_flow = next_flow + prev_flow
+  fb_norm = np.linalg.norm(fb_flow , axis=2) 
+
+  zero_flow_mask = np.clip(1 - np.linalg.norm(prev_flow, axis=-1)[...,None] * 20, 0, 1)
+  diff_mask_flow = fb_norm[..., None] * zero_flow_mask
+
+  # resize flow
+  next_flow = cv2.resize(next_flow, (w, h)) 
+  next_flow = (next_flow * np.array([h,w])).astype(np.float32)
+  prev_flow = cv2.resize(prev_flow, (w, h))
+  prev_flow = (prev_flow  * np.array([h,w])).astype(np.float32)
+
+  # Generate sampling grids
+  grid_y, grid_x = torch.meshgrid(torch.arange(0, h), torch.arange(0, w))
+  flow_grid = torch.stack((grid_x, grid_y), dim=0).float()
+  flow_grid += torch.from_numpy(prev_flow).permute(2, 0, 1)
+  flow_grid = flow_grid.unsqueeze(0)
+  flow_grid[:, 0, :, :] = 2 * flow_grid[:, 0, :, :] / (w - 1) - 1
+  flow_grid[:, 1, :, :] = 2 * flow_grid[:, 1, :, :] / (h - 1) - 1
+  flow_grid = flow_grid.permute(0, 2, 3, 1)
+
+  
+  prev_frame_torch = torch.from_numpy(prev_frame).float().unsqueeze(0).permute(0, 3, 1, 2) #N, C, H, W
+  prev_frame_styled_torch = torch.from_numpy(prev_frame_styled).float().unsqueeze(0).permute(0, 3, 1, 2) #N, C, H, W
+
+  warped_frame = torch.nn.functional.grid_sample(prev_frame_torch, flow_grid, mode="nearest", padding_mode="reflection", align_corners=True).permute(0, 2, 3, 1)[0].numpy()
+  warped_frame_styled = torch.nn.functional.grid_sample(prev_frame_styled_torch, flow_grid, mode="nearest", padding_mode="reflection", align_corners=True).permute(0, 2, 3, 1)[0].numpy()
+
+  #warped_frame = cv2.remap(prev_frame, flow_map, None, cv2.INTER_NEAREST, borderMode = cv2.BORDER_REFLECT)
+  #warped_frame_styled = cv2.remap(prev_frame_styled, flow_map, None, cv2.INTER_NEAREST, borderMode = cv2.BORDER_REFLECT)
+
+  
+  diff_mask_org = np.abs(warped_frame.astype(np.float32) - cur_frame.astype(np.float32)) / 255
+  diff_mask_org = diff_mask_org.max(axis = -1, keepdims=True)
+
+  diff_mask_stl = np.abs(warped_frame_styled.astype(np.float32) - cur_frame.astype(np.float32)) / 255
+  diff_mask_stl = diff_mask_stl.max(axis = -1, keepdims=True)
+
+  alpha_mask = np.maximum.reduce([diff_mask_flow * args_dict['occlusion_mask_flow_multiplier'] * 10, \
+                                  diff_mask_org * args_dict['occlusion_mask_difo_multiplier'], \
+                                  diff_mask_stl * args_dict['occlusion_mask_difs_multiplier']]) #
+  alpha_mask = alpha_mask.repeat(3, axis = -1)
+
+  #alpha_mask_blured = cv2.dilate(alpha_mask, np.ones((5, 5), np.float32))
+  if args_dict['occlusion_mask_blur'] > 0:
+    blur_filter_size = min(w,h) // 15 | 1
+    alpha_mask = cv2.GaussianBlur(alpha_mask, (blur_filter_size, blur_filter_size) , args_dict['occlusion_mask_blur'], cv2.BORDER_REFLECT)
+
+  alpha_mask = np.clip(alpha_mask, 0, 1)
+
+  return alpha_mask, warped_frame_styled
+
+def frames_norm(frame): return frame / 127.5 - 1
+
+def flow_norm(flow): return flow / 255
+
+def occl_norm(occl): return occl / 127.5 - 1
+
+def frames_renorm(frame): return (frame + 1) * 127.5
+
+def flow_renorm(flow): return flow * 255
+
+def occl_renorm(occl): return (occl + 1) * 127.5

SD-CN-Animation/scripts/core/txt2vid.py

@@ -0,0 +1,240 @@
+import sys, os
+
+import torch
+import gc
+import numpy as np
+from PIL import Image
+
+import modules.paths as ph
+from modules.shared import devices
+
+from scripts.core import utils, flow_utils
+from FloweR.model import FloweR
+
+import skimage
+import datetime
+import cv2
+import gradio as gr
+import time
+
+FloweR_model = None
+DEVICE = 'cpu'
+def FloweR_clear_memory():
+  global FloweR_model
+  del FloweR_model
+  gc.collect()
+  torch.cuda.empty_cache()
+  FloweR_model = None
+
+def FloweR_load_model(w, h):
+  global DEVICE, FloweR_model
+  DEVICE = devices.get_optimal_device()
+
+  model_path = ph.models_path + '/FloweR/FloweR_0.1.2.pth'
+  #remote_model_path = 'https://drive.google.com/uc?id=1K7gXUosgxU729_l-osl1HBU5xqyLsALv' #FloweR_0.1.1.pth
+  remote_model_path = 'https://drive.google.com/uc?id=1-UYsTXkdUkHLgtPK1Y5_7kKzCgzL_Z6o' #FloweR_0.1.2.pth
+
+  if not os.path.isfile(model_path):
+    from basicsr.utils.download_util import load_file_from_url
+    os.makedirs(os.path.dirname(model_path), exist_ok=True)
+    load_file_from_url(remote_model_path, file_name=model_path)
+
+
+  FloweR_model = FloweR(input_size = (h, w))
+  FloweR_model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+  # Move the model to the device
+  FloweR_model = FloweR_model.to(DEVICE)
+  FloweR_model.eval()
+
+def read_frame_from_video(input_video):
+  if input_video is None: return None
+
+  # Reading video file
+  if input_video.isOpened():
+    ret, cur_frame = input_video.read()
+    if cur_frame is not None: 
+      cur_frame = cv2.cvtColor(cur_frame, cv2.COLOR_BGR2RGB) 
+  else:
+    cur_frame = None
+    input_video.release()
+    input_video = None
+  
+  return cur_frame
+
+def start_process(*args):
+    processing_start_time = time.time()
+    args_dict = utils.args_to_dict(*args)
+    args_dict = utils.get_mode_args('t2v', args_dict)
+
+    # Open the input video file
+    input_video = None
+    if args_dict['file'] is not None:
+      input_video = cv2.VideoCapture(args_dict['file'].name)
+
+    # Create an output video file with the same fps, width, and height as the input video
+    output_video_name = f'outputs/sd-cn-animation/txt2vid/{datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.mp4'
+    output_video_folder = os.path.splitext(output_video_name)[0]
+    os.makedirs(os.path.dirname(output_video_name), exist_ok=True)
+
+    #if args_dict['save_frames_check']: 
+    os.makedirs(output_video_folder, exist_ok=True)
+
+    # Writing to current params to params.json
+    setts_json = utils.export_settings(*args)
+    with open(os.path.join(output_video_folder, "params.json"), "w") as outfile:
+      outfile.write(setts_json)
+
+    curr_frame = None
+    prev_frame = None
+    
+    def save_result_to_image(image, ind):
+      if args_dict['save_frames_check']: 
+        cv2.imwrite(os.path.join(output_video_folder, f'{ind:05d}.png'), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+
+    def set_cn_frame_input():
+      if args_dict['cn_frame_send'] == 0: # Current generated frame"
+        pass
+      elif args_dict['cn_frame_send'] == 1: # Current generated frame"
+        if curr_frame is not None:
+          utils.set_CNs_input_image(args_dict, Image.fromarray(curr_frame), set_references=True)
+      elif args_dict['cn_frame_send'] == 2: # Previous generated frame
+        if prev_frame is not None:
+          utils.set_CNs_input_image(args_dict, Image.fromarray(prev_frame), set_references=True)
+      elif args_dict['cn_frame_send'] == 3: # Current reference video frame
+        if input_video is not None:
+          curr_video_frame = read_frame_from_video(input_video)
+          curr_video_frame = cv2.resize(curr_video_frame, (args_dict['width'], args_dict['height']))
+          utils.set_CNs_input_image(args_dict, Image.fromarray(curr_video_frame), set_references=True)
+        else:
+          raise Exception('There is no input video! Set it up first.')
+      else:
+        raise Exception('Incorrect cn_frame_send mode!')
+
+    set_cn_frame_input()
+
+    if args_dict['init_image'] is not None:
+      #resize array to args_dict['width'], args_dict['height']
+      image_array=args_dict['init_image']#this is a numpy array
+      init_frame = np.array(Image.fromarray(image_array).resize((args_dict['width'], args_dict['height'])).convert('RGB'))
+      processed_frame = init_frame.copy()
+    else:
+      processed_frames, _, _, _ = utils.txt2img(args_dict)
+      processed_frame = np.array(processed_frames[0])[...,:3]
+      #if input_video is not None:
+      #  processed_frame = skimage.exposure.match_histograms(processed_frame, curr_video_frame, channel_axis=-1)
+      processed_frame = np.clip(processed_frame, 0, 255).astype(np.uint8)
+      init_frame = processed_frame.copy()
+
+    output_video = cv2.VideoWriter(output_video_name, cv2.VideoWriter_fourcc(*'mp4v'), args_dict['fps'], (args_dict['width'], args_dict['height']))
+    output_video.write(cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR))
+
+    stat = f"Frame: 1 / {args_dict['length']}; " + utils.get_time_left(1, args_dict['length'], processing_start_time)
+    utils.shared.is_interrupted = False
+
+    save_result_to_image(processed_frame, 1)
+    yield stat, init_frame, None, None, processed_frame, None, gr.Button.update(interactive=False), gr.Button.update(interactive=True)
+
+    org_size = args_dict['width'], args_dict['height']
+    size = args_dict['width'] // 128 * 128, args_dict['height'] // 128 * 128
+    FloweR_load_model(size[0], size[1])
+
+    clip_frames = np.zeros((4, size[1], size[0], 3), dtype=np.uint8)
+
+    prev_frame = init_frame
+
+    for ind in range(args_dict['length'] - 1): 
+      if utils.shared.is_interrupted: break
+
+      args_dict = utils.args_to_dict(*args)
+      args_dict = utils.get_mode_args('t2v', args_dict)
+
+      clip_frames = np.roll(clip_frames, -1, axis=0)
+      clip_frames[-1] = cv2.resize(prev_frame[...,:3], size)
+      clip_frames_torch = flow_utils.frames_norm(torch.from_numpy(clip_frames).to(DEVICE, dtype=torch.float32))
+
+      with torch.no_grad():
+        pred_data = FloweR_model(clip_frames_torch.unsqueeze(0))[0]
+
+      pred_flow = flow_utils.flow_renorm(pred_data[...,:2]).cpu().numpy()
+      pred_occl = flow_utils.occl_renorm(pred_data[...,2:3]).cpu().numpy().repeat(3, axis = -1)
+      pred_next = flow_utils.frames_renorm(pred_data[...,3:6]).cpu().numpy()
+      
+      pred_occl = np.clip(pred_occl * 10, 0, 255).astype(np.uint8)
+      pred_next = np.clip(pred_next, 0, 255).astype(np.uint8)
+      
+      pred_flow = cv2.resize(pred_flow, org_size)
+      pred_occl = cv2.resize(pred_occl, org_size)
+      pred_next = cv2.resize(pred_next, org_size)
+
+      curr_frame = pred_next.copy()
+
+      '''
+      pred_flow = pred_flow / (1 + np.linalg.norm(pred_flow, axis=-1, keepdims=True) * 0.05) 
+      pred_flow = cv2.GaussianBlur(pred_flow, (31,31), 1, cv2.BORDER_REFLECT_101)
+    
+      pred_occl = cv2.GaussianBlur(pred_occl, (21,21), 2, cv2.BORDER_REFLECT_101)
+      pred_occl = (np.abs(pred_occl / 255) ** 1.5) * 255
+      pred_occl = np.clip(pred_occl * 25, 0, 255).astype(np.uint8)
+
+      flow_map = pred_flow.copy()
+      flow_map[:,:,0] += np.arange(args_dict['width'])
+      flow_map[:,:,1] += np.arange(args_dict['height'])[:,np.newaxis]
+
+      warped_frame = cv2.remap(prev_frame, flow_map, None, cv2.INTER_NEAREST, borderMode = cv2.BORDER_REFLECT_101)
+      alpha_mask = pred_occl / 255.
+      #alpha_mask = np.clip(alpha_mask + np.random.normal(0, 0.4, size = alpha_mask.shape), 0, 1)
+      curr_frame = pred_next.astype(float) * alpha_mask + warped_frame.astype(float) * (1 - alpha_mask)
+      curr_frame = np.clip(curr_frame, 0, 255).astype(np.uint8)
+      #curr_frame = warped_frame.copy()
+      '''
+
+      set_cn_frame_input()
+
+      args_dict['mode'] = 4
+      args_dict['init_img'] = Image.fromarray(pred_next)
+      args_dict['mask_img'] = Image.fromarray(pred_occl)
+      args_dict['seed'] = -1
+      args_dict['denoising_strength'] = args_dict['processing_strength']
+
+      processed_frames, _, _, _ = utils.img2img(args_dict)
+      processed_frame = np.array(processed_frames[0])[...,:3]
+      #if input_video is not None:
+      #  processed_frame = skimage.exposure.match_histograms(processed_frame, curr_video_frame, channel_axis=-1)
+      #else: 
+      processed_frame = skimage.exposure.match_histograms(processed_frame, init_frame, channel_axis=-1)
+      processed_frame = np.clip(processed_frame, 0, 255).astype(np.uint8)
+
+      args_dict['mode'] = 0
+      args_dict['init_img'] = Image.fromarray(processed_frame)
+      args_dict['mask_img'] = None
+      args_dict['seed'] = -1
+      args_dict['denoising_strength'] = args_dict['fix_frame_strength']
+
+      #utils.set_CNs_input_image(args_dict, Image.fromarray(curr_frame))
+      processed_frames, _, _, _ = utils.img2img(args_dict)
+      processed_frame = np.array(processed_frames[0])[...,:3]
+      #if input_video is not None:
+      #  processed_frame = skimage.exposure.match_histograms(processed_frame, curr_video_frame, channel_axis=-1)
+      #else: 
+      processed_frame = skimage.exposure.match_histograms(processed_frame, init_frame, channel_axis=-1)
+      processed_frame = np.clip(processed_frame, 0, 255).astype(np.uint8)
+
+      output_video.write(cv2.cvtColor(processed_frame, cv2.COLOR_RGB2BGR))
+      prev_frame = processed_frame.copy()
+
+      save_result_to_image(processed_frame, ind + 2)
+      stat = f"Frame: {ind + 2} / {args_dict['length']}; " + utils.get_time_left(ind+2, args_dict['length'], processing_start_time)
+      yield stat, curr_frame, pred_occl, pred_next, processed_frame, None, gr.Button.update(interactive=False), gr.Button.update(interactive=True)
+
+    if input_video is not None: input_video.release()
+    output_video.release()
+    FloweR_clear_memory()
+
+    curr_frame = gr.Image.update()
+    occlusion_mask = gr.Image.update()
+    warped_styled_frame_ = gr.Image.update() 
+    processed_frame = gr.Image.update()
+
+    # print('TOTAL TIME:', int(time.time() - processing_start_time))
+
+    yield 'done', curr_frame, occlusion_mask, warped_styled_frame_, processed_frame, output_video_name, gr.Button.update(interactive=True), gr.Button.update(interactive=False)

SD-CN-Animation/scripts/core/utils.py

@@ -0,0 +1,432 @@
+class shared:
+  is_interrupted = False
+  v2v_custom_inputs_size = 0
+  t2v_custom_inputs_size = 0
+
+def get_component_names():
+  components_list = [
+    'glo_sdcn_process_mode',
+    'v2v_file', 'v2v_width', 'v2v_height', 'v2v_prompt', 'v2v_n_prompt', 'v2v_cfg_scale', 'v2v_seed', 'v2v_processing_strength', 'v2v_fix_frame_strength', 
+    'v2v_sampler_index', 'v2v_steps', 'v2v_override_settings',
+    'v2v_occlusion_mask_blur', 'v2v_occlusion_mask_trailing', 'v2v_occlusion_mask_flow_multiplier', 'v2v_occlusion_mask_difo_multiplier', 'v2v_occlusion_mask_difs_multiplier',
+    'v2v_step_1_processing_mode', 'v2v_step_1_blend_alpha', 'v2v_step_1_seed', 'v2v_step_2_seed',
+    't2v_file','t2v_init_image', 't2v_width', 't2v_height', 't2v_prompt', 't2v_n_prompt', 't2v_cfg_scale', 't2v_seed', 't2v_processing_strength', 't2v_fix_frame_strength',
+    't2v_sampler_index', 't2v_steps', 't2v_length', 't2v_fps', 't2v_cn_frame_send',
+    'glo_save_frames_check'
+  ]
+
+  return components_list
+
+def args_to_dict(*args): # converts list of argumets into dictionary for better handling of it
+  args_list = get_component_names()
+
+  # set default values for params that were not specified
+  args_dict = {
+    # video to video params
+    'v2v_mode': 0,
+    'v2v_prompt': '',
+    'v2v_n_prompt': '',
+    'v2v_prompt_styles': [],
+    'v2v_init_video': None, # Always required
+
+    'v2v_steps': 15,
+    'v2v_sampler_index': 0, # 'Euler a'    
+    'v2v_mask_blur': 0,
+
+    'v2v_inpainting_fill': 1, # original
+    'v2v_restore_faces': False,
+    'v2v_tiling': False,
+    'v2v_n_iter': 1,
+    'v2v_batch_size': 1,
+    'v2v_cfg_scale': 5.5,
+    'v2v_image_cfg_scale': 1.5,
+    'v2v_denoising_strength': 0.75,
+    'v2v_processing_strength': 0.85,
+    'v2v_fix_frame_strength': 0.15,
+    'v2v_seed': -1,
+    'v2v_subseed': -1,
+    'v2v_subseed_strength': 0,
+    'v2v_seed_resize_from_h': 512,
+    'v2v_seed_resize_from_w': 512,
+    'v2v_seed_enable_extras': False,
+    'v2v_height': 512,
+    'v2v_width': 512,
+    'v2v_resize_mode': 1,
+    'v2v_inpaint_full_res': True,
+    'v2v_inpaint_full_res_padding': 0,
+    'v2v_inpainting_mask_invert': False,
+
+    # text to video params
+    't2v_mode': 4,
+    't2v_prompt': '',
+    't2v_n_prompt': '',
+    't2v_prompt_styles': [],
+    't2v_init_img': None,
+    't2v_mask_img': None,
+
+    't2v_steps': 15,
+    't2v_sampler_index': 0, # 'Euler a'    
+    't2v_mask_blur': 0,
+
+    't2v_inpainting_fill': 1, # original
+    't2v_restore_faces': False,
+    't2v_tiling': False,
+    't2v_n_iter': 1,
+    't2v_batch_size': 1,
+    't2v_cfg_scale': 5.5,
+    't2v_image_cfg_scale': 1.5,
+    't2v_denoising_strength': 0.75,
+    't2v_processing_strength': 0.85,
+    't2v_fix_frame_strength': 0.15,
+    't2v_seed': -1,
+    't2v_subseed': -1,
+    't2v_subseed_strength': 0,
+    't2v_seed_resize_from_h': 512,
+    't2v_seed_resize_from_w': 512,
+    't2v_seed_enable_extras': False,
+    't2v_height': 512,
+    't2v_width': 512,
+    't2v_resize_mode': 1,
+    't2v_inpaint_full_res': True,
+    't2v_inpaint_full_res_padding': 0,
+    't2v_inpainting_mask_invert': False,
+
+    't2v_override_settings': [],
+    #'t2v_script_inputs': [0],
+
+    't2v_fps': 12,
+  }
+
+  args = list(args)
+
+  for i in range(len(args_list)):
+    if (args[i] is None) and (args_list[i] in args_dict):
+      #args[i] = args_dict[args_list[i]] 
+      pass
+    else:
+      args_dict[args_list[i]] = args[i]
+
+  args_dict['v2v_script_inputs'] = args[len(args_list):len(args_list)+shared.v2v_custom_inputs_size]
+  #print('v2v_script_inputs', args_dict['v2v_script_inputs'])
+  args_dict['t2v_script_inputs'] = args[len(args_list)+shared.v2v_custom_inputs_size:]
+  #print('t2v_script_inputs', args_dict['t2v_script_inputs'])
+  return args_dict
+
+def get_mode_args(mode, args_dict):
+  mode_args_dict = {}
+  for key, value in args_dict.items():
+    if key[:3] in [mode, 'glo'] :
+      mode_args_dict[key[4:]] = value
+
+  return mode_args_dict
+
+def set_CNs_input_image(args_dict, image, set_references = False):
+  for script_input in args_dict['script_inputs']:
+    if type(script_input).__name__ == 'UiControlNetUnit':
+      if script_input.module not in ["reference_only", "reference_adain", "reference_adain+attn"] or set_references:
+        script_input.image = np.array(image)
+        script_input.batch_images = [np.array(image)]
+
+import time
+import datetime
+
+def get_time_left(ind, length, processing_start_time):
+  s_passed = int(time.time() - processing_start_time)
+  time_passed = datetime.timedelta(seconds=s_passed)
+  s_left = int(s_passed / ind * (length - ind))
+  time_left = datetime.timedelta(seconds=s_left)
+  return f"Time elapsed: {time_passed}; Time left: {time_left};"
+
+import numpy as np
+from PIL import Image, ImageOps, ImageFilter, ImageEnhance, ImageChops
+from types import SimpleNamespace  
+
+from modules.generation_parameters_copypaste import create_override_settings_dict
+from modules.processing import Processed, StableDiffusionProcessingImg2Img, StableDiffusionProcessingTxt2Img, process_images
+import modules.processing as processing
+from modules.ui import plaintext_to_html
+import modules.images as images
+import modules.scripts
+from modules.shared import opts, devices, state
+from modules import devices, sd_samplers, img2img
+from modules import shared, sd_hijack, lowvram
+
+# TODO: Refactor all the code below
+
+def process_img(p, input_img, output_dir, inpaint_mask_dir, args):
+    processing.fix_seed(p)
+
+    #images = shared.listfiles(input_dir)
+    images = [input_img]
+
+    is_inpaint_batch = False
+    #if inpaint_mask_dir:
+    #    inpaint_masks = shared.listfiles(inpaint_mask_dir)
+    #    is_inpaint_batch = len(inpaint_masks) > 0
+    #if is_inpaint_batch:
+    #    print(f"\nInpaint batch is enabled. {len(inpaint_masks)} masks found.")
+
+    #print(f"Will process {len(images)} images, creating {p.n_iter * p.batch_size} new images for each.")
+
+    save_normally = output_dir == ''
+
+    p.do_not_save_grid = True
+    p.do_not_save_samples = not save_normally
+
+    state.job_count = len(images) * p.n_iter
+
+    generated_images = []
+    for i, image in enumerate(images):
+        state.job = f"{i+1} out of {len(images)}"
+        if state.skipped:
+            state.skipped = False
+
+        if state.interrupted:
+            break
+
+        img = image #Image.open(image)
+        # Use the EXIF orientation of photos taken by smartphones.
+        img = ImageOps.exif_transpose(img)
+        p.init_images = [img] * p.batch_size
+
+        #if is_inpaint_batch:
+        #    # try to find corresponding mask for an image using simple filename matching
+        #    mask_image_path = os.path.join(inpaint_mask_dir, os.path.basename(image))
+        #    # if not found use first one ("same mask for all images" use-case)
+        #    if not mask_image_path in inpaint_masks:
+        #        mask_image_path = inpaint_masks[0]
+        #    mask_image = Image.open(mask_image_path)
+        #    p.image_mask = mask_image
+
+        proc = modules.scripts.scripts_img2img.run(p, *args)
+        if proc is None:
+            proc = process_images(p)
+            generated_images.append(proc.images[0])
+
+        #for n, processed_image in enumerate(proc.images):
+        #    filename = os.path.basename(image)
+
+        #    if n > 0:
+        #        left, right = os.path.splitext(filename)
+        #        filename = f"{left}-{n}{right}"
+
+        #    if not save_normally:
+        #        os.makedirs(output_dir, exist_ok=True)
+        #        if processed_image.mode == 'RGBA':
+        #            processed_image = processed_image.convert("RGB")
+        #        processed_image.save(os.path.join(output_dir, filename))
+
+    return generated_images
+
+def img2img(args_dict):  
+    args = SimpleNamespace(**args_dict)
+    override_settings = create_override_settings_dict(args.override_settings)
+
+    is_batch = args.mode == 5
+
+    if args.mode == 0:  # img2img
+        image = args.init_img.convert("RGB")
+        mask = None
+    elif args.mode == 1:  # img2img sketch
+        image = args.sketch.convert("RGB")
+        mask = None
+    elif args.mode == 2:  # inpaint
+        image, mask = args.init_img_with_mask["image"], args.init_img_with_mask["mask"]
+        alpha_mask = ImageOps.invert(image.split()[-1]).convert('L').point(lambda x: 255 if x > 0 else 0, mode='1')
+        mask = ImageChops.lighter(alpha_mask, mask.convert('L')).convert('L')
+        image = image.convert("RGB")
+    elif args.mode == 3:  # inpaint sketch
+        image = args.inpaint_color_sketch
+        orig = args.inpaint_color_sketch_orig or args.inpaint_color_sketch
+        pred = np.any(np.array(image) != np.array(orig), axis=-1)
+        mask = Image.fromarray(pred.astype(np.uint8) * 255, "L")
+        mask = ImageEnhance.Brightness(mask).enhance(1 - args.mask_alpha / 100)
+        blur = ImageFilter.GaussianBlur(args.mask_blur)
+        image = Image.composite(image.filter(blur), orig, mask.filter(blur))
+        image = image.convert("RGB")
+    elif args.mode == 4:  # inpaint upload mask
+        #image = args.init_img_inpaint
+        #mask = args.init_mask_inpaint
+
+        image = args.init_img.convert("RGB")
+        mask = args.mask_img.convert("L")
+    else:
+        image = None
+        mask = None
+
+    # Use the EXIF orientation of photos taken by smartphones.
+    if image is not None:
+        image = ImageOps.exif_transpose(image)
+
+    assert 0. <= args.denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'
+
+    p = StableDiffusionProcessingImg2Img(
+        sd_model=shared.sd_model,
+        outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples,
+        outpath_grids=opts.outdir_grids or opts.outdir_img2img_grids,
+        prompt=args.prompt,
+        negative_prompt=args.n_prompt,
+        styles=args.prompt_styles,
+        seed=args.seed,
+        subseed=args.subseed,
+        subseed_strength=args.subseed_strength,
+        seed_resize_from_h=args.seed_resize_from_h,
+        seed_resize_from_w=args.seed_resize_from_w,
+        seed_enable_extras=args.seed_enable_extras,
+        sampler_name=sd_samplers.samplers_for_img2img[args.sampler_index].name,
+        batch_size=args.batch_size,
+        n_iter=args.n_iter,
+        steps=args.steps,
+        cfg_scale=args.cfg_scale,
+        width=args.width,
+        height=args.height,
+        restore_faces=args.restore_faces,
+        tiling=args.tiling,
+        init_images=[image],
+        mask=mask,
+        mask_blur=args.mask_blur,
+        inpainting_fill=args.inpainting_fill,
+        resize_mode=args.resize_mode,
+        denoising_strength=args.denoising_strength,
+        image_cfg_scale=args.image_cfg_scale,
+        inpaint_full_res=args.inpaint_full_res,
+        inpaint_full_res_padding=args.inpaint_full_res_padding,
+        inpainting_mask_invert=args.inpainting_mask_invert,
+        override_settings=override_settings,
+    )
+
+    p.scripts = modules.scripts.scripts_img2img
+    p.script_args = args.script_inputs
+
+    #if shared.cmd_opts.enable_console_prompts:
+    #    print(f"\nimg2img: {args.prompt}", file=shared.progress_print_out)
+
+    if mask:
+        p.extra_generation_params["Mask blur"] = args.mask_blur
+    
+    '''
+    if is_batch:
+        ...
+    #    assert not shared.cmd_opts.hide_ui_dir_config, "Launched with --hide-ui-dir-config, batch img2img disabled"
+    #    process_batch(p, img2img_batch_input_dir, img2img_batch_output_dir, img2img_batch_inpaint_mask_dir, args.script_inputs)
+    #    processed = Processed(p, [], p.seed, "")
+    else:
+        processed = modules.scripts.scripts_img2img.run(p, *args.script_inputs)
+        if processed is None:
+            processed = process_images(p)
+    '''
+
+    generated_images = process_img(p, image, None, '', args.script_inputs)
+    processed = Processed(p, [], p.seed, "")
+    p.close()
+
+    shared.total_tqdm.clear()
+
+    generation_info_js = processed.js()
+    #if opts.samples_log_stdout:
+    #    print(generation_info_js)
+
+    #if opts.do_not_show_images:
+    #    processed.images = []
+
+    #print(generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments))
+    return generated_images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
+
+def txt2img(args_dict):  
+    args = SimpleNamespace(**args_dict)
+    override_settings = create_override_settings_dict(args.override_settings)
+
+    p = StableDiffusionProcessingTxt2Img(
+        sd_model=shared.sd_model,
+        outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
+        outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids,
+        prompt=args.prompt,
+        styles=args.prompt_styles,
+        negative_prompt=args.n_prompt,
+        seed=args.seed,
+        subseed=args.subseed,
+        subseed_strength=args.subseed_strength,
+        seed_resize_from_h=args.seed_resize_from_h,
+        seed_resize_from_w=args.seed_resize_from_w,
+        seed_enable_extras=args.seed_enable_extras,
+        sampler_name=sd_samplers.samplers[args.sampler_index].name,
+        batch_size=args.batch_size,
+        n_iter=args.n_iter,
+        steps=args.steps,
+        cfg_scale=args.cfg_scale,
+        width=args.width,
+        height=args.height,
+        restore_faces=args.restore_faces,
+        tiling=args.tiling,
+        #enable_hr=args.enable_hr,
+        #denoising_strength=args.denoising_strength if enable_hr else None,
+        #hr_scale=hr_scale,
+        #hr_upscaler=hr_upscaler,
+        #hr_second_pass_steps=hr_second_pass_steps,
+        #hr_resize_x=hr_resize_x,
+        #hr_resize_y=hr_resize_y,
+        override_settings=override_settings,
+    )
+
+    p.scripts = modules.scripts.scripts_txt2img
+    p.script_args = args.script_inputs
+
+    #if cmd_opts.enable_console_prompts:
+    #    print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
+
+    processed = modules.scripts.scripts_txt2img.run(p, *args.script_inputs)
+
+    if processed is None:
+        processed = process_images(p)
+
+    p.close()
+
+    shared.total_tqdm.clear()
+
+    generation_info_js = processed.js()
+    #if opts.samples_log_stdout:
+    #    print(generation_info_js)
+
+    #if opts.do_not_show_images:
+    #    processed.images = []
+
+    return processed.images, generation_info_js, plaintext_to_html(processed.info), plaintext_to_html(processed.comments)
+
+
+import json
+def get_json(obj):
+  return json.loads(
+    json.dumps(obj, default=lambda o: getattr(o, '__dict__', str(o)))
+  )
+
+def export_settings(*args):
+  args_dict = args_to_dict(*args)
+  if args[0] == 'vid2vid':
+    args_dict = get_mode_args('v2v', args_dict)
+  elif args[0] == 'txt2vid':
+    args_dict = get_mode_args('t2v', args_dict)
+  else:
+    msg = f"Unsupported processing mode: '{args[0]}'"
+    raise Exception(msg)
+  
+  # convert CN params into a readable dict
+  cn_remove_list = ['low_vram', 'is_ui', 'input_mode', 'batch_images', 'output_dir', 'loopback', 'image']
+
+  args_dict['ControlNets'] = []
+  for script_input in args_dict['script_inputs']:
+    if type(script_input).__name__ == 'UiControlNetUnit':
+      cn_values_dict = get_json(script_input)
+      if cn_values_dict['enabled']:
+        for key in cn_remove_list:
+          if key in cn_values_dict: del cn_values_dict[key]
+        args_dict['ControlNets'].append(cn_values_dict)
+  
+  # remove unimportant values
+  remove_list = ['save_frames_check', 'restore_faces', 'prompt_styles', 'mask_blur', 'inpainting_fill', 'tiling', 'n_iter', 'batch_size', 'subseed', 'subseed_strength', 'seed_resize_from_h', \
+                 'seed_resize_from_w', 'seed_enable_extras', 'resize_mode', 'inpaint_full_res', 'inpaint_full_res_padding', 'inpainting_mask_invert', 'file', 'denoising_strength', \
+                 'override_settings', 'script_inputs', 'init_img', 'mask_img', 'mode', 'init_video']
+  
+  for key in remove_list:
+    if key in args_dict: del args_dict[key]
+
+  return json.dumps(args_dict, indent=2, default=lambda o: getattr(o, '__dict__', str(o)))

SD-CN-Animation/scripts/core/vid2vid.py

@@ -0,0 +1,270 @@
+import sys, os
+
+import math
+import os
+import sys
+import traceback
+
+import numpy as np
+from PIL import Image
+
+from modules import devices, sd_samplers
+from modules import shared, sd_hijack, lowvram
+
+from modules.shared import devices
+import modules.shared as shared
+
+import gc
+import cv2
+import gradio as gr
+
+import time
+import skimage
+import datetime
+
+from scripts.core.flow_utils import RAFT_estimate_flow, RAFT_clear_memory, compute_diff_map
+from scripts.core import utils
+
+class sdcn_anim_tmp:
+  prepear_counter = 0
+  process_counter = 0
+  input_video = None
+  output_video = None
+  curr_frame = None
+  prev_frame = None
+  prev_frame_styled = None
+  prev_frame_alpha_mask = None
+  fps = None
+  total_frames = None
+  prepared_frames = None
+  prepared_next_flows = None
+  prepared_prev_flows = None
+  frames_prepared = False
+
+def read_frame_from_video():
+  # Reading video file
+  if sdcn_anim_tmp.input_video.isOpened():
+    ret, cur_frame = sdcn_anim_tmp.input_video.read()
+    if cur_frame is not None: 
+      cur_frame = cv2.cvtColor(cur_frame, cv2.COLOR_BGR2RGB) 
+  else:
+    cur_frame = None
+    sdcn_anim_tmp.input_video.release()
+  
+  return cur_frame
+
+def get_cur_stat():
+  stat =  f'Frames prepared: {sdcn_anim_tmp.prepear_counter + 1} / {sdcn_anim_tmp.total_frames}; '
+  stat += f'Frames processed: {sdcn_anim_tmp.process_counter + 1} / {sdcn_anim_tmp.total_frames}; '
+  return stat
+
+def clear_memory_from_sd():
+  if shared.sd_model is not None:
+    sd_hijack.model_hijack.undo_hijack(shared.sd_model)
+    try:
+      lowvram.send_everything_to_cpu()
+    except Exception as e:
+      ...
+    del shared.sd_model
+    shared.sd_model = None
+  gc.collect()
+  devices.torch_gc()
+
+def start_process(*args):
+  processing_start_time = time.time()
+  args_dict = utils.args_to_dict(*args)
+  args_dict = utils.get_mode_args('v2v', args_dict)
+  
+  sdcn_anim_tmp.process_counter = 0
+  sdcn_anim_tmp.prepear_counter = 0
+
+  # Open the input video file
+  sdcn_anim_tmp.input_video = cv2.VideoCapture(args_dict['file'].name)
+  
+  # Get useful info from the source video
+  sdcn_anim_tmp.fps = int(sdcn_anim_tmp.input_video.get(cv2.CAP_PROP_FPS))
+  sdcn_anim_tmp.total_frames = int(sdcn_anim_tmp.input_video.get(cv2.CAP_PROP_FRAME_COUNT))
+  loop_iterations = (sdcn_anim_tmp.total_frames-1) * 2
+
+  # Create an output video file with the same fps, width, and height as the input video
+  output_video_name = f'outputs/sd-cn-animation/vid2vid/{datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.mp4'
+  output_video_folder = os.path.splitext(output_video_name)[0]
+  os.makedirs(os.path.dirname(output_video_name), exist_ok=True)
+  
+  if args_dict['save_frames_check']:
+    os.makedirs(output_video_folder, exist_ok=True)
+
+  def save_result_to_image(image, ind):
+    if args_dict['save_frames_check']: 
+      cv2.imwrite(os.path.join(output_video_folder, f'{ind:05d}.png'), cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
+
+  sdcn_anim_tmp.output_video = cv2.VideoWriter(output_video_name, cv2.VideoWriter_fourcc(*'mp4v'), sdcn_anim_tmp.fps, (args_dict['width'], args_dict['height']))
+
+  curr_frame = read_frame_from_video()
+  curr_frame = cv2.resize(curr_frame, (args_dict['width'], args_dict['height']))
+  sdcn_anim_tmp.prepared_frames = np.zeros((11, args_dict['height'], args_dict['width'], 3), dtype=np.uint8)
+  sdcn_anim_tmp.prepared_next_flows = np.zeros((10, args_dict['height'], args_dict['width'], 2))
+  sdcn_anim_tmp.prepared_prev_flows = np.zeros((10, args_dict['height'], args_dict['width'], 2))
+  sdcn_anim_tmp.prepared_frames[0] = curr_frame
+
+  args_dict['init_img'] = Image.fromarray(curr_frame)
+  utils.set_CNs_input_image(args_dict, Image.fromarray(curr_frame))
+  processed_frames, _, _, _ = utils.img2img(args_dict)
+  processed_frame = np.array(processed_frames[0])[...,:3]
+  processed_frame = skimage.exposure.match_histograms(processed_frame, curr_frame, channel_axis=None)
+  processed_frame = np.clip(processed_frame, 0, 255).astype(np.uint8)
+  #print('Processed frame ', 0)
+
+  sdcn_anim_tmp.curr_frame = curr_frame
+  sdcn_anim_tmp.prev_frame = curr_frame.copy()
+  sdcn_anim_tmp.prev_frame_styled = processed_frame.copy()
+  utils.shared.is_interrupted = False
+
+  save_result_to_image(processed_frame, 1)
+  stat = get_cur_stat() + utils.get_time_left(1, loop_iterations, processing_start_time)
+  yield stat, sdcn_anim_tmp.curr_frame, None, None, processed_frame, None, gr.Button.update(interactive=False), gr.Button.update(interactive=True)
+
+  for step in range(loop_iterations):
+    if utils.shared.is_interrupted: break
+    
+    args_dict = utils.args_to_dict(*args)
+    args_dict = utils.get_mode_args('v2v', args_dict)
+
+    occlusion_mask = None
+    prev_frame = None
+    curr_frame = sdcn_anim_tmp.curr_frame
+    warped_styled_frame_ = gr.Image.update()
+    processed_frame = gr.Image.update()
+
+    prepare_steps = 10
+    if sdcn_anim_tmp.process_counter % prepare_steps == 0 and not sdcn_anim_tmp.frames_prepared: # prepare next 10 frames for processing
+        #clear_memory_from_sd()
+        device = devices.get_optimal_device()
+
+        curr_frame = read_frame_from_video()
+        if curr_frame is not None: 
+            curr_frame = cv2.resize(curr_frame, (args_dict['width'], args_dict['height']))
+            prev_frame = sdcn_anim_tmp.prev_frame.copy()
+
+            next_flow, prev_flow, occlusion_mask = RAFT_estimate_flow(prev_frame, curr_frame, device=device)
+            occlusion_mask = np.clip(occlusion_mask * 0.1 * 255, 0, 255).astype(np.uint8)
+
+            cn = sdcn_anim_tmp.prepear_counter % 10
+            if sdcn_anim_tmp.prepear_counter % 10 == 0:
+                sdcn_anim_tmp.prepared_frames[cn] = sdcn_anim_tmp.prev_frame
+            sdcn_anim_tmp.prepared_frames[cn + 1] = curr_frame.copy()
+            sdcn_anim_tmp.prepared_next_flows[cn] = next_flow.copy()
+            sdcn_anim_tmp.prepared_prev_flows[cn] = prev_flow.copy()
+            #print('Prepared frame ', cn+1)
+
+            sdcn_anim_tmp.prev_frame = curr_frame.copy()
+
+        sdcn_anim_tmp.prepear_counter += 1
+        if sdcn_anim_tmp.prepear_counter % prepare_steps == 0 or \
+        sdcn_anim_tmp.prepear_counter >= sdcn_anim_tmp.total_frames - 1 or \
+        curr_frame is None:
+            # Remove RAFT from memory
+            RAFT_clear_memory()
+            sdcn_anim_tmp.frames_prepared = True
+    else:
+        # process frame
+        sdcn_anim_tmp.frames_prepared = False
+
+        cn = sdcn_anim_tmp.process_counter % 10 
+        curr_frame = sdcn_anim_tmp.prepared_frames[cn+1][...,:3]
+        prev_frame = sdcn_anim_tmp.prepared_frames[cn][...,:3]
+        next_flow = sdcn_anim_tmp.prepared_next_flows[cn]
+        prev_flow = sdcn_anim_tmp.prepared_prev_flows[cn]
+
+        ### STEP 1
+        alpha_mask, warped_styled_frame = compute_diff_map(next_flow, prev_flow, prev_frame, curr_frame, sdcn_anim_tmp.prev_frame_styled, args_dict)
+        warped_styled_frame_ = warped_styled_frame.copy()
+
+        #fl_w, fl_h = prev_flow.shape[:2]
+        #prev_flow_n = prev_flow / np.array([fl_h,fl_w])
+        #flow_mask = np.clip(1 - np.linalg.norm(prev_flow_n, axis=-1)[...,None] * 20, 0, 1)
+        #alpha_mask = alpha_mask * flow_mask
+
+        if sdcn_anim_tmp.process_counter > 0 and args_dict['occlusion_mask_trailing']:
+            alpha_mask = alpha_mask + sdcn_anim_tmp.prev_frame_alpha_mask * 0.5
+        sdcn_anim_tmp.prev_frame_alpha_mask = alpha_mask
+
+        # alpha_mask = np.round(alpha_mask * 8) / 8 #> 0.3
+        alpha_mask = np.clip(alpha_mask, 0, 1)
+        occlusion_mask = np.clip(alpha_mask * 255, 0, 255).astype(np.uint8)
+
+        # fix warped styled frame from duplicated that occures on the places where flow is zero, but only because there is no place to get the color from
+        warped_styled_frame = curr_frame.astype(float) * alpha_mask + warped_styled_frame.astype(float) * (1 - alpha_mask)
+
+        # process current frame
+        # TODO: convert args_dict into separate dict that stores only params necessery for img2img processing
+        img2img_args_dict = args_dict #copy.deepcopy(args_dict)
+        img2img_args_dict['denoising_strength'] = args_dict['processing_strength']
+        if args_dict['step_1_processing_mode'] == 0: # Process full image then blend in occlusions
+          img2img_args_dict['mode'] = 0
+          img2img_args_dict['mask_img'] = None #Image.fromarray(occlusion_mask)
+        elif args_dict['step_1_processing_mode'] == 1: # Inpaint occlusions
+          img2img_args_dict['mode'] = 4
+          img2img_args_dict['mask_img'] = Image.fromarray(occlusion_mask)
+        else:
+           raise Exception('Incorrect step 1 processing mode!')
+        
+        blend_alpha = args_dict['step_1_blend_alpha']
+        init_img = warped_styled_frame * (1 - blend_alpha) + curr_frame * blend_alpha
+        img2img_args_dict['init_img'] = Image.fromarray(np.clip(init_img, 0, 255).astype(np.uint8))
+        img2img_args_dict['seed'] = args_dict['step_1_seed']
+        utils.set_CNs_input_image(img2img_args_dict, Image.fromarray(curr_frame))
+        processed_frames, _, _, _ = utils.img2img(img2img_args_dict)
+        processed_frame = np.array(processed_frames[0])[...,:3]
+
+        # normalizing the colors
+        processed_frame = skimage.exposure.match_histograms(processed_frame, curr_frame, channel_axis=None)
+        processed_frame = processed_frame.astype(float) * alpha_mask + warped_styled_frame.astype(float) * (1 - alpha_mask)
+        
+        #processed_frame = processed_frame * 0.94 + curr_frame * 0.06
+        processed_frame = np.clip(processed_frame, 0, 255).astype(np.uint8)
+        sdcn_anim_tmp.prev_frame_styled = processed_frame.copy()
+
+        ### STEP 2
+        if args_dict['fix_frame_strength'] > 0:
+          img2img_args_dict = args_dict #copy.deepcopy(args_dict)
+          img2img_args_dict['mode'] = 0
+          img2img_args_dict['init_img'] = Image.fromarray(processed_frame)
+          img2img_args_dict['mask_img'] = None
+          img2img_args_dict['denoising_strength'] = args_dict['fix_frame_strength']
+          img2img_args_dict['seed'] = args_dict['step_2_seed']
+          utils.set_CNs_input_image(img2img_args_dict, Image.fromarray(curr_frame))
+          processed_frames, _, _, _ = utils.img2img(img2img_args_dict)
+          processed_frame = np.array(processed_frames[0])
+          processed_frame = skimage.exposure.match_histograms(processed_frame, curr_frame, channel_axis=None)
+
+        processed_frame = np.clip(processed_frame, 0, 255).astype(np.uint8)
+        warped_styled_frame_ = np.clip(warped_styled_frame_, 0, 255).astype(np.uint8)
+        
+        # Write the frame to the output video
+        frame_out = np.clip(processed_frame, 0, 255).astype(np.uint8)
+        frame_out = cv2.cvtColor(frame_out, cv2.COLOR_RGB2BGR) 
+        sdcn_anim_tmp.output_video.write(frame_out)
+
+        sdcn_anim_tmp.process_counter += 1
+        #if sdcn_anim_tmp.process_counter >= sdcn_anim_tmp.total_frames - 1:
+        #    sdcn_anim_tmp.input_video.release()
+        #    sdcn_anim_tmp.output_video.release()
+        #    sdcn_anim_tmp.prev_frame = None
+
+        save_result_to_image(processed_frame, sdcn_anim_tmp.process_counter + 1)
+
+    stat = get_cur_stat() + utils.get_time_left(step+2, loop_iterations+1, processing_start_time)
+    yield stat, curr_frame, occlusion_mask, warped_styled_frame_, processed_frame, None, gr.Button.update(interactive=False), gr.Button.update(interactive=True)
+
+  RAFT_clear_memory()
+
+  sdcn_anim_tmp.input_video.release()
+  sdcn_anim_tmp.output_video.release()
+
+  curr_frame = gr.Image.update()
+  occlusion_mask = gr.Image.update()
+  warped_styled_frame_ = gr.Image.update() 
+  processed_frame = gr.Image.update()
+
+  yield get_cur_stat(), curr_frame, occlusion_mask, warped_styled_frame_, processed_frame, output_video_name, gr.Button.update(interactive=True), gr.Button.update(interactive=False)