NativeAngels/blitz_diffusion
1
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, Callable, Literal13 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 36HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.37server_timeout = 60038 39 40@document()41def load(42 name: str,43 src: str | None = None,44 hf_token: str | Literal[False] | None = None,45 alias: str | None = None,46 **kwargs,47) -> Blocks:48 """49 Constructs a demo from a Hugging Face repo. Can accept model repos (if src is "models") or Space repos (if src is "spaces"). The input50 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.51 custom `css`, `js`, and `head` attributes) will not be loaded.52 Parameters:53 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")54 src: the source of the model: `models` or `spaces` (or leave empty if source is provided as a prefix in `name`)55 hf_token: optional access token for loading private Hugging Face Hub models or spaces. Will default to the locally saved token if not provided. Pass `token=False` if you don't want to send your token to the server. Find your token here: https://huggingface.co/settings/tokens. Warning: only provide a token if you are loading a trusted private Space as it can be read by the Space you are loading.56 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)57 Returns:58 a Gradio Blocks object for the given model59 Example:60 import gradio as gr61 demo = gr.load("gradio/question-answering", src="spaces")62 demo.launch()63 """64 return load_blocks_from_repo(65 name=name, src=src, hf_token=hf_token, alias=alias, **kwargs66 )67 68 69def load_blocks_from_repo(70 name: str,71 src: str | None = None,72 hf_token: str | Literal[False] | None = None,73 alias: str | None = None,74 **kwargs,75) -> Blocks:76 """Creates and returns a Blocks instance from a Hugging Face model or Space repo."""77 if src is None:78 # Separate the repo type (e.g. "model") from repo name (e.g. "google/vit-base-patch16-224")79 tokens = name.split("/")80 if len(tokens) <= 1:81 raise ValueError(82 "Either `src` parameter must be provided, or `name` must be formatted as {src}/{repo name}"83 )84 src = tokens[0]85 name = "/".join(tokens[1:])86 87 factory_methods: dict[str, Callable] = {88 # for each repo type, we have a method that returns the Interface given the model name & optionally an hf_token89 "huggingface": from_model,90 "models": from_model,91 "spaces": from_spaces,92 }93 if src.lower() not in factory_methods:94 raise ValueError(f"parameter: src must be one of {factory_methods.keys()}")95 96 if hf_token is not None and hf_token is not False:97 if Context.hf_token is not None and Context.hf_token != hf_token:98 warnings.warn(99 """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."""100 )101 Context.hf_token = hf_token102 103 blocks: gradio.Blocks = factory_methods[src](name, hf_token, alias, **kwargs)104 return blocks105 106 107def from_model(108 model_name: str, hf_token: str | Literal[False] | None, alias: str | None, **kwargs109):110 model_url = f"https://huggingface.co/{model_name}"111 api_url = f"https://api-inference.huggingface.co/models/{model_name}"112 print(f"Fetching model from: {model_url}")113 114 headers = (115 {} if hf_token in [False, None] else {"Authorization": f"Bearer {hf_token}"}116 )117 response = httpx.request("GET", api_url, headers=headers)118 if response.status_code != 200:119 raise ModelNotFoundError(120 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."121 )122 p = response.json().get("pipeline_tag")123 124 headers["X-Wait-For-Model"] = "true"125 client = huggingface_hub.InferenceClient(126 model=model_name, headers=headers, token=hf_token, timeout=server_timeout,127 )128 129 # For tasks that are not yet supported by the InferenceClient130 GRADIO_CACHE = os.environ.get("GRADIO_TEMP_DIR") or str( # noqa: N806131 Path(tempfile.gettempdir()) / "gradio"132 )133 134 def custom_post_binary(data):135 data = to_binary({"path": data})136 response = httpx.request("POST", api_url, headers=headers, content=data)137 return save_base64_to_cache(138 external_utils.encode_to_base64(response), cache_dir=GRADIO_CACHE139 )140 141 preprocess = None142 postprocess = None143 examples = None144 145 # example model: ehcalabres/wav2vec2-lg-xlsr-en-speech-emotion-recognition146 if p == "audio-classification":147 inputs = components.Audio(type="filepath", label="Input")148 outputs = components.Label(label="Class")149 postprocess = external_utils.postprocess_label150 examples = [151 "https://gradio-builds.s3.amazonaws.com/demo-files/audio_sample.wav"152 ]153 fn = client.audio_classification154 # example model: facebook/xm_transformer_sm_all-en155 elif p == "audio-to-audio":156 inputs = components.Audio(type="filepath", label="Input")157 outputs = components.Audio(label="Output")158 examples = [159 "https://gradio-builds.s3.amazonaws.com/demo-files/audio_sample.wav"160 ]161 fn = custom_post_binary162 # example model: facebook/wav2vec2-base-960h163 elif p == "automatic-speech-recognition":164 inputs = components.Audio(type="filepath", label="Input")165 outputs = components.Textbox(label="Output")166 examples = [167 "https://gradio-builds.s3.amazonaws.com/demo-files/audio_sample.wav"168 ]169 fn = client.automatic_speech_recognition170 # example model: microsoft/DialoGPT-medium171 elif p == "conversational":172 inputs = [173 components.Textbox(render=False),174 components.State(render=False),175 ]176 outputs = [177 components.Chatbot(render=False),178 components.State(render=False),179 ]180 examples = [["Hello World"]]181 preprocess = external_utils.chatbot_preprocess182 postprocess = external_utils.chatbot_postprocess183 fn = client.conversational184 # example model: julien-c/distilbert-feature-extraction185 elif p == "feature-extraction":186 inputs = components.Textbox(label="Input")187 outputs = components.Dataframe(label="Output")188 fn = client.feature_extraction189 postprocess = utils.resolve_singleton190 # example model: distilbert/distilbert-base-uncased191 elif p == "fill-mask":192 inputs = components.Textbox(label="Input")193 outputs = components.Label(label="Classification")194 examples = [195 "Hugging Face is the AI community, working together, to [MASK] the future."196 ]197 postprocess = external_utils.postprocess_mask_tokens198 fn = client.fill_mask199 # Example: google/vit-base-patch16-224200 elif p == "image-classification":201 inputs = components.Image(type="filepath", label="Input Image")202 outputs = components.Label(label="Classification")203 postprocess = external_utils.postprocess_label204 examples = ["https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg"]205 fn = client.image_classification206 # Example: deepset/xlm-roberta-base-squad2207 elif p == "question-answering":208 inputs = [209 components.Textbox(label="Question"),210 components.Textbox(lines=7, label="Context"),211 ]212 outputs = [213 components.Textbox(label="Answer"),214 components.Label(label="Score"),215 ]216 examples = [217 [218 "What entity was responsible for the Apollo program?",219 "The Apollo program, also known as Project Apollo, was the third United States human spaceflight"220 " program carried out by the National Aeronautics and Space Administration (NASA), which accomplished"221 " landing the first humans on the Moon from 1969 to 1972.",222 ]223 ]224 postprocess = external_utils.postprocess_question_answering225 fn = client.question_answering226 # Example: facebook/bart-large-cnn227 elif p == "summarization":228 inputs = components.Textbox(label="Input")229 outputs = components.Textbox(label="Summary")230 examples = [231 [232 "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."233 ]234 ]235 fn = client.summarization236 # Example: distilbert-base-uncased-finetuned-sst-2-english237 elif p == "text-classification":238 inputs = components.Textbox(label="Input")239 outputs = components.Label(label="Classification")240 examples = ["I feel great"]241 postprocess = external_utils.postprocess_label242 fn = client.text_classification243 # Example: gpt2244 elif p == "text-generation":245 inputs = components.Textbox(label="Text")246 outputs = inputs247 examples = ["Once upon a time"]248 fn = external_utils.text_generation_wrapper(client)249 # Example: valhalla/t5-small-qa-qg-hl250 elif p == "text2text-generation":251 inputs = components.Textbox(label="Input")252 outputs = components.Textbox(label="Generated Text")253 examples = ["Translate English to Arabic: How are you?"]254 fn = client.text_generation255 # Example: Helsinki-NLP/opus-mt-en-ar256 elif p == "translation":257 inputs = components.Textbox(label="Input")258 outputs = components.Textbox(label="Translation")259 examples = ["Hello, how are you?"]260 fn = client.translation261 # Example: facebook/bart-large-mnli262 elif p == "zero-shot-classification":263 inputs = [264 components.Textbox(label="Input"),265 components.Textbox(label="Possible class names (" "comma-separated)"),266 components.Checkbox(label="Allow multiple true classes"),267 ]268 outputs = components.Label(label="Classification")269 postprocess = external_utils.postprocess_label270 examples = [["I feel great", "happy, sad", False]]271 fn = external_utils.zero_shot_classification_wrapper(client)272 # Example: sentence-transformers/distilbert-base-nli-stsb-mean-tokens273 elif p == "sentence-similarity":274 inputs = [275 components.Textbox(276 label="Source Sentence",277 placeholder="Enter an original sentence",278 ),279 components.Textbox(280 lines=7,281 placeholder="Sentences to compare to -- separate each sentence by a newline",282 label="Sentences to compare to",283 ),284 ]285 outputs = components.JSON(label="Similarity scores")286 examples = [["That is a happy person", "That person is very happy"]]287 fn = external_utils.sentence_similarity_wrapper(client)288 # Example: julien-c/ljspeech_tts_train_tacotron2_raw_phn_tacotron_g2p_en_no_space_train289 elif p == "text-to-speech":290 inputs = components.Textbox(label="Input")291 outputs = components.Audio(label="Audio")292 examples = ["Hello, how are you?"]293 fn = client.text_to_speech294 # example model: osanseviero/BigGAN-deep-128295 elif p == "text-to-image":296 inputs = components.Textbox(label="Input")297 outputs = components.Image(label="Output")298 examples = ["A beautiful sunset"]299 fn = client.text_to_image300 # example model: huggingface-course/bert-finetuned-ner301 elif p == "token-classification":302 inputs = components.Textbox(label="Input")303 outputs = components.HighlightedText(label="Output")304 examples = [305 "Hugging Face is a company based in Paris and New York City that acquired Gradio in 2021."306 ]307 fn = external_utils.token_classification_wrapper(client)308 # example model: impira/layoutlm-document-qa309 elif p == "document-question-answering":310 inputs = [311 components.Image(type="filepath", label="Input Document"),312 components.Textbox(label="Question"),313 ]314 postprocess = external_utils.postprocess_label315 outputs = components.Label(label="Label")316 fn = client.document_question_answering317 # example model: dandelin/vilt-b32-finetuned-vqa318 elif p == "visual-question-answering":319 inputs = [320 components.Image(type="filepath", label="Input Image"),321 components.Textbox(label="Question"),322 ]323 outputs = components.Label(label="Label")324 postprocess = external_utils.postprocess_visual_question_answering325 examples = [326 [327 "https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg",328 "What animal is in the image?",329 ]330 ]331 fn = client.visual_question_answering332 # example model: Salesforce/blip-image-captioning-base333 elif p == "image-to-text":334 inputs = components.Image(type="filepath", label="Input Image")335 outputs = components.Textbox(label="Generated Text")336 examples = ["https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg"]337 fn = client.image_to_text338 # example model: rajistics/autotrain-Adult-934630783339 elif p in ["tabular-classification", "tabular-regression"]:340 examples = external_utils.get_tabular_examples(model_name)341 col_names, examples = external_utils.cols_to_rows(examples) # type: ignore342 examples = [[examples]] if examples else None343 inputs = components.Dataframe(344 label="Input Rows",345 type="pandas",346 headers=col_names,347 col_count=(len(col_names), "fixed"),348 render=False,349 )350 outputs = components.Dataframe(351 label="Predictions", type="array", headers=["prediction"]352 )353 fn = external_utils.tabular_wrapper354 # example model: microsoft/table-transformer-detection355 elif p == "object-detection":356 inputs = components.Image(type="filepath", label="Input Image")357 outputs = components.AnnotatedImage(label="Annotations")358 fn = external_utils.object_detection_wrapper(client)359 # example model: stabilityai/stable-diffusion-xl-refiner-1.0360 elif p == "image-to-image":361 inputs = [362 components.Image(type="filepath", label="Input Image"),363 components.Textbox(label="Input"),364 ]365 outputs = components.Image(label="Output")366 examples = [367 [368 "https://gradio-builds.s3.amazonaws.com/demo-files/cheetah-002.jpg",369 "Photo of a cheetah with green eyes",370 ]371 ]372 fn = client.image_to_image373 else:374 raise ValueError(f"Unsupported pipeline type: {p}")375 376 def query_huggingface_inference_endpoints(*data, **kwargs):377 if preprocess is not None:378 data = preprocess(*data)379 try:380 data = fn(*data, **kwargs) # type: ignore381 except huggingface_hub.utils.HfHubHTTPError as e:382 if "429" in str(e):383 raise TooManyRequestsError() from e384 if postprocess is not None:385 data = postprocess(data) # type: ignore386 return data387 388 query_huggingface_inference_endpoints.__name__ = alias or model_name389 390 interface_info = {391 "fn": query_huggingface_inference_endpoints,392 "inputs": inputs,393 "outputs": outputs,394 "title": model_name,395 #"examples": examples,396 }397 398 kwargs = dict(interface_info, **kwargs)399 interface = gradio.Interface(**kwargs)400 return interface401 402 403def from_spaces(404 space_name: str, hf_token: str | None, alias: str | None, **kwargs405) -> Blocks:406 space_url = f"https://huggingface.co/spaces/{space_name}"407 408 print(f"Fetching Space from: {space_url}")409 410 headers = {}411 if hf_token not in [False, None]:412 headers["Authorization"] = f"Bearer {hf_token}"413 414 iframe_url = (415 httpx.get(416 f"https://huggingface.co/api/spaces/{space_name}/host", headers=headers417 )418 .json()419 .get("host")420 )421 422 if iframe_url is None:423 raise ValueError(424 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."425 )426 427 r = httpx.get(iframe_url, headers=headers)428 429 result = re.search(430 r"window.gradio_config = (.*?);[\s]*</script>", r.text431 ) # some basic regex to extract the config432 try:433 config = json.loads(result.group(1)) # type: ignore434 except AttributeError as ae:435 raise ValueError(f"Could not load the Space: {space_name}") from ae436 if "allow_flagging" in config: # Create an Interface for Gradio 2.x Spaces437 return from_spaces_interface(438 space_name, config, alias, hf_token, iframe_url, **kwargs439 )440 else: # Create a Blocks for Gradio 3.x Spaces441 if kwargs:442 warnings.warn(443 "You cannot override parameters for this Space by passing in kwargs. "444 "Instead, please load the Space as a function and use it to create a "445 "Blocks or Interface locally. You may find this Guide helpful: "446 "https://gradio.app/using_blocks_like_functions/"447 )448 return from_spaces_blocks(space=space_name, hf_token=hf_token)449 450 451def from_spaces_blocks(space: str, hf_token: str | None) -> Blocks:452 client = Client(453 space,454 hf_token=hf_token,455 download_files=False,456 _skip_components=False,457 )458 # We set deserialize to False to avoid downloading output files from the server.459 # Instead, we serve them as URLs using the /proxy/ endpoint directly from the server.460 461 if client.app_version < version.Version("4.0.0b14"):462 raise GradioVersionIncompatibleError(463 f"Gradio version 4.x cannot load spaces with versions less than 4.x ({client.app_version})."464 "Please downgrade to version 3 to load this space."465 )466 467 # Use end_to_end_fn here to properly upload/download all files468 predict_fns = []469 for fn_index, endpoint in client.endpoints.items():470 if not isinstance(endpoint, Endpoint):471 raise TypeError(472 f"Expected endpoint to be an Endpoint, but got {type(endpoint)}"473 )474 helper = client.new_helper(fn_index)475 if endpoint.backend_fn:476 predict_fns.append(endpoint.make_end_to_end_fn(helper))477 else:478 predict_fns.append(None)479 return gradio.Blocks.from_config(client.config, predict_fns, client.src) # type: ignore480 481 482def from_spaces_interface(483 model_name: str,484 config: dict,485 alias: str | None,486 hf_token: str | None,487 iframe_url: str,488 **kwargs,489) -> Interface:490 config = external_utils.streamline_spaces_interface(config)491 api_url = f"{iframe_url}/api/predict/"492 headers = {"Content-Type": "application/json"}493 if hf_token not in [False, None]:494 headers["Authorization"] = f"Bearer {hf_token}"495 496 # The function should call the API with preprocessed data497 def fn(*data):498 data = json.dumps({"data": data})499 response = httpx.post(api_url, headers=headers, data=data) # type: ignore500 result = json.loads(response.content.decode("utf-8"))501 if "error" in result and "429" in result["error"]:502 raise TooManyRequestsError("Too many requests to the Hugging Face API")503 try:504 output = result["data"]505 except KeyError as ke:506 raise KeyError(507 f"Could not find 'data' key in response from external Space. Response received: {result}"508 ) from ke509 if (510 len(config["outputs"]) == 1511 ): # if the fn is supposed to return a single value, pop it512 output = output[0]513 if (514 len(config["outputs"]) == 1 and isinstance(output, list)515 ): # Needed to support Output.Image() returning bounding boxes as well (TODO: handle different versions of gradio since they have slightly different APIs)516 output = output[0]517 return output518 519 fn.__name__ = alias if (alias is not None) else model_name520 config["fn"] = fn521 522 kwargs = dict(config, **kwargs)523 kwargs["_api_mode"] = True524 interface = gradio.Interface(**kwargs)525 return interface526 527 528def gr_Interface_load(529 name: str,530 src: str | None = None,531 hf_token: str | None = None,532 alias: str | None = None,533 **kwargs, # ignore534) -> Blocks:535 try:536 return load_blocks_from_repo(name, src, hf_token, alias)537 except Exception as e:538 print(e)539 return gradio.Interface(lambda: None, ['text'], ['image'])540 541 542def list_uniq(l):543 return sorted(set(l), key=l.index)544 545 546def get_status(model_name: str):547 from huggingface_hub import AsyncInferenceClient548 client = AsyncInferenceClient(token=HF_TOKEN, timeout=10)549 return client.get_model_status(model_name)550 551 552def is_loadable(model_name: str, force_gpu: bool = False):553 try:554 status = get_status(model_name)555 except Exception as e:556 print(e)557 print(f"Couldn't load {model_name}.")558 return False559 gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()560 if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):561 print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")562 return status is not None and status.state in ["Loadable", "Loaded"] and (not force_gpu or gpu_state)563 564 565def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False):566 from huggingface_hub import HfApi567 api = HfApi(token=HF_TOKEN)568 default_tags = ["diffusers"]569 if not sort: sort = "last_modified"570 limit = limit * 20 if check_status and force_gpu else limit * 5571 models = []572 try:573 model_infos = api.list_models(author=author, #task="text-to-image",574 tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)575 except Exception as e:576 print(f"Error: Failed to list models.")577 print(e)578 return models579 for model in model_infos:580 if not model.private and not model.gated or HF_TOKEN is not None:581 loadable = is_loadable(model.id, force_gpu) if check_status else True582 if not_tag and not_tag in model.tags or not loadable: continue583 models.append(model.id)584 if len(models) == limit: break585 return models586 587 588def save_image(image, savefile, modelname, prompt, nprompt, height=0, width=0, steps=0, cfg=0, seed=-1):589 from PIL import Image, PngImagePlugin590 import json591 try:592 metadata = {"prompt": prompt, "negative_prompt": nprompt, "Model": {"Model": modelname.split("/")[-1]}}593 if steps > 0: metadata["num_inference_steps"] = steps594 if cfg > 0: metadata["guidance_scale"] = cfg595 if seed != -1: metadata["seed"] = seed596 if width > 0 and height > 0: metadata["resolution"] = f"{width} x {height}"597 metadata_str = json.dumps(metadata)598 info = PngImagePlugin.PngInfo()599 info.add_text("metadata", metadata_str)600 image.save(savefile, "PNG", pnginfo=info)601 return str(Path(savefile).resolve())602 except Exception as e:603 print(f"Failed to save image file: {e}")604 raise Exception(f"Failed to save image file:") from e605 606 607def randomize_seed():608 from random import seed, randint609 MAX_SEED = 2**32-1610 seed()611 rseed = randint(0, MAX_SEED)612 return rseed