Upload 3 files

model.py

@@ -14,9 +14,9 @@ models = [
     'digiplay/majicMIX_realistic_v7',
     'votepurchase/counterfeitV30_v30',
     'Meina/MeinaMix_V11',
-    'John6666/cute-illustration-style-reinforced-model-v61-sd15',
     'KBlueLeaf/Kohaku-XL-Epsilon-rev3',
-    'kayfahaarukku/UrangDiffusion-1.1',
+    'KBlueLeaf/Kohaku-XL-Zeta',
+    'kayfahaarukku/UrangDiffusion-1.2',
     'Eugeoter/artiwaifu-diffusion-1.0',
     'Raelina/Rae-Diffusion-XL-V2',
     'Raelina/Raemu-XL-V4',

multit2i.py

@@ -33,22 +33,43 @@ def is_repo_name(s):
     return re.fullmatch(r'^[^/]+?/[^/]+?$', s)
 
 
-def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30):
+def get_status(model_name: str):
+    from huggingface_hub import InferenceClient
+    client = InferenceClient(timeout=10)
+    return client.get_model_status(model_name)
+
+
+def is_loadable(model_name: str, force_gpu: bool = False):
+    try:
+        status = get_status(model_name)
+    except Exception as e:
+        print(e)
+        print(f"Couldn't load {model_name}.")
+        return False
+    gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()
+    if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):
+        print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")
+    return status is not None and status.state in ["Loadable", "Loaded"] and (not force_gpu or gpu_state)
+
+
+def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False):
     from huggingface_hub import HfApi
     api = HfApi()
     default_tags = ["diffusers"]
     if not sort: sort = "last_modified"
+    limit = limit * 20 if check_status and force_gpu else limit * 5
     models = []
     try:
-        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)
+        model_infos = api.list_models(author=author, task="text-to-image",
+                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)
     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 and HF_TOKEN is None:
-           if not_tag and not_tag in model.tags: continue
+        if not model.private and not model.gated:
+           loadable = is_loadable(model.id, force_gpu) if check_status else True
+           if not_tag and not_tag in model.tags or not loadable: continue
            models.append(model.id)
            if len(models) == limit: break
     return models