File size: 3,260 Bytes
add8f0b |
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 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
"""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_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.__class__.__name__
if str(type(pipeline).__module__).startswith("transformers.pipelines")
else pipeline.__class__.__name__
)
return interface_info
|