Upload multit2i.py

multit2i.py

@@ -3,8 +3,10 @@ import asyncio
 from threading import RLock
 from pathlib import Path
 from huggingface_hub import InferenceClient
+import os
 
 
+HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None
 server_timeout = 600
 inference_timeout = 300
 
@@ -38,14 +40,14 @@ def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="l
     if not sort: sort = "last_modified"
     models = []
     try:
-        model_infos = api.list_models(author=author, pipeline_tag="text-to-image",
+        model_infos = api.list_models(author=author, pipeline_tag="text-to-image", token=HF_TOKEN,
                                        tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit * 5)
     except Exception as e:
         print(f"Error: Failed to list models.")
         print(e)
         return models
     for model in model_infos:
-        if not model.private and not model.gated:
+        if not model.private and not model.gated and HF_TOKEN is None:
            if not_tag and not_tag in model.tags: continue
            models.append(model.id)
            if len(models) == limit: break
@@ -58,7 +60,7 @@ def get_t2i_model_info_dict(repo_id: str):
     info = {"md": "None"}
     try:
         if not is_repo_name(repo_id) or not api.repo_exists(repo_id=repo_id): return info
-        model = api.model_info(repo_id=repo_id)
+        model = api.model_info(repo_id=repo_id, token=HF_TOKEN)
     except Exception as e:
         print(f"Error: Failed to get {repo_id}'s info.")
         print(e)
@@ -156,7 +158,7 @@ def load_model(model_name: str):
     global model_info_dict
     if model_name in loaded_models.keys(): return loaded_models[model_name]
     try:
-        loaded_models[model_name] = load_from_model(model_name)
+        loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
         print(f"Loaded: {model_name}")
     except Exception as e:
         if model_name in loaded_models.keys(): del loaded_models[model_name]
@@ -179,12 +181,12 @@ def load_model_api(model_name: str):
     if model_name in loaded_models.keys(): return loaded_models[model_name]
     try:
         client = InferenceClient(timeout=5)
-        status = client.get_model_status(model_name)
+        status = client.get_model_status(model_name, token=HF_TOKEN)
         if status is None or status.framework != "diffusers" or status.state not in ["Loadable", "Loaded"]:
             print(f"Failed to load by API: {model_name}")
             return None
         else:
-            loaded_models[model_name] = InferenceClient(model_name, timeout=server_timeout)
+            loaded_models[model_name] = InferenceClient(model_name, token=HF_TOKEN, timeout=server_timeout)
             print(f"Loaded by API: {model_name}")
     except Exception as e:
         if model_name in loaded_models.keys(): del loaded_models[model_name]
@@ -340,9 +342,9 @@ def infer_body(client: InferenceClient | gr.Interface, prompt: str, neg_prompt:
     if cfg is not None and cfg > 0: cfg = kwargs["guidance_scale"] = cfg
     try:
         if isinstance(client, InferenceClient):
-            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
+            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
         elif isinstance(client, gr.Interface):
-            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs)
+            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
         else: return None
         image.save(png_path)
         return str(Path(png_path).resolve())