Перевел импорты модулей в относительные пути

main.py

@@ -1,199 +1,7 @@
 import logging
 from pathlib import Path
-from typing import TYPE_CHECKING
-
-from cv2 import imwrite
-import tqdm
-
+from src.runner import runner
 from src.config import presets
-from src.utils.fs import FileSystem
-from src.utils.video import VideoMaker
-from src.interpolator import (
-    ImageInterpolator,
-    Anchor,
-    get_device,
-    get_vram_available,
-    ModelRunner,
-)
-
-
-if TYPE_CHECKING:
-    import torch
-    import numpy as np
-
-
-def performing_warning_message(device: "torch.device"):
-    if device.type in ("cpu", "mps"):
-        if device.type == "mps":
-            logging.warning(
-                "Running on Apple Silicon GPU (MPS) may have limited performance. Consider using a CUDA-enabled GPU for better performance."
-            )
-        else:
-            logging.warning(
-                "Running on CPU may be very slow. Consider using a GPU for better performance."
-            )
-    elif device.type == "cuda":
-        pass
-    else:
-        raise Exception(f"Unsupported device type: {device.type}")
-
-
-def init_fs(base_path: Path) -> FileSystem:
-    fs = FileSystem(base_path)
-    fs.clear_directory(fs.frames_path)
-    fs.clear_directory(fs.interpolated_path)
-    fs.clear_directory(fs.moved_path)
-    fs.clear_directory(fs.video_part_path)
-    return fs
-
-
-def init_video_maker() -> VideoMaker:
-    return VideoMaker()
-
-
-def init_device() -> "torch.device":
-    device = get_device()
-    performing_warning_message(device)
-    vram_available = get_vram_available(device)
-    logging.info(f"Available VRAM: {vram_available / (1024**3):.2f} GB")
-    return device
-
-
-def init_anchor(device: "torch.device") -> Anchor:
-    if device.type in ("cpu", "mps"):
-        return Anchor(resolution=8192 * 8192, memory=1, memory_bias=0)
-    elif device.type == "cuda":
-        return Anchor(
-            resolution=1024 * 512, memory=1500 * 1024**2, memory_bias=2500 * 1024**2
-        )
-    else:
-        raise Exception(f"Unsupported device type: {device.type}")
-
-
-def init_model_runner(
-    config: Path, checkpoint_path: Path, device: "torch.device"
-) -> ModelRunner:
-    return ModelRunner(config, checkpoint_path, device)
-
-
-def init_interpolator(
-    model_runner: ModelRunner, device: "torch.device"
-) -> ImageInterpolator:
-    anchor = init_anchor(device)
-    return ImageInterpolator(device, anchor, model_runner)
-
-
-class InterpolationPipeline:
-    def __init__(
-        self,
-        config: Path,
-        checkpoint_path: Path,
-        base_path: Path,
-    ):
-        self.fs = init_fs(base_path)
-        self.video_maker = init_video_maker()
-        self.device = init_device()
-        self.model_runner = init_model_runner(config, checkpoint_path, self.device)
-        self.interpolator = init_interpolator(self.model_runner, self.device)
-
-    def run(self, video_path: Path, output_video: str):
-        prev_frames = tuple()
-        interpolated_frames: list["np.ndarray"] = []
-        part = 0
-        chunk_seconds = 10
-        length = self.video_maker.get_video_duration(video_path)
-        last_part_seconds = 1 if length % chunk_seconds else 0
-        total_parts = int(length // chunk_seconds) + last_part_seconds
-        fps = self.video_maker.get_fps(video_path)
-        logging.info(f"Video FPS: {fps}")
-        fps *= 2  # Doubling FPS
-        width, height = self.video_maker.get_size(video_path)
-        for frames in self.video_maker.video_to_frames_generator(
-            video_path, self.fs.frames_path, chunk_seconds
-        ):
-            logging.info(f"Processing frames: {len(frames)}")
-            if prev_frames:
-                img1 = prev_frames[-1]
-                img2 = frames[0]
-                img1_2 = self.interpolator.interpolate(img1, img2)
-                interpolated_frames.append(img1_2)
-                generator = self._frame_generator(prev_frames, interpolated_frames)
-                part_path = self.fs.video_part_path / f"video_{part:08d}.mp4"
-                self.video_maker.images_to_video_pipeline(
-                    generator, part_path, width, height, fps
-                )
-                interpolated_frames = []
-                logging.info(f"Finished processing part {part:08d}")
-                part += 1
-            for i in tqdm.tqdm(
-                range(len(frames) - 1),
-                desc=f"Processing video frames {part + 1} / {total_parts}",
-            ):
-                img1 = frames[i]
-                img2 = frames[i + 1]
-                img1_2 = self.interpolator.interpolate(img1, img2)
-                interpolated_frames.append(img1_2)
-            prev_frames = frames
-
-        generator = self._frame_generator(prev_frames, interpolated_frames)
-        part_path = self.fs.video_part_path / f"video_{part:08d}.mp4"
-        self.video_maker.images_to_video_pipeline(
-            generator, part_path, width, height, fps
-        )
-        logging.info(f"Finished processing part {part:08d}")
-        self._merge_video_parts(self.fs.output_path / output_video)
-        logging.info(
-            f"Video interpolation completed. Output saved to: {self.fs.output_path / output_video}"
-        )
-
-    def _save_images(
-        self,
-        source: tuple["np.ndarray", ...],
-        interpolated: list["np.ndarray"],
-    ):
-        logging.info("Saving images...")
-        self.fs.clear_directory(self.fs.moved_path)
-        index = 0
-        for i, frame in enumerate(source):
-            name = self.fs.moved_path / f"img_{index:08d}.png"
-            index += 1
-            imwrite(name, frame)
-            if i < len(interpolated):
-                name = self.fs.moved_path / f"img_{index:08d}.png"
-                index += 1
-                imwrite(name, interpolated[i])
-        logging.info("Success...")
-
-    def _merge_frames_to_video(self, output_video: Path, fps: float):
-        self.video_maker.images_to_video(self.fs.moved_path, output_video, fps)
-
-    def _merge_video_parts(self, output_video: Path):
-        self.video_maker.concatenate_videos(self.fs.video_part_path, output_video)
-        self.fs.clear_directory(self.fs.video_part_path)
-
-    def _frame_generator(
-        self,
-        source: tuple["np.ndarray", ...],
-        interpolated: list["np.ndarray"],
-    ):
-        for i, frame in enumerate(source):
-            yield frame
-            if i < len(interpolated):
-                yield interpolated[i]
-
-
-def runner(
-    base_path: Path,
-    video_path: Path,
-    output_video: str,
-    preset: presets.Preset = presets.LARGE,
-):
-    pipeline = InterpolationPipeline(
-        config=preset.config,
-        checkpoint_path=preset.checkpoint,
-        base_path=base_path,
-    )
-    pipeline.run(video_path, output_video)
 
 
 def main():