Cool.

nostoslabs · Jan 17, 2024 · 5d2cd01 · 5d2cd01
1 parent e5d1545
commit 5d2cd01
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/video_upscaler/main.py b/src/video_upscaler/main.py
@@ -3,10 +3,10 @@
 from PIL import Image
 import numpy as np
 import av
-from RealESRGAN import RealESRGAN
 from argparse import ArgumentParser
 from pathlib import Path
 from tqdm import tqdm
+from RealESRGAN import RealESRGAN
 
 def upscale_image(image: Image, device: torch.device, weights: Path, scale: int = 4) -> Image:
     """Upscale a single image."""
@@ -23,7 +23,7 @@ def upscale_frame(frame: av.VideoFrame, upscaler: RealESRGAN) -> av.VideoFrame:
 def get_upscaler(device: torch.device, model_weights_path: Path, scale: int = 4) -> RealESRGAN:
     """Get the upscaler model."""
     upscaler = RealESRGAN(device, scale)
-    upscaler.load_weights(str(model_weights_path), download=True)
+    upscaler.load_weights(model_weights_path, download=True)
     return upscaler
 
 def upscale_video(input_path: Path, upscaler: RealESRGAN, scale: int = 4) -> None:

diff --git a/tests/data/heidi.mp4 b/tests/data/heidi.mp4