Improve support

.github/workflows/check-code.yml

@@ -50,4 +50,4 @@ jobs:
       - name: Install test dependencies
         run: python -m pip install -r dev-requirements.txt
       - name: Test
-        run: pytest tests
+        run: pytest tests --full-trace

.vscode/settings.json

@@ -1,13 +1,6 @@
 {
-    "python.linting.enabled": true,
-    "python.linting.flake8Enabled": true,
-    "python.linting.flake8Args": [],
-    "python.linting.mypyEnabled": true,
-    "python.linting.mypyArgs": ["--follow-imports=silent"],
-    "python.linting.pylintEnabled": true,
-    "python.linting.pylintArgs": ["--load-plugins", "pylint_pytest"],
-    "python.testing.pytestEnabled": true,
-    "files.associations": {
-        "*.yaml": "home-assistant"
-    }
+    "flake8.args": [],
+    "mypy-type-checker.args": ["--follow-imports=silent"],
+    "pylint.args": ["--load-plugins", "pylint_pytest"],
+    "python.testing.pytestEnabled": true
 }

buganime/buganime.py

@@ -44,9 +44,12 @@ class Movie:
 
 def parse_streams(streams: Any) -> transcode.VideoInfo:
     def _get_video_stream() -> Any:
-        for stream in streams:
+        video_streams = [stream for stream in streams if stream['codec_type'] == 'video']
+        if len(video_streams) == 1:
+            return video_streams[0]
+        for stream in video_streams:
             match stream:
-                case {'codec_type': 'video', 'disposition': {'default': 1}}:
+                case {'disposition': {'default': 1}}:
                     return stream
         raise RuntimeError('No default video stream found')
 
@@ -79,15 +82,16 @@ def _get_subtitle_stream_index() -> int:
     video = _get_video_stream()
     return transcode.VideoInfo(audio_index=_get_audio_stream()['index'], subtitle_index=_get_subtitle_stream_index(),
                                width=video['width'], height=video['height'], fps=video['r_frame_rate'],
-                               frames=int(video['tags'].get('NUMBER_OF_FRAMES') or video['tags'].get('NUMBER_OF_FRAMES-eng') or 0))
+                               frames=int(video.get('tags', {}).get('NUMBER_OF_FRAMES') or video.get('tags', {}).get('NUMBER_OF_FRAMES-eng') or 0))
 
 
 def parse_filename(input_path: str) -> TVShow | Movie:
     # Remove metadata in brackets/parentheses and extension (e.g. hash, resolution, etc.)
-    input_path = input_path.replace('_', ' ')
+    input_path = re.sub(r'[_+]', ' ', input_path)
     input_path = re.sub(r'\[[^\]]*\]', '', input_path)
     input_path = re.sub(r'\([^\)]*\)', '', input_path)
     input_path = re.sub(r'\d{3,4}p[ -][^\\]*', '', input_path)
+    input_path = re.sub(r'[ -]*\\[ -]*', r'\\', input_path)
     input_path = os.path.splitext(input_path)[0].strip(' -')
 
     # Remove extension and directories
@@ -103,13 +107,13 @@ def parse_filename(input_path: str) -> TVShow | Movie:
         return TVShow(name=match.group('name'), season=int(match.group('season')), episode=int(match.group('episode')))
 
     # Other standalone TV Shows
-    if match := re.match(r'^(?P<name>.+?)[ -]+(?:S(?:eason ?)?(?P<season>\d{1,2})[ -]*)?E?(?P<episode>\d{1,3})(?:v\d+)?$', input_name):
+    if match := re.match(r'^(?P<name>.+?)[ -]* (?:S(?:eason ?)?(?P<season>\d{1,2})[ -]*)?E?(?P<episode>\d{1,3})(?:v\d+)?(?:[ -].*)?$', input_name):
         return TVShow(name=match.group('name'), season=int(match.group('season') or '1'), episode=int(match.group('episode')))
 
     # Structured TV Shows
