| 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) |
|
|
