File size: 1,437 Bytes
8b78124 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 | import os
from pathlib import Path
from huggingface_hub import snapshot_download
from .configuration import PredNetConfig
from .infer.predictor import Predictor
class PredNetModel:
"""Black-box TF SavedModel wrapper compatible with AutoModel via trust_remote_code."""
def __init__(self, config: PredNetConfig, predictor: Predictor):
self.config = config
self._predictor = predictor
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
if os.path.isdir(pretrained_model_name_or_path):
local_dir = pretrained_model_name_or_path
else:
local_dir = snapshot_download(pretrained_model_name_or_path)
config = PredNetConfig.from_pretrained(local_dir)
savedmodel_dir = str(Path(local_dir) / "savedmodel")
predictor = Predictor(
model_dir=savedmodel_dir,
resize_hw=tuple(config.resize_hw),
)
return cls(config, predictor)
def predict_sequence(self, frames, pad_last_frame="none"):
return self._predictor.predict_sequence(frames, pad_last_frame=pad_last_frame)
def predict_last_frame(self, frames, pad_last_frame="none"):
return self._predictor.predict_last_frame(frames, pad_last_frame=pad_last_frame)
def predict_outputs(self, frames, pad_last_frame="none"):
return self._predictor.predict_outputs(frames, pad_last_frame=pad_last_frame)
|