import gradio as gr
import torch
from transformers import AutoModel

@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).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()