Переместил interpolator.py внутрь src | добавил пресеты | добавил новые модели

main.py

@@ -4,9 +4,10 @@ from typing import TYPE_CHECKING
 
 import tqdm
 
+from src.config import presets
 from src.utils.fs import FileSystem
 from src.utils.video import VideoMaker
-from interpolator import (
+from src.interpolator import (
     ImageInterpolator,
     Anchor,
     get_device,
@@ -104,7 +105,7 @@ class InterpolationPipeline:
         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
+        fps *= 2  # Doubling FPS
         for frame_paths in self.video_maker.video_to_frames_generator(
             video_path, self.fs.frames_path, chunk_seconds
         ):
@@ -125,7 +126,7 @@ class InterpolationPipeline:
                 part += 1
             for i in tqdm.tqdm(
                 range(len(frame_paths) - 1),
-                desc=f"Processing video frames {part} / {total_parts}",
+                desc=f"Processing video frames {part + 1} / {total_parts}",
             ):
                 img1 = frame_paths[i]
                 img2 = frame_paths[i + 1]
@@ -175,14 +176,16 @@ class InterpolationPipeline:
         )
 
 
-def main():
-    config = Path("src/config/AMT-G.yaml")
-    checkpoint_path = Path("src/pretrained/amt-g.pth")
+def main(preset: presets.Preset = presets.LARGE):
     base_path = Path("output")
     video_path = Path("example/video.mp4")
     output_video = "interpolated_video.mp4"
 
-    pipeline = InterpolationPipeline(config, checkpoint_path, base_path)
+    pipeline = InterpolationPipeline(
+        config=preset.config,
+        checkpoint_path=preset.checkpoint,
+        base_path=base_path,
+    )
     pipeline.run(video_path, output_video)