Spaces:
Runtime error
Runtime error
| 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, | |
| } | |
| def from_model_manager(model_manager): | |
| # TODO: fetch GPU ID from 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 | |
| ) | |
| # left part | |
| 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") | |
| # right part | |
| 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] | |
| # 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 = [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 |