from PIL import Image
import cupy as cp
import numpy as np
from tqdm import tqdm
from ..extensions.FastBlend.patch_match import PyramidPatchMatcher
from ..extensions.FastBlend.runners.fast import TableManager
from .base import VideoProcessor
class FastBlendSmoother(VideoProcessor):
def __init__(
self,
inference_mode="fast", batch_size=8, window_size=60,
minimum_patch_size=5, threads_per_block=8, num_iter=5, gpu_id=0, guide_weight=10.0, initialize="identity", tracking_window_size=0
):
self.inference_mode = inference_mode
self.batch_size = batch_size
self.window_size = window_size
self.ebsynth_config = {
"minimum_patch_size": minimum_patch_size,
"threads_per_block": threads_per_block,
"num_iter": num_iter,
"gpu_id": gpu_id,
"guide_weight": guide_weight,
"initialize": initialize,
"tracking_window_size": tracking_window_size
}
@staticmethod
def from_model_manager(model_manager, **kwargs):
# TODO: fetch GPU ID from model_manager
return FastBlendSmoother(**kwargs)
def inference_fast(self, frames_guide, frames_style):
table_manager = TableManager()
patch_match_engine = PyramidPatchMatcher(
image_height=frames_style[0].shape[0],
image_width=frames_style[0].shape[1],
channel=3,
**self.ebsynth_config
)
# left part
table_l = table_manager.build_remapping_table(frames_guide, frames_style, patch_match_engine, self.batch_size, desc="Fast Mode Step 1/4")
table_l = table_manager.remapping_table_to_blending_table(table_l)
table_l = table_manager.process_window_sum(frames_guide, table_l, patch_match_engine, self.window_size, self.batch_size, desc="Fast Mode Step 2/4")
# right part
table_r = table_manager.build_remapping_table(frames_guide[::-1], frames_style[::-1], patch_match_engine, self.batch_size, desc="Fast Mode Step 3/4")
table_r = table_manager.remapping_table_to_blending_table(table_r)
table_r = table_manager.process_window_sum(frames_guide[::-1], table_r, patch_match_engine, self.window_size, self.batch_size, desc="Fast Mode Step 4/4")[::-1]
# merge
frames = []
for (frame_l, weight_l), frame_m, (frame_r, weight_r) in zip(table_l, frames_style, table_r):
weight_m = -1
weight = weight_l + weight_m + weight_r
frame = frame_l * (weight_l / weight) + frame_m * (weight_m / weight) + frame_r * (weight_r / weight)
frames.append(frame)
frames = [frame.clip(0, 255).astype("uint8") for frame in frames]
frames = [Image.fromarray(frame) for frame in frames]
return frames
def inference_balanced(self, frames_guide, frames_style):
patch_match_engine = PyramidPatchMatcher(
image_height=frames_style[0].shape[0],
image_width=frames_style[0].shape[1],
channel=3,
**self.ebsynth_config
)
output_frames = []
# tasks
n = len(frames_style)
tasks = []
for target in range(n):
for source in range(target - self.window_size, target + self.window_size + 1):
if source >= 0 and source < n and source != target:
tasks.append((source, target))
# run
frames = [(None, 1) for i in range(n)]
for batch_id in tqdm(range(0, len(tasks), self.batch_size), desc="Balanced Mode"):
tasks_batch = tasks[batch_id: min(batch_id+self.batch_size, len(tasks))]
source_guide = np.stack([frames_guide[source] for source, target in tasks_batch])
target_guide = np.stack([frames_guide[target] for source, target in tasks_batch])
source_style = np.stack([frames_style[source] for source, target in tasks_batch])
_, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style)
for (source, target), result in zip(tasks_batch, target_style):
frame, weight = frames[target]
if frame is None:
frame = frames_style[target]
frames[target] = (
frame * (weight / (weight + 1)) + result / (weight + 1),
weight + 1
)
if weight + 1 == min(n, target + self.window_size + 1) - max(0, target - self.window_size):
frame = frame.clip(0, 255).astype("uint8")
output_frames.append(Image.fromarray(frame))
frames[target] = (None, 1)
return output_frames
def inference_accurate(self, frames_guide, frames_style):
patch_match_engine = PyramidPatchMatcher(
image_height=frames_style[0].shape[0],
image_width=frames_style[0].shape[1],
channel=3,
use_mean_target_style=True,
**self.ebsynth_config
)
output_frames = []
# run
n = len(frames_style)
for target in tqdm(range(n), desc="Accurate Mode"):
l, r = max(target - self.window_size, 0), min(target + self.window_size + 1, n)
remapped_frames = []
for i in range(l, r, self.batch_size):
j = min(i + self.batch_size, r)
source_guide = np.stack([frames_guide[source] for source in range(i, j)])
target_guide = np.stack([frames_guide[target]] * (j - i))
source_style = np.stack([frames_style[source] for source in range(i, j)])
_, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style)
remapped_frames.append(target_style)
frame = np.concatenate(remapped_frames, axis=0).mean(axis=0)
frame = frame.clip(0, 255).astype("uint8")
output_frames.append(Image.fromarray(frame))
return output_frames
def release_vram(self):
mempool = cp.get_default_memory_pool()
pinned_mempool = cp.get_default_pinned_memory_pool()
mempool.free_all_blocks()
pinned_mempool.free_all_blocks()
def __call__(self, rendered_frames, original_frames=None, **kwargs):
rendered_frames = [np.array(frame) for frame in rendered_frames]
original_frames = [np.array(frame) for frame in original_frames]
if self.inference_mode == "fast":
output_frames = self.inference_fast(original_frames, rendered_frames)
elif self.inference_mode == "balanced":
output_frames = self.inference_balanced(original_frames, rendered_frames)
elif self.inference_mode == "accurate":
output_frames = self.inference_accurate(original_frames, rendered_frames)
else:
raise ValueError("inference_mode must be fast, balanced or accurate")
self.release_vram()
return output_frames