Skip to content

Commit

Permalink
Add fp16 filename check
Browse files Browse the repository at this point in the history
  • Loading branch information
LightArrowsEXE committed Sep 21, 2024
1 parent 033de5f commit b749a18
Show file tree
Hide file tree
Showing 2 changed files with 20 additions and 5 deletions.
10 changes: 9 additions & 1 deletion lvsfunc/models/base.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from warnings import warn
import importlib.resources as pkg_resources

from stgpytools import FileWasNotFoundError, SPath
Expand All @@ -24,7 +25,14 @@ def _get_model_path(sub_dir: str, model_name: str, fp16: bool = True) -> SPath:
if not fp16:
return model_path

return model_path.with_name(model_path.stem.split('_fp32')[0] + '_fp16.onnx')
new_path = model_path.with_name(model_path.stem.split('_fp32')[0] + '_fp16.onnx')

if new_path.exists():
return new_path

warn(f'{model_name}: Could not find fp16 model! Returning fp32 model instead.')

return model_path


# Paths to every model. Should always end with _fp32.onnx, we swap it out later if necessary.
Expand Down
15 changes: 11 additions & 4 deletions lvsfunc/models/delowpass.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,11 @@ def apply(self, clip: vs.VideoNode, **kwargs: Any) -> vs.VideoNode:
"""
Apply the delowpass model to the clip.
WARNING: These models are only for horizontal lowpassing!
They will also mess up any random high-frequency content in the image.
Practically speaking, this almost always means compression noise.
Make sure you properly preprocess your video before using these models!
Example usage:
.. code-block:: python
Expand All @@ -73,15 +78,17 @@ def apply(self, clip: vs.VideoNode, **kwargs: Any) -> vs.VideoNode:
except ImportError:
raise DependencyNotFoundError("vsmlrt", self._func)

func = FunctionUtil(clip, self.__class__, kwargs.pop('planes', 0), vs.RGB, 32)
fp16 = kwargs.pop('fp16', False)

model_path = _get_model_path('delowpass', str(self).split('(')[0], fp16)
fp16 = 'fp16' in model_path.stem

func = FunctionUtil(clip, self.__class__, kwargs.pop('planes', 0), vs.RGB, 16 if fp16 else 32)

fp16 = kwargs.pop('fp16', func.bitdepth != 32)
matrix = Matrix.from_param_or_video(kwargs.pop('matrix', None), clip)

work_clip = Lanczos.resample(func.work_clip, vs.RGBS, matrix_in=matrix)

model_path = _get_model_path('delowpass', str(self).split('(')[0], fp16)

if self.backend is None:
self.backend = autoselect_backend(fp16=fp16)

Expand Down

0 comments on commit b749a18

Please sign in to comment.