app.py

import gradio as gr
import os
import torch
import spaces
from transformers import AutoModel
from huggingface_hub import login

hf_token = os.getenv("HF_TOKEN")
login(token=hf_token, add_to_git_credential=True)

@spaces.GPU
def get_model_summary(model_name):
    # Check if CUDA is available and set the device accordingly
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Load the model and move it to the selected device
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to(device)
    
    # Return the model's architecture as a string
    return str(model)

# Create the Gradio interface
interface = gr.Interface(
    fn=get_model_summary,
    inputs=gr.inputs.Textbox(lines=1, placeholder="Enter model name...", label="Model Name"),
    outputs="text"),
    examples=[
        ["google/gemma-7b"],
        ["microsoft/Phi-3-mini-4k-instruct"],
        ["google/paligemma-3b-mix-224"],
        ["microsoft/Phi-3-vision-128k-instruct"]
    ]
interface.launch()