from huggingface_hub import list_models
import gradio as gr
from transformers import pipeline

# Fetch timm models from Hugging Face Hub
timm_models = list_models(filter="timm", sort="downloads", limit=20)  # Fetch top 20 based on downloads
model_ids = [model.modelId for model in timm_models]

# Initialize a pipeline with a default model
default_model = model_ids[0]
pipe = pipeline("image-classification", model=default_model)

# Function for classification
def classify(image, model_name):
    pipe.model = model_name  # Update model dynamically
    results = pipe(image)
    return {result["label"]: round(result["score"], 2) for result in results}

# Gradio Interface
demo = gr.Interface(
    fn=classify,
    inputs=[
        gr.Image(type="pil", label="Upload an Image"),
        gr.Dropdown(choices=model_ids, label="Select timm Model", value=default_model)
    ],
    outputs=gr.Label(num_top_classes=3, label="Top Predictions"),
    title="timm Model Image Classifier",
    description="Select a timm model and upload an image for classification."
)

demo.launch()