Upload multit2i.py

multit2i.py

@@ -143,20 +143,23 @@ def save_gallery(image_path: str | None, images: list[tuple] | None):
 
 # https://github.com/gradio-app/gradio/blob/main/gradio/external.py
 # https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
-def load_from_model(model_name: str, hf_token: str = None):
+from typing import Literal
+def load_from_model(model_name: str, hf_token: str | Literal[False] | None = None):
     import httpx
     import huggingface_hub
-    from gradio.exceptions import ModelNotFoundError
+    from gradio.exceptions import ModelNotFoundError, TooManyRequestsError
     model_url = f"https://huggingface.co/{model_name}"
     api_url = f"https://api-inference.huggingface.co/models/{model_name}"
     print(f"Fetching model from: {model_url}")
 
-    headers = {"Authorization": f"Bearer {hf_token}"} if hf_token is not None else {}
+    headers = ({} if hf_token in [False, None] else {"Authorization": f"Bearer {hf_token}"})
     response = httpx.request("GET", api_url, headers=headers)
     if response.status_code != 200:
         raise ModelNotFoundError(
             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."
         )
+    p = response.json().get("pipeline_tag")
+    if p != "text-to-image": raise ModelNotFoundError(f"This model isn't for text-to-image or unsupported: {model_name}.")
     headers["X-Wait-For-Model"] = "true"
     client = huggingface_hub.InferenceClient(model=model_name, headers=headers,
                                               token=hf_token, timeout=server_timeout)
@@ -165,7 +168,14 @@ def load_from_model(model_name: str, hf_token: str = None):
     fn = client.text_to_image
 
     def query_huggingface_inference_endpoints(*data, **kwargs):
-        return fn(*data, **kwargs)
+        try:
+            data = fn(*data, **kwargs)  # type: ignore
+        except huggingface_hub.utils.HfHubHTTPError as e:
+            if "429" in str(e):
+                raise TooManyRequestsError() from e
+        except Exception as e:
+            raise Exception(e)
+        return data
 
     interface_info = {
         "fn": query_huggingface_inference_endpoints,
@@ -370,6 +380,7 @@ def infer_body(client: InferenceClient | gr.Interface | object, prompt: str, neg
         elif isinstance(client, gr.Interface):
             image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
         else: return None
+        if isinstance(image, tuple): return None
         image.save(png_path)
         return str(Path(png_path).resolve())
     except Exception as e: