Update app.py

app.py

@@ -6,21 +6,39 @@ from threading import Timer
 
 HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
 def get_available_free():
-    models = InferenceClient(token=HUGGINGFACE_TOKEN).list_deployed_models("text-generation-inference")['text-generation']
+    models_dict = InferenceClient(token=HUGGINGFACE_TOKEN).list_deployed_models("text-generation-inference")
+    models = models_dict['text-generation'] + models_dict['text2text-generation'] #token=HUGGINGFACE_TOKEN
+    models_vision = models_dict['image-text-to-text']
+    
     models_conclusion = {
         "Model": [],
         "API": [],
         "Text Completion": [],
-        "Chat Completion": []
+        "Chat Completion": [],
+        "Vision": []
     }
-    for m in models:
+    for m in models + models_vision:
         text_available = False
         chat_available = False
+        vision_available = False
+        if m in models_vision:
+            vision_available = True
         pro_sub = False
         try:
             InferenceClient(m, timeout=10, token=HUGGINGFACE_TOKEN).text_generation("Hi.", max_new_tokens=1)
             text_available = True
-            InferenceClient(m, timeout=10, token=HUGGINGFACE_TOKEN).chat_completion(messages=[{'role': 'user', 'content': 'Hi.'}], max_tokens=1)
+        except Exception as e:
+            print(e)
+            if e and "Model requires a Pro subscription" in str(e):
+                pro_sub = True
+            if e and "Rate limit reached" in str(e):
+                print("Rate Limited!!")
+                if os.path.exists("data.csv"):
+                    print("Loading data from file...")
+                    return pd.read_csv("data.csv").to_dict(orient='list')
+                return []
+        try:
+            InferenceClient(m, timeout=10).chat_completion(messages=[{'role': 'user', 'content': 'Hi.'}], max_tokens=1) #token=HUGGINGFACE_TOKEN
             chat_available = True
         except Exception as e:
             print(e)
@@ -36,6 +54,7 @@ def get_available_free():
         models_conclusion["API"].append("Free" if chat_available or text_available else ("Pro Subscription" if pro_sub else "Not Responding"))
         models_conclusion["Chat Completion"].append("---" if (pro_sub or (not chat_available and not text_available)) else ("✓" if chat_available else "⌀"))
         models_conclusion["Text Completion"].append("---" if (pro_sub or (not chat_available and not text_available)) else ("✓" if text_available else "⌀"))
+        models_conclusion["Vision"].append("---" if (pro_sub or (not chat_available and not text_available)) else ("✓" if vision_available else "⌀"))
     pd.DataFrame(models_conclusion).to_csv("data.csv", index=False)
     return models_conclusion
 
@@ -48,7 +67,7 @@ def update_data():
     df['Text Completion'] = df['Text Completion'].map(status_mapping)
     df['Chat Completion'] = df['Chat Completion'].map(status_mapping)
     
-    df = df.sort_values(by=['API', 'Text Completion', 'Chat Completion'])
+    df = df.sort_values(by=['API', 'Text Completion', 'Chat Completion', 'Vision'])
     
     df['Text Completion'] = df['Text Completion'].map({v: k for k, v in status_mapping.items()})
     df['Chat Completion'] = df['Chat Completion'].map({v: k for k, v in status_mapping.items()})
@@ -62,7 +81,7 @@ def display_table(search_query=""):
     else:
         filtered_df = df
     
-    styled_df = filtered_df.style.apply(apply_row_styles, axis=1, subset=["Model", "API", "Text Completion", "Chat Completion"])
+    styled_df = filtered_df.style.apply(apply_row_styles, axis=1, subset=["Model", "API", "Text Completion", "Chat Completion", "Vision"])
     return styled_df
 
 def apply_row_styles(row):
@@ -71,7 +90,8 @@ def apply_row_styles(row):
         color_status(api_value, row["Model"]),
         color_status(api_value, row["API"]),
         color_status(api_value, row["Text Completion"]),
-        color_status(api_value, row["Chat Completion"])
+        color_status(api_value, row["Chat Completion"]),
+        color_status(api_value, row["Vision"])
     ]
 
 def color_status(api_value, cell_value):