-    dir_re = r'(?P<name>[^\\]+?)[ -]+S(?:eason ?)?(?P<season>\d{1,2})[ -][^\\]*'
-    file_re = r'[^\\]*S\d{1,2}?E(?P<episode>\d{1,3})(?:[ -][^\\]*)?'
-    if match := re.match(fr'^.*\\{dir_re}\\{file_re}$', input_path):
+    dir_re = r'(?P<name>[^\\]+?)[ -]+S(?:eason ?)?\d{1,2}(?:[ -][^\\]*)?'
+    file_re = r'[^\\]*S(?P<season>\d{1,2})E(?P<episode>\d{1,3})(?:[ -][^\\]*)?'
+    if match := re.match(fr'^.*\\{dir_re}(?:\\.*)?\\{file_re}$', input_path):
         return TVShow(name=match.group('name'), season=int(match.group('season')), episode=int(match.group('episode')))
 
     return Movie(name=input_name)
@@ -137,20 +141,20 @@ def process_file(input_path: str) -> None:
     logging.info('ffprobe %s wrote %s, %s', str(proc.args), proc.stderr, proc.stdout)
     video_info = parse_streams(json.loads(proc.stdout)['streams'])
 
-    try:
-        with lock_mutex(name=UPSCALE_MUTEX_NAME):
-            logging.info('Running Upscaler')
-            asyncio.run(transcode.Transcoder(input_path=input_path, output_path=output_path, height_out=2160, video_info=video_info).run())
-            logging.info('Upscaler for %s finished', input_path)
-    except Exception:
-        logging.exception('Failed to convert %s', input_path)
+    with lock_mutex(name=UPSCALE_MUTEX_NAME):
+        logging.info('Running Upscaler')
+        asyncio.run(transcode.Transcoder(input_path=input_path, output_path=output_path, height_out=2160, width_out=3840, video_info=video_info).run())
+        logging.info('Upscaler for %s finished', input_path)
 
 
 def process_path(input_path: str) -> None:
     if os.path.isdir(input_path):
         for root, _, files in os.walk(input_path):
             for file in files:
-                process_file(input_path=os.path.join(root, file))
+                try:
+                    process_file(input_path=os.path.join(root, file))
+                except Exception:
+                    logging.exception('Failed to convert %s', input_path)
     else:
         process_file(input_path=input_path)
 

buganime/transcode.py

@@ -3,6 +3,8 @@
 import tempfile
 import asyncio
 import logging
+import warnings
+import shutil
 from dataclasses import dataclass
 from typing import AsyncIterator, cast, Optional
 
@@ -44,17 +46,21 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
                 tensor = body(tensor)
             return cast(torch.Tensor, self.__upsampler(tensor) + base)
 
-    def __init__(self, input_path: str, output_path: str, height_out: int, video_info: VideoInfo) -> None:
+    def __init__(self, input_path: str, output_path: str, height_out: int, width_out: int, video_info: VideoInfo) -> None:
         if not os.path.isfile(MODEL_PATH):
             with open(MODEL_PATH, 'wb') as file:
                 file.write(requests.get(MODEL_URL, timeout=600).content)
         self.__input_path, self.__output_path = input_path, output_path
         self.__video_info = video_info
         self.__height_out = height_out
-        self.__width_out = round(self.__video_info.width * self.__height_out / self.__video_info.height)
+        self.__width_out = width_out
         model = Transcoder.Module(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4)
-        model.load_state_dict(torch.load(MODEL_PATH)['params'], strict=True)
-        self.__model = model.eval().cuda().half()
+        if torch.cuda.is_available():
+            model.load_state_dict(torch.load(MODEL_PATH)['params'], strict=True)
+            self.__model = model.eval().cuda().half()
+        else:
+            model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu'))['params'], strict=True)
+            self.__model = model.eval()
         self.__gpu_lock: Optional[asyncio.Lock] = None
         self.__frame_tasks_queue: Optional[asyncio.Queue[Optional[asyncio.Task[bytes]]]] = None
 
