| | from .runners.fast import TableManager, PyramidPatchMatcher |
| | from PIL import Image |
| | import numpy as np |
| | import cupy as cp |
| |
|
| |
|
| | class FastBlendSmoother: |
| | def __init__(self): |
| | self.batch_size = 8 |
| | self.window_size = 64 |
| | self.ebsynth_config = { |
| | "minimum_patch_size": 5, |
| | "threads_per_block": 8, |
| | "num_iter": 5, |
| | "gpu_id": 0, |
| | "guide_weight": 10.0, |
| | "initialize": "identity", |
| | "tracking_window_size": 0, |
| | } |
| |
|
| | @staticmethod |
| | def from_model_manager(model_manager): |
| | |
| | return FastBlendSmoother() |
| |
|
| | def run(self, frames_guide, frames_style, batch_size, window_size, ebsynth_config): |
| | frames_guide = [np.array(frame) for frame in frames_guide] |
| | frames_style = [np.array(frame) for frame in 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, |
| | **ebsynth_config |
| | ) |
| | |
| | table_l = table_manager.build_remapping_table(frames_guide, frames_style, patch_match_engine, batch_size, desc="FastBlend 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, window_size, batch_size, desc="FastBlend Step 2/4") |
| | |
| | table_r = table_manager.build_remapping_table(frames_guide[::-1], frames_style[::-1], patch_match_engine, batch_size, desc="FastBlend 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, window_size, batch_size, desc="FastBlend 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 = [Image.fromarray(frame.clip(0, 255).astype("uint8")) for frame in frames] |
| | return frames |
| | |
| | def __call__(self, rendered_frames, original_frames=None, **kwargs): |
| | frames = self.run( |
| | original_frames, rendered_frames, |
| | self.batch_size, self.window_size, self.ebsynth_config |
| | ) |
| | mempool = cp.get_default_memory_pool() |
| | pinned_mempool = cp.get_default_pinned_memory_pool() |
| | mempool.free_all_blocks() |
| | pinned_mempool.free_all_blocks() |
| | return frames |
