Skip to content

Commit

Permalink
Base1xModel: Allow custom models to be passed
Browse files Browse the repository at this point in the history
This is intended for simple implementations by users that simply point to a model in a local directory and runs `apply`.

For example:

```py
class ExampleModel(Base1xModel):
    _model_filepath = "_assets/example_model_fp32.onnx"
```
  • Loading branch information
LightArrowsEXE committed Nov 23, 2024
1 parent 2c5ff5c commit 9883c5e
Showing 1 changed file with 17 additions and 3 deletions.
20 changes: 17 additions & 3 deletions lvsfunc/models/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,18 @@ class Base1xModel:
"""

_model_filename = ''
"""Filename of the model."""
"""
Filename of the model. If it contains a slash, it will be treated as a relative path.
For example:
```python
>>> class ExampleModel(Base1xModel):
... _model_filename = "./example_model_fp32.onnx"
...
>>> ExampleModel.apply(clip)
```
"""

def __str__(self) -> str:
return self.__class__.__name__
Expand Down Expand Up @@ -90,7 +101,7 @@ def _set_func(self, clip: vs.VideoNode) -> None:
color_range = ColorRange.from_param_or_video(self._kwargs.pop('color_range', None), clip)

self._matrix = Matrix.from_param_or_video(matrix, clip)
self._planes = normalize_planes(clip, self._kwargs.pop('planes', 0))
self._planes = normalize_planes(clip, self._kwargs.pop('planes', None))

# Pre-resample using the same method I use during training.
proc_clip = self._scale_based_on_planes(clip)
Expand Down Expand Up @@ -157,7 +168,10 @@ def _select_planes(self, processed: vs.VideoNode, ref: vs.VideoNode | None = Non
def _set_model_path(self) -> None:
"""Set the path of the model."""

model_path = get_models_path() / str(self).lower() / self._model_filename
if any(x in self._model_filename for x in ['/', '\\']):
model_path = SPath(self._model_filename)
else:
model_path = get_models_path() / str(self).lower() / self._model_filename

if not model_path.exists():
raise FileWasNotFoundError(
Expand Down

0 comments on commit 9883c5e

Please sign in to comment.