@@ -78,10 +84,20 @@ async def __read_input_frames(self) -> AsyncIterator[bytes]:
 
     async def __write_output_frames(self, frames: AsyncIterator[bytes]) -> None:
         with tempfile.TemporaryDirectory() as temp_dir:
-            os.link(self.__input_path, os.path.join(temp_dir, 'input.mkv'))
-            args = ('-f', 'rawvideo', '-framerate', str(self.__video_info.fps), '-pix_fmt', 'rgb24', '-s', f'{self.__width_out}x{self.__height_out}',
+            if os.path.splitdrive(self.__input_path)[0] == os.path.splitdrive(temp_dir)[0]:
+                os.link(self.__input_path, os.path.join(temp_dir, 'input.mkv'))
+            else:
+                shutil.copy(self.__input_path, os.path.join(temp_dir, 'input.mkv'))
+            width_out = self.__width_out
+            height_out = self.__height_out
+            if self.__video_info.width / self.__video_info.height > self.__width_out / self.__height_out:
+                height_out = round(self.__video_info.height * self.__width_out / self.__video_info.width)
+            else:
+                width_out = round(self.__video_info.width * self.__height_out / self.__video_info.height)
+            args = ('-f', 'rawvideo', '-framerate', str(self.__video_info.fps), '-pix_fmt', 'rgb24', '-s', f'{width_out}x{height_out}',
                     '-i', 'pipe:', '-i', 'input.mkv',
-                    '-map', '0', '-map', f'1:{self.__video_info.audio_index}', '-vf', f'subtitles=input.mkv:si={self.__video_info.subtitle_index}',
+                    '-map', '0', '-map', f'1:{self.__video_info.audio_index}',
+                    '-vf', f'subtitles=input.mkv:si={self.__video_info.subtitle_index}, pad={self.__width_out}:{self.__height_out}:(ow-iw)/2:(oh-ih)/2:black',
                     *FFMPEG_OUTPUT_ARGS, self.__output_path,
                     '-loglevel', 'warning', '-y')
             proc = await asyncio.subprocess.create_subprocess_exec('ffmpeg', *args, stdin=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
@@ -103,15 +119,19 @@ async def __write_output_frames(self, frames: AsyncIterator[bytes]) -> None:
     @retry.retry(RuntimeError, tries=10, delay=1)
     def __gpu_upscale(self, frame: torch.Tensor) -> torch.Tensor:
         with torch.no_grad():
-            frame_float = frame.cuda().permute(2, 0, 1).half() / 255
+            if torch.cuda.is_available():
+                frame_float = frame.cuda().permute(2, 0, 1).half() / 255
+            else:
+                frame_float = frame.permute(2, 0, 1) / 255
             frame_upscaled_float = self.__model(frame_float.unsqueeze(0)).data.squeeze().clamp_(0, 1)
             return cast(torch.Tensor, (frame_upscaled_float * 255.0).round().byte().permute(1, 2, 0).cpu())
 
     async def __upscale_frame(self, frame: bytes) -> bytes:
         if self.__video_info.height == self.__height_out:
             return frame
         with torch.no_grad():
-            frame_arr = torch.frombuffer(frame, dtype=torch.uint8).reshape([self.__video_info.height, self.__video_info.width, 3])
+            with warnings.catch_warnings(action='ignore'):
+                frame_arr = torch.frombuffer(frame, dtype=torch.uint8).reshape([self.__video_info.height, self.__video_info.width, 3])
         assert self.__gpu_lock
         async with self.__gpu_lock:
             frame_cpu = await asyncio.to_thread(self.__gpu_upscale, frame_arr)

dev-requirements.txt

@@ -7,4 +7,5 @@ mypy
 tqdm-stubs
 types-retry
 types-requests
+torch<2.3.0  # Due to https://github.com/pytorch/pytorch/issues/124897
 -r requirements.txt