check valid hub id

app.py

@@ -10,8 +10,7 @@ import pandas as pd
 import os
 import backoff
 from functools import lru_cache
-
-import os
+from huggingface_hub import list_models, ModelFilter
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 
@@ -65,15 +64,20 @@ def return_random_sample(k=27):
     images = dataset[sample]["image"]
     return [resize_image(image).convert("RGB") for image in images]
 
+@lru_cache()
+def get_valid_hub_image_classification_model_ids():
+    models = list_models(limit=None, filter=ModelFilter(task="image-classification"))
+    return {model.id for model in models}
 
 def predict_subset(model_id, token):
     API_URL = f"https://api-inference.huggingface.co/models/{model_id}"
     headers = {"Authorization": f"Bearer {token}"}
-
+    valid_model_ids = get_valid_hub_image_classification_model_ids()
+    if model_id not in valid_model_ids:
+        gr.Error(f"model_id {model_id} is not a valid image classification model id")
     @backoff.on_predicate(backoff.expo, lambda x: x.status_code == 503, max_time=30)
     def _query(url):
         r = requests.post(API_URL, headers=headers, data=url)
-        print(r)
         return r
 
     @lru_cache(maxsize=1000)