import gradio as gr
import pandas as pd
from huggingface_hub.utils import HfHubHTTPError
from accelerate.commands.estimate import check_has_model, create_empty_model
from accelerate.utils import calculate_maximum_sizes
from model_utils import get_model

# We need to store them as globals because gradio doesn't have a way for us to pass them in to the button
MODEL = None


def get_results(model_name: str, library: str, precision: list, training: list,  access_token: str, zero_stage: int, num_nodes: int, num_gpus: int, offloading: list, zero_init: list, additional_buffer_factor: float):
    global MODEL
    
    MODEL = get_model(model_name, library, access_token)

    data = MODEL.__repr__()
    
    title = f"## Model Representation for '{model_name}'"
    return [title, gr.update(visible=True, value=data)]

with gr.Blocks() as demo:
    with gr.Column():
        out_text = gr.Markdown()
        out = gr.Code()

        with gr.Row():
            inp = gr.Textbox(label="Model Name or URL", value="bert-base-cased")
        with gr.Row():
            library = gr.Radio(["auto", "transformers", "timm"], label="Library", value="auto")
            access_token = gr.Textbox(label="API Token", placeholder="Optional (for gated models)")
        with gr.Row():
            btn = gr.Button("Calculate Memory Usage")
        

    btn.click(
        get_results,
        inputs=[inp, library, access_token,],
        outputs=[out_text, out],
    )

demo.launch()