Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 
+import sys
 from huggingface_hub import ModelCard, HfApi
 import requests
 import networkx as nx
@@ -20,6 +21,23 @@ TITLE = """
 """
 
 
+# We should first try to cache models
+class CachedModelCard(ModelCard):
+  _cache = {}
+
+  @classmethod
+  def load(cls, model_id: str, **kwargs) -> "ModelCard":
+    if model_id not in cls._cache:
+      try:
+        print('REQUEST ModelCard:', model_id)
+        cls._cache[model_id] = super().load(model_id, **kwargs)
+      except:
+        cls._cache[model_id] = None
+    else:
+      print('CACHED:', model_id)
+    return cls._cache[model_id]
+
+
 def get_model_names_from_yaml(url):
     """Get a list of parent model names from the yaml file."""
     model_tags = []
@@ -32,7 +50,7 @@ def get_model_names_from_yaml(url):
 def get_license_color(model):
     """Get the color of the model based on its license."""
     try:
-        card = ModelCard.load(model)
+        card = CachedModelCard.load(model)
         license = card.data.to_dict()['license'].lower()
         # Define permissive licenses
         permissive_licenses = ['mit', 'bsd', 'apache-2.0', 'openrail']  # Add more as needed
@@ -46,47 +64,45 @@ def get_license_color(model):
         return 'lightgray'
 
 
-def get_model_names(model, genealogy, found_models=None):
-    """Get a list of parent model names from the model id."""
-    model_tags = []
-
+def get_model_names(model, genealogy, found_models=None, visited_models=None):
+    print('---')
+    print(model)
     if found_models is None:
-        found_models = []
+        found_models = set()
+    if visited_models is None:
+        visited_models = set()
+
+    if model in visited_models:
+        print("Model already visited...")
+        return found_models
+    visited_models.add(model)
 
     try:
-        card = ModelCard.load(model)
-        card_dict = card.data.to_dict()  # Convert the ModelCard object to a dictionary
+        card = CachedModelCard.load(model)
+        card_dict = card.data.to_dict()
         license = card_dict['license']
 
-        # Check the base_model in metadata
+        model_tags = []
         if 'base_model' in card_dict:
             model_tags = card_dict['base_model']
 
-        # Check the tags in metadata
         if 'tags' in card_dict and not model_tags:
             tags = card_dict['tags']
             model_tags = [model_name for model_name in tags if '/' in model_name]
 
-        # Check for merge.yml and mergekit_config.yml if no model_tags found in the tags
         if not model_tags:
             model_tags.extend(get_model_names_from_yaml(f"https://huggingface.co/{model}/blob/main/merge.yml"))
         if not model_tags:
             model_tags.extend(get_model_names_from_yaml(f"https://huggingface.co/{model}/blob/main/mergekit_config.yml"))
 
-        # Convert to a list if tags is not None or empty, else set to an empty list
         if not isinstance(model_tags, list):
             model_tags = [model_tags] if model_tags else []
 
-        # Add found model names to the list
-        found_models.extend(model_tags)
+        found_models.add(model)
 
-        # Record the genealogy
         for model_tag in model_tags:
             genealogy[model_tag].append(model)
-
-        # Recursively check for more models
-        for model_tag in model_tags:
-            get_model_names(model_tag, genealogy, found_models)
+            get_model_names(model_tag, genealogy, found_models, visited_models)
 
     except Exception as e:
         print(f"Could not find model names for {model}: {e}")
@@ -125,6 +141,8 @@ def create_family_tree(start_model):
     genealogy = defaultdict(list)
     get_model_names(start_model, genealogy)  # Assuming this populates the genealogy
 
+    print("Number of models:", len(CachedModelCard._cache))
+
     # Create a directed graph
     G = nx.DiGraph()
 
@@ -133,11 +151,14 @@ def create_family_tree(start_model):
         for child in children:
             G.add_edge(parent, child)
 
-    # Get max depth
-    max_depth = nx.dag_longest_path_length(G) + 1
-
-    # Get max width
-    max_width = max_width_of_tree(G) + 1
+    try:
+        # Get max depth and width
+        max_depth = nx.dag_longest_path_length(G) + 1
+        max_width = max_width_of_tree(G) + 1
+    except:
+        # Get max depth and width
+        max_depth = 21
+        max_width = 9
 
     # Estimate plot size
     height = max(8, 1.6 * max_depth)
@@ -149,6 +170,7 @@ def create_family_tree(start_model):
 
     # Determine node colors based on license
     node_colors = [get_license_color(node) for node in G.nodes()]
+    clear_output()
 
     # Create a label mapping with line breaks
     labels = {node: node.replace("/", "\n") for node in G.nodes()}