Spaces:
Sleeping
Sleeping
from modelscope.pipelines import pipeline | |
from .tool import Tool | |
class ModelscopePipelineTool(Tool): | |
default_model: str = '' | |
task: str = '' | |
model_revision = None | |
def __init__(self, cfg): | |
super().__init__(cfg) | |
self.model = self.cfg.get('model', None) or self.default_model | |
self.model_revision = self.cfg.get('model_revision', | |
None) or self.model_revision | |
self.pipeline_params = self.cfg.get('pipeline_params', {}) | |
self.pipeline = None | |
self.is_initialized = False | |
def setup(self): | |
# only initialize when this tool is really called to save memory | |
if not self.is_initialized: | |
self.pipeline = pipeline( | |
task=self.task, | |
model=self.model, | |
model_revision=self.model_revision, | |
**self.pipeline_params) | |
self.is_initialized = True | |
def _local_call(self, *args, **kwargs): | |
self.setup() | |
parsed_args, parsed_kwargs = self._local_parse_input(*args, **kwargs) | |
origin_result = self.pipeline(*parsed_args, **parsed_kwargs) | |
final_result = self._parse_output(origin_result, remote=False) | |
return final_result | |