| 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): |
| |
| 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 |
| ) |
| |
| 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") |
| |
| 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] |
| |
| 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 = [] |
| |
| 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)) |
| |
| 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 = [] |
| |
| 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 |
|
|
