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, hf_token=hf_token).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="text", outputs="text")
interface.launch()