import gradio as gr

from huggingface_hub import HfApi, hf_hub_download
from transformers import pipeline

def get_model_ids():
    api = HfApi()
    models = api.list_models(filter="llama-leaderboard")
    model_ids = [x.modelId for x in models]
    return model_ids
    
models = {}
for model_id in get_model_ids():
  models[model_id] = pipeline("image-classification", model=model_id)
    
def predict(img, model_id):
  preds =  models[model_id](img)
  res = {}
  for pred in preds:
    res[pred["label"]] = pred["score"]
  return res
  
gr.Interface(
  fn=predict, 
   inputs=[
     gr.inputs.Image(type="pil"),
     gr.inputs.Dropdown(get_model_ids()),
   ],
   outputs=gr.outputs.Label(num_top_classes=3),
   examples=[["llama.jpg", "osanseviero/llama-or-potato"], 
     ["potato.jpg", "osanseviero/llama-or-potato"],
     ["horse.jpg", "osanseviero/llama-horse-zebra"]]
).launch()