hugging-apps/echo-memory

1from PIL import Image2import cupy as cp3import numpy as np4from tqdm import tqdm5from ..extensions.FastBlend.patch_match import PyramidPatchMatcher6from ..extensions.FastBlend.runners.fast import TableManager7from .base import VideoProcessor8 9 10class FastBlendSmoother(VideoProcessor):11    def __init__(12        self,13        inference_mode="fast", batch_size=8, window_size=60,14        minimum_patch_size=5, threads_per_block=8, num_iter=5, gpu_id=0, guide_weight=10.0, initialize="identity", tracking_window_size=015    ):16        self.inference_mode = inference_mode17        self.batch_size = batch_size18        self.window_size = window_size19        self.ebsynth_config = {20            "minimum_patch_size": minimum_patch_size,21            "threads_per_block": threads_per_block,22            "num_iter": num_iter,23            "gpu_id": gpu_id,24            "guide_weight": guide_weight,25            "initialize": initialize,26            "tracking_window_size": tracking_window_size27        }28 29    @staticmethod30    def from_model_manager(model_manager, **kwargs):31        # TODO: fetch GPU ID from model_manager32        return FastBlendSmoother(**kwargs)33 34    def inference_fast(self, frames_guide, frames_style):35        table_manager = TableManager()36        patch_match_engine = PyramidPatchMatcher(37            image_height=frames_style[0].shape[0],38            image_width=frames_style[0].shape[1],39            channel=3,40            **self.ebsynth_config41        )42        # left part43        table_l = table_manager.build_remapping_table(frames_guide, frames_style, patch_match_engine, self.batch_size, desc="Fast Mode Step 1/4")44        table_l = table_manager.remapping_table_to_blending_table(table_l)45        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")46        # right part47        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")48        table_r = table_manager.remapping_table_to_blending_table(table_r)49        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]50        # merge51        frames = []52        for (frame_l, weight_l), frame_m, (frame_r, weight_r) in zip(table_l, frames_style, table_r):53            weight_m = -154            weight = weight_l + weight_m + weight_r55            frame = frame_l * (weight_l / weight) + frame_m * (weight_m / weight) + frame_r * (weight_r / weight)56            frames.append(frame)57        frames = [frame.clip(0, 255).astype("uint8") for frame in frames]58        frames = [Image.fromarray(frame) for frame in frames]59        return frames60    61    def inference_balanced(self, frames_guide, frames_style):62        patch_match_engine = PyramidPatchMatcher(63            image_height=frames_style[0].shape[0],64            image_width=frames_style[0].shape[1],65            channel=3,66            **self.ebsynth_config67        )68        output_frames = []69        # tasks70        n = len(frames_style)71        tasks = []72        for target in range(n):73            for source in range(target - self.window_size, target + self.window_size + 1):74                if source >= 0 and source < n and source != target:75                    tasks.append((source, target))76        # run77        frames = [(None, 1) for i in range(n)]78        for batch_id in tqdm(range(0, len(tasks), self.batch_size), desc="Balanced Mode"):79            tasks_batch = tasks[batch_id: min(batch_id+self.batch_size, len(tasks))]80            source_guide = np.stack([frames_guide[source] for source, target in tasks_batch])81            target_guide = np.stack([frames_guide[target] for source, target in tasks_batch])82            source_style = np.stack([frames_style[source] for source, target in tasks_batch])83            _, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style)84            for (source, target), result in zip(tasks_batch, target_style):85                frame, weight = frames[target]86                if frame is None:87                    frame = frames_style[target]88                frames[target] = (89                    frame * (weight / (weight + 1)) + result / (weight + 1),90                    weight + 191                )92                if weight + 1 == min(n, target + self.window_size + 1) - max(0, target - self.window_size):93                    frame = frame.clip(0, 255).astype("uint8")94                    output_frames.append(Image.fromarray(frame))95                    frames[target] = (None, 1)96        return output_frames97    98    def inference_accurate(self, frames_guide, frames_style):99        patch_match_engine = PyramidPatchMatcher(100            image_height=frames_style[0].shape[0],101            image_width=frames_style[0].shape[1],102            channel=3,103            use_mean_target_style=True,104            **self.ebsynth_config105        )106        output_frames = []107        # run108        n = len(frames_style)109        for target in tqdm(range(n), desc="Accurate Mode"):110            l, r = max(target - self.window_size, 0), min(target + self.window_size + 1, n)111            remapped_frames = []112            for i in range(l, r, self.batch_size):113                j = min(i + self.batch_size, r)114                source_guide = np.stack([frames_guide[source] for source in range(i, j)])115                target_guide = np.stack([frames_guide[target]] * (j - i))116                source_style = np.stack([frames_style[source] for source in range(i, j)])117                _, target_style = patch_match_engine.estimate_nnf(source_guide, target_guide, source_style)118                remapped_frames.append(target_style)119            frame = np.concatenate(remapped_frames, axis=0).mean(axis=0)120            frame = frame.clip(0, 255).astype("uint8")121            output_frames.append(Image.fromarray(frame))122        return output_frames123    124    def release_vram(self):125        mempool = cp.get_default_memory_pool()126        pinned_mempool = cp.get_default_pinned_memory_pool()127        mempool.free_all_blocks()128        pinned_mempool.free_all_blocks()129    130    def __call__(self, rendered_frames, original_frames=None, **kwargs):131        rendered_frames = [np.array(frame) for frame in rendered_frames]132        original_frames = [np.array(frame) for frame in original_frames]133        if self.inference_mode == "fast":134            output_frames = self.inference_fast(original_frames, rendered_frames)135        elif self.inference_mode == "balanced":136            output_frames = self.inference_balanced(original_frames, rendered_frames)137        elif self.inference_mode == "accurate":138            output_frames = self.inference_accurate(original_frames, rendered_frames)139        else:140            raise ValueError("inference_mode must be fast, balanced or accurate")141        self.release_vram()142        return output_frames143