Blane187/multi-diffusion
3
1"""This module should not be used directly as its API is subject to change. Instead,2use the `gr.Blocks.load()` or `gr.load()` functions."""3 4from __future__ import annotations5 6import json7import os8import re9import tempfile10import warnings11from pathlib import Path12from typing import TYPE_CHECKING, Callable13 14import httpx15import huggingface_hub16from gradio_client import Client17from gradio_client.client import Endpoint18from gradio_client.documentation import document19from packaging import version20 21import gradio22from gradio import components, external_utils, utils23from gradio.context import Context24from gradio.exceptions import (25 GradioVersionIncompatibleError,26 ModelNotFoundError,27 TooManyRequestsError,28)29from gradio.processing_utils import save_base64_to_cache, to_binary30 31if TYPE_CHECKING:32 from gradio.blocks import Blocks33 from gradio.interface import Interface34 35 36server_timeout = 60037 38 39@document()40def load(41 name: str,42 src: str | None = None,43 hf_token: str | None = None,44 alias: str | None = None,45 **kwargs,46) -> Blocks:47 """48 Constructs a demo from a Hugging Face repo. Can accept model repos (if src is "models") or Space repos (if src is "spaces"). The input49 and output components are automatically loaded from the repo. Note that if a Space is loaded, certain high-level attributes of the Blocks (e.g.50 custom `css`, `js`, and `head` attributes) will not be loaded.51 Parameters:52 name: the name of the model (e.g. "gpt2" or "facebook/bart-base") or space (e.g. "flax-community/spanish-gpt2"), can include the `src` as prefix (e.g. "models/facebook/bart-base")53 src: the source of the model: `models` or `spaces` (or leave empty if source is provided as a prefix in `name`)54 hf_token: optional access token for loading private Hugging Face Hub models or spaces. Find your token here: https://huggingface.co/settings/tokens. Warning: only provide this if you are loading a trusted private Space as it can be read by the Space you are loading.55 alias: optional string used as the name of the loaded model instead of the default name (only applies if loading a Space running Gradio 2.x)56 Returns:57 a Gradio Blocks object for the given model58 Example:59 import gradio as gr60 demo = gr.load("gradio/question-answering", src="spaces")61 demo.launch()62 """63 return load_blocks_from_repo(64 name=name, src=src, hf_token=hf_token, alias=alias, **kwargs65 )66 67 68def load_blocks_from_repo(69 name: str,70 src: str | None = None,71 hf_token: str | None = None,72 alias: str | None = None,73 **kwargs,74) -> Blocks:75 """Creates and returns a Blocks instance from a Hugging Face model or Space repo."""76 if src is None:77 # Separate the repo type (e.g. "model") from repo name (e.g. "google/vit-base-patch16-224")78 tokens = name.split("/")79 if len(tokens) <= 1:80 raise ValueError(81 "Either `src` parameter must be provided, or `name` must be formatted as {src}/{repo name}"82 )83 src = tokens[0]84 name = "/".join(tokens[1:])85 86 factory_methods: dict[str, Callable] = {87 # for each repo type, we have a method that returns the Interface given the model name & optionally an hf_token88 "huggingface": from_model,89 "models": from_model,90 "spaces": from_spaces,91 }92 if src.lower() not in factory_methods:93 raise ValueError(f"parameter: src must be one of {factory_methods.keys()}")94 95 if hf_token is not None:96 if Context.hf_token is not None and Context.hf_token != hf_token:97 warnings.warn(98 """You are loading a model/Space with a different access token than the one you used to load a previous model/Space. This is not recommended, as it may cause unexpected behavior."""99 )100 Context.hf_token = hf_token101 102 blocks: gradio.Blocks = factory_methods[src](name, hf_token, alias, **kwargs)103 return blocks104 105 106def from_model(model_name: str, hf_token: str | None, alias: str | None, **kwargs):107 model_url = f"https://huggingface.co/{model_name}"108 api_url = f"https://api-inference.huggingface.co/models/{model_name}"109 print(f"Fetching model from: {model_url}")110 111 headers = {"Authorization": f"Bearer {hf_token}"} if hf_token is not None else {}112 response = httpx.request("GET", api_url, headers=headers)113 if response.status_code != 200:114 raise ModelNotFoundError(115 f"Could not find model: {model_name}. If it is a private or gated model, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."116 )117 p = response.json().get("pipeline_tag")118 119 headers["X-Wait-For-Model"] = "true"120 client = huggingface_hub.InferenceClient(121 model=model_name, headers=headers, token=hf_token, timeout=server_timeout,122 )123 124 # For tasks that are not yet supported by the InferenceClient125 GRADIO_CACHE = os.environ.get("GRADIO_TEMP_DIR") or str( # noqa: N806126 Path(tempfile.gettempdir()) / "gradio"127 )128 129 def custom_post_binary(data):130 data = to_binary({"path": data})131 response = httpx.request("POST", api_url, headers=headers, content=data)132 return save_base64_to_cache(133 external_utils.encode_to_base64(response), cache_dir=GRADIO_CACHE134 )135 136 preprocess = None137 postprocess = None138 examples = None139 140 # example model: ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition141 if p == "audio-classification":142 inputs = components.Audio(type="filepath", label="Input")143 outputs = components.Label(label="Class")144 postprocess = external_utils.postprocess_label145 examples = [146 "https://gradio-builds.s3.amazonaws.com/demo-files/audio_sample.wav"147 ]148 fn = client.audio_classification149 # example model: facebook/xm_transformer_sm_all-en150 elif p == "audio-to-audio":151 inputs = components.Audio(type="filepath", label="Input")152 outputs = components.Audio(label="Output")153 examples = [154 "https://gradio-builds.s3.amazonaws.com/demo-files/audio_sample.wav"155 ]156 fn = custom_post_binary157 # example model: facebook/wav2vec2-base-960h158 elif p == "automatic-speech-recognition":159 inputs = components.Audio(type="filepath", label="Input")160 outputs = components.Textbox(label="Output")161 examples = [162 "https://gradio-builds.s3.amazonaws.com/demo-files/audio_sample.wav"163 ]164 fn = client.automatic_speech_recognition165 # example model: microsoft/DialoGPT-medium166 elif p == "conversational":167 inputs = [168 components.Textbox(render=False),169 components.State(render=False),170 ]171 outputs = [172 components.Chatbot(render=False),173 components.State(render=False),174 ]175 examples = [["Hello World"]]176 preprocess = external_utils.chatbot_preprocess177 postprocess = external_utils.chatbot_postprocess178 fn = client.conversational179 # example model: julien-c/distilbert-feature-extraction180 elif p == "feature-extraction":181 inputs = components.Textbox(label="Input")182 outputs = components.Dataframe(label="Output")183 fn = client.feature_extraction184 postprocess = utils.resolve_singleton185 # example model: distilbert/distilbert-base-uncased186 elif p == "fill-mask":187 inputs = components.Textbox(label="Input")188 outputs = components.Label(label="Classification")189 examples = [190 "Hugging Face is the AI community, working together, to [MASK] the future."191 ]192 postprocess = external_utils.postprocess_mask_tokens193 fn = client.fill_mask194 # Example: google/vit-base-patch16-224195 elif p == "image-classification":196 inputs = components.Image(type="filepath", label="Input Image")197 outputs = components.Label(label="Classification")198 postprocess = external_utils.postprocess_label199 examples = ["https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg"]200 fn = client.image_classification201 # Example: deepset/xlm-roberta-base-squad2202 elif p == "question-answering":203 inputs = [204 components.Textbox(label="Question"),205 components.Textbox(lines=7, label="Context"),206 ]207 outputs = [208 components.Textbox(label="Answer"),209 components.Label(label="Score"),210 ]211 examples = [212 [213 "What entity was responsible for the Apollo program?",214 "The Apollo program, also known as Project Apollo, was the third United States human spaceflight"215 " program carried out by the National Aeronautics and Space Administration (NASA), which accomplished"216 " landing the first humans on the Moon from 1969 to 1972.",217 ]218 ]219 postprocess = external_utils.postprocess_question_answering220 fn = client.question_answering221 # Example: facebook/bart-large-cnn222 elif p == "summarization":223 inputs = components.Textbox(label="Input")224 outputs = components.Textbox(label="Summary")225 examples = [226 [227 "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct."228 ]229 ]230 fn = client.summarization231 # Example: distilbert-base-uncased-finetuned-sst-2-english232 elif p == "text-classification":233 inputs = components.Textbox(label="Input")234 outputs = components.Label(label="Classification")235 examples = ["I feel great"]236 postprocess = external_utils.postprocess_label237 fn = client.text_classification238 # Example: gpt2239 elif p == "text-generation":240 inputs = components.Textbox(label="Text")241 outputs = inputs242 examples = ["Once upon a time"]243 fn = external_utils.text_generation_wrapper(client)244 # Example: valhalla/t5-small-qa-qg-hl245 elif p == "text2text-generation":246 inputs = components.Textbox(label="Input")247 outputs = components.Textbox(label="Generated Text")248 examples = ["Translate English to Arabic: How are you?"]249 fn = client.text_generation250 # Example: Helsinki-NLP/opus-mt-en-ar251 elif p == "translation":252 inputs = components.Textbox(label="Input")253 outputs = components.Textbox(label="Translation")254 examples = ["Hello, how are you?"]255 fn = client.translation256 # Example: facebook/bart-large-mnli257 elif p == "zero-shot-classification":258 inputs = [259 components.Textbox(label="Input"),260 components.Textbox(label="Possible class names (" "comma-separated)"),261 components.Checkbox(label="Allow multiple true classes"),262 ]263 outputs = components.Label(label="Classification")264 postprocess = external_utils.postprocess_label265 examples = [["I feel great", "happy, sad", False]]266 fn = external_utils.zero_shot_classification_wrapper(client)267 # Example: sentence-transformers/distilbert-base-nli-stsb-mean-tokens268 elif p == "sentence-similarity":269 inputs = [270 components.Textbox(271 label="Source Sentence",272 placeholder="Enter an original sentence",273 ),274 components.Textbox(275 lines=7,276 placeholder="Sentences to compare to -- separate each sentence by a newline",277 label="Sentences to compare to",278 ),279 ]280 outputs = components.JSON(label="Similarity scores")281 examples = [["That is a happy person", "That person is very happy"]]282 fn = external_utils.sentence_similarity_wrapper(client)283 # Example: julien-c/ljspeech_tts_train_tacotron2_raw_phn_tacotron_g2p_en_no_space_train284 elif p == "text-to-speech":285 inputs = components.Textbox(label="Input")286 outputs = components.Audio(label="Audio")287 examples = ["Hello, how are you?"]288 fn = client.text_to_speech289 # example model: osanseviero/BigGAN-deep-128290 elif p == "text-to-image":291 inputs = components.Textbox(label="Input")292 outputs = components.Image(label="Output")293 examples = ["A beautiful sunset"]294 fn = client.text_to_image295 # example model: huggingface-course/bert-finetuned-ner296 elif p == "token-classification":297 inputs = components.Textbox(label="Input")298 outputs = components.HighlightedText(label="Output")299 examples = [300 "Hugging Face is a company based in Paris and New York City that acquired Gradio in 2021."301 ]302 fn = external_utils.token_classification_wrapper(client)303 # example model: impira/layoutlm-document-qa304 elif p == "document-question-answering":305 inputs = [306 components.Image(type="filepath", label="Input Document"),307 components.Textbox(label="Question"),308 ]309 postprocess = external_utils.postprocess_label310 outputs = components.Label(label="Label")311 fn = client.document_question_answering312 # example model: dandelin/vilt-b32-finetuned-vqa313 elif p == "visual-question-answering":314 inputs = [315 components.Image(type="filepath", label="Input Image"),316 components.Textbox(label="Question"),317 ]318 outputs = components.Label(label="Label")319 postprocess = external_utils.postprocess_visual_question_answering320 examples = [321 [322 "https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg",323 "What animal is in the image?",324 ]325 ]326 fn = client.visual_question_answering327 # example model: Salesforce/blip-image-captioning-base328 elif p == "image-to-text":329 inputs = components.Image(type="filepath", label="Input Image")330 outputs = components.Textbox(label="Generated Text")331 examples = ["https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg"]332 fn = client.image_to_text333 # example model: rajistics/autotrain-Adult-934630783334 elif p in ["tabular-classification", "tabular-regression"]:335 examples = external_utils.get_tabular_examples(model_name)336 col_names, examples = external_utils.cols_to_rows(examples) # type: ignore337 examples = [[examples]] if examples else None338 inputs = components.Dataframe(339 label="Input Rows",340 type="pandas",341 headers=col_names,342 col_count=(len(col_names), "fixed"),343 render=False,344 )345 outputs = components.Dataframe(346 label="Predictions", type="array", headers=["prediction"]347 )348 fn = external_utils.tabular_wrapper349 # example model: microsoft/table-transformer-detection350 elif p == "object-detection":351 inputs = components.Image(type="filepath", label="Input Image")352 outputs = components.AnnotatedImage(label="Annotations")353 fn = external_utils.object_detection_wrapper(client)354 # example model: stabilityai/stable-diffusion-xl-refiner-1.0355 elif p == "image-to-image":356 inputs = [357 components.Image(type="filepath", label="Input Image"),358 components.Textbox(label="Input"),359 ]360 outputs = components.Image(label="Output")361 examples = [362 [363 "https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg",364 "Photo of a cheetah with green eyes",365 ]366 ]367 fn = client.image_to_image368 else:369 raise ValueError(f"Unsupported pipeline type: {p}")370 371 def query_huggingface_inference_endpoints(*data, **kwargs):372 if preprocess is not None:373 data = preprocess(*data)374 data = fn(*data, **kwargs) # type: ignore375 if postprocess is not None:376 data = postprocess(data) # type: ignore377 return data378 379 query_huggingface_inference_endpoints.__name__ = alias or model_name380 381 interface_info = {382 "fn": query_huggingface_inference_endpoints,383 "inputs": inputs,384 "outputs": outputs,385 "title": model_name,386 # "examples": examples,387 }388 389 kwargs = dict(interface_info, **kwargs)390 interface = gradio.Interface(**kwargs)391 return interface392 393 394def from_spaces(395 space_name: str, hf_token: str | None, alias: str | None, **kwargs396) -> Blocks:397 client = Client(398 space_name,399 hf_token=hf_token,400 download_files=False,401 _skip_components=False,402 )403 404 space_url = f"https://huggingface.co/spaces/{space_name}"405 406 print(f"Fetching Space from: {space_url}")407 408 headers = {}409 if hf_token is not None:410 headers["Authorization"] = f"Bearer {hf_token}"411 412 iframe_url = (413 httpx.get(414 f"https://huggingface.co/api/spaces/{space_name}/host", headers=headers415 )416 .json()417 .get("host")418 )419 420 if iframe_url is None:421 raise ValueError(422 f"Could not find Space: {space_name}. If it is a private or gated Space, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."423 )424 425 r = httpx.get(iframe_url, headers=headers)426 427 result = re.search(428 r"window.gradio_config = (.*?);[\s]*</script>", r.text429 ) # some basic regex to extract the config430 try:431 config = json.loads(result.group(1)) # type: ignore432 except AttributeError as ae:433 raise ValueError(f"Could not load the Space: {space_name}") from ae434 if "allow_flagging" in config: # Create an Interface for Gradio 2.x Spaces435 return from_spaces_interface(436 space_name, config, alias, hf_token, iframe_url, **kwargs437 )438 else: # Create a Blocks for Gradio 3.x Spaces439 if kwargs:440 warnings.warn(441 "You cannot override parameters for this Space by passing in kwargs. "442 "Instead, please load the Space as a function and use it to create a "443 "Blocks or Interface locally. You may find this Guide helpful: "444 "https://gradio.app/using_blocks_like_functions/"445 )446 if client.app_version < version.Version("4.0.0b14"):447 return from_spaces_blocks(space=space_name, hf_token=hf_token)448 449 450def from_spaces_blocks(space: str, hf_token: str | None) -> Blocks:451 client = Client(452 space,453 hf_token=hf_token,454 download_files=False,455 _skip_components=False,456 )457 # We set deserialize to False to avoid downloading output files from the server.458 # Instead, we serve them as URLs using the /proxy/ endpoint directly from the server.459 460 if client.app_version < version.Version("4.0.0b14"):461 raise GradioVersionIncompatibleError(462 f"Gradio version 4.x cannot load spaces with versions less than 4.x ({client.app_version})."463 "Please downgrade to version 3 to load this space."464 )465 466 # Use end_to_end_fn here to properly upload/download all files467 predict_fns = []468 for fn_index, endpoint in client.endpoints.items():469 if not isinstance(endpoint, Endpoint):470 raise TypeError(471 f"Expected endpoint to be an Endpoint, but got {type(endpoint)}"472 )473 helper = client.new_helper(fn_index)474 if endpoint.backend_fn:475 predict_fns.append(endpoint.make_end_to_end_fn(helper))476 else:477 predict_fns.append(None)478 return gradio.Blocks.from_config(client.config, predict_fns, client.src) # type: ignore479 480 481def from_spaces_interface(482 model_name: str,483 config: dict,484 alias: str | None,485 hf_token: str | None,486 iframe_url: str,487 **kwargs,488) -> Interface:489 config = external_utils.streamline_spaces_interface(config)490 api_url = f"{iframe_url}/api/predict/"491 headers = {"Content-Type": "application/json"}492 if hf_token is not None:493 headers["Authorization"] = f"Bearer {hf_token}"494 495 # The function should call the API with preprocessed data496 def fn(*data):497 data = json.dumps({"data": data})498 response = httpx.post(api_url, headers=headers, data=data) # type: ignore499 result = json.loads(response.content.decode("utf-8"))500 if "error" in result and "429" in result["error"]:501 raise TooManyRequestsError("Too many requests to the Hugging Face API")502 try:503 output = result["data"]504 except KeyError as ke:505 raise KeyError(506 f"Could not find 'data' key in response from external Space. Response received: {result}"507 ) from ke508 if (509 len(config["outputs"]) == 1510 ): # if the fn is supposed to return a single value, pop it511 output = output[0]512 if (513 len(config["outputs"]) == 1 and isinstance(output, list)514 ): # Needed to support Output.Image() returning bounding boxes as well (TODO: handle different versions of gradio since they have slightly different APIs)515 output = output[0]516 return output517 518 fn.__name__ = alias if (alias is not None) else model_name519 config["fn"] = fn520 521 kwargs = dict(config, **kwargs)522 kwargs["_api_mode"] = True523 interface = gradio.Interface(**kwargs)524 return interface525 526 527def gr_Interface_load(528 name: str,529 src: str | None = None,530 hf_token: str | None = None,531 alias: str | None = None,532 **kwargs,533) -> Blocks:534 try:535 return load_blocks_from_repo(name, src, hf_token, alias)536 except Exception as e:537 print(e)538 return gradio.Interface(lambda: None, ['text'], ['image'])539 540 541def list_uniq(l):542 return sorted(set(l), key=l.index)543 544 545def get_status(model_name: str):546 from huggingface_hub import InferenceClient547 client = InferenceClient(timeout=10)548 return client.get_model_status(model_name)549 550 551def is_loadable(model_name: str, force_gpu: bool = False):552 try:553 status = get_status(model_name)554 except Exception as e:555 print(e)556 print(f"Couldn't load {model_name}.")557 return False558 gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()559 if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):560 print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")561 return status is not None and status.state in ["Loadable", "Loaded"] and (not force_gpu or gpu_state)562 563 564def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False):565 from huggingface_hub import HfApi566 api = HfApi()567 default_tags = ["diffusers"]568 if not sort: sort = "last_modified"569 limit = limit * 20 if check_status and force_gpu else limit * 5570 models = []571 try:572 model_infos = api.list_models(author=author, task="text-to-image",573 tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)574 except Exception as e:575 print(f"Error: Failed to list models.")576 print(e)577 return models578 for model in model_infos:579 if not model.private and not model.gated:580 loadable = is_loadable(model.id, force_gpu) if check_status else True581 if not_tag and not_tag in model.tags or not loadable: continue582 models.append(model.id)583 if len(models) == limit: break584 return models585