Spaces:
Running
Running
| """This module should not be used directly as its API is subject to change. Instead, | |
| please use the `gr.Interface.from_pipeline()` function.""" | |
| from __future__ import annotations | |
| from typing import TYPE_CHECKING | |
| from gradio.pipelines_utils import ( | |
| handle_diffusers_pipeline, | |
| handle_transformers_js_pipeline, | |
| handle_transformers_pipeline, | |
| ) | |
| if TYPE_CHECKING: | |
| import diffusers | |
| import transformers | |
| def load_from_pipeline( | |
| pipeline: transformers.Pipeline | diffusers.DiffusionPipeline, # type: ignore | |
| ) -> dict: | |
| """ | |
| Gets the appropriate Interface kwargs for a given Hugging Face transformers.Pipeline or diffusers.DiffusionPipeline. | |
| pipeline (transformers.Pipeline): the transformers.Pipeline from which to create an interface | |
| Returns: | |
| (dict): a dictionary of kwargs that can be used to construct an Interface object | |
| """ | |
| if str(type(pipeline).__module__).startswith("transformers.pipelines."): | |
| pipeline_info = handle_transformers_pipeline(pipeline) | |
| elif str(type(pipeline).__module__).startswith("diffusers.pipelines."): | |
| pipeline_info = handle_diffusers_pipeline(pipeline) | |
| else: | |
| raise ValueError( | |
| "pipeline must be a transformers.pipeline or diffusers.pipeline" | |
| ) | |
| def fn(*params): | |
| if pipeline_info: | |
| data = pipeline_info["preprocess"](*params) | |
| if str(type(pipeline).__module__).startswith("transformers.pipelines"): | |
| from transformers import pipelines | |
| # special cases that needs to be handled differently | |
| if isinstance( | |
| pipeline, | |
| ( | |
| pipelines.text_classification.TextClassificationPipeline, | |
| pipelines.text2text_generation.Text2TextGenerationPipeline, | |
| pipelines.text2text_generation.TranslationPipeline, | |
| ), | |
| ): | |
| data = pipeline(*data) | |
| else: | |
| data = pipeline(**data) # type: ignore | |
| # special case for object-detection | |
| # original input image sent to postprocess function | |
| if isinstance( | |
| pipeline, | |
| pipelines.object_detection.ObjectDetectionPipeline, | |
| ): | |
| output = pipeline_info["postprocess"](data, params[0]) | |
| else: | |
| output = pipeline_info["postprocess"](data) | |
| return output | |
| elif str(type(pipeline).__module__).startswith("diffusers.pipelines"): | |
| data = pipeline(**data) # type: ignore | |
| output = pipeline_info["postprocess"](data) | |
| return output | |
| else: | |
| raise ValueError("pipeline_info can not be None.") | |
| interface_info = pipeline_info.copy() if pipeline_info else {} | |
| interface_info["fn"] = fn | |
| del interface_info["preprocess"] | |
| del interface_info["postprocess"] | |
| # define the title/description of the Interface | |
| interface_info["title"] = ( | |
| pipeline.model.config.name_or_path | |
| if str(type(pipeline).__module__).startswith("transformers.pipelines") | |
| else pipeline.__class__.__name__ | |
| ) | |
| return interface_info | |
| def load_from_js_pipeline(pipeline) -> dict: | |
| if str(type(pipeline).__module__).startswith("transformers_js_py."): | |
| pipeline_info = handle_transformers_js_pipeline(pipeline) | |
| else: | |
| raise ValueError("pipeline must be a transformers_js_py's pipeline") | |
| async def fn(*params): | |
| preprocess = pipeline_info["preprocess"] | |
| postprocess = pipeline_info["postprocess"] | |
| postprocess_takes_inputs = pipeline_info.get("postprocess_takes_inputs", False) | |
| preprocessed_params = preprocess(*params) if preprocess else params | |
| pipeline_output = await pipeline(*preprocessed_params) | |
| postprocessed_output = ( | |
| postprocess(pipeline_output, *(params if postprocess_takes_inputs else ())) | |
| if postprocess | |
| else pipeline_output | |
| ) | |
| return postprocessed_output | |
| interface_info = { | |
| "fn": fn, | |
| "inputs": pipeline_info["inputs"], | |
| "outputs": pipeline_info["outputs"], | |
| "title": f"{pipeline.task} ({pipeline.model.config._name_or_path})", | |
| } | |
| return interface_info | |