import gradio as gr

def format_params(params):
    if params >= 1e9:
        return f"{params / 1e9:.2f}B"
    elif params >= 1e6:
        return f"{params / 1e6:.2f}M"
    return str(params)

def calculate_training_memory(params, precision, batch_size, seq_length, num_heads, head_dim, num_layers):
    bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
    
    # Model Weights
    model_memory = params * bytes_per_param
    
    # Optimizer States (Adam)
    optimizer_memory = model_memory * 2
    
    # Gradients
    gradient_memory = model_memory
    
    # Activation Memory (approximate formula)
    activation_memory = batch_size * seq_length * num_heads * head_dim * num_layers * bytes_per_param
    
    # Total Training Memory
    total_memory = model_memory + optimizer_memory + gradient_memory + activation_memory
    
    return f"Model Weights: {model_memory / 1e9:.2f} GB\nOptimizer: {optimizer_memory / 1e9:.2f} GB\nGradients: {gradient_memory / 1e9:.2f} GB\nActivation Memory: {activation_memory / 1e9:.2f} GB\nTotal Training Memory: {total_memory / 1e9:.2f} GB"

def calculate_inference_memory(params, precision, batch_size, seq_length, num_heads, head_dim, num_layers):
    bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
    
    # Model Weights
    model_memory = params * bytes_per_param
    
    # KV Cache
    kv_cache_memory = batch_size * seq_length * num_heads * head_dim * 2 * num_layers * bytes_per_param
    
    # Total Inference Memory
    total_memory = model_memory + kv_cache_memory
    
    return f"Model Weights: {model_memory / 1e9:.2f} GB\nKV Cache: {kv_cache_memory / 1e9:.2f} GB\nTotal Inference Memory: {total_memory / 1e9:.2f} GB"

def calculate_kv_cache(batch_size, seq_length, num_heads, head_dim, num_layers, precision):
    bytes_per_param = 2 if precision in ["FP16/BF16", "BF16"] else 4
    
    # KV Cache Calculation
    kv_cache_memory = batch_size * seq_length * num_heads * head_dim * 2 * num_layers * bytes_per_param
    
    return f"KV Cache Memory: {kv_cache_memory / 1e9:.2f} GB"

with gr.Blocks() as app:
    gr.Markdown("# GPU Memory Calculator for Transformer Models")
    
    with gr.Tabs():
        with gr.Tab("Training Memory Calculation"):
            with gr.Row():
                params = gr.Number(label="Number of Parameters (e.g., 175B = 175e9)", value=175e9)
                precision = gr.Radio(["FP16/BF16", "FP32"], label="Precision", value="FP16/BF16")
            with gr.Row():
                batch_size = gr.Number(label="Batch Size", value=1)
                seq_length = gr.Number(label="Sequence Length", value=2048)
            with gr.Row():
                num_heads = gr.Number(label="Number of Attention Heads", value=96)
                head_dim = gr.Number(label="Head Dimension", value=128)
                num_layers = gr.Number(label="Number of Layers", value=96)
            train_button = gr.Button("Calculate Training Memory")
            train_output = gr.Textbox(label="Training Memory Usage")
            train_button.click(calculate_training_memory, [params, precision, batch_size, seq_length, num_heads, head_dim, num_layers], train_output)
        
        with gr.Tab("Inference Memory Calculation"):
            with gr.Row():
                params_inf = gr.Number(label="Number of Parameters (e.g., 175B = 175e9)", value=175e9)
                precision_inf = gr.Radio(["FP16/BF16", "FP32"], label="Precision", value="FP16/BF16")
            with gr.Row():
                batch_size_inf = gr.Number(label="Batch Size", value=1)
                seq_length_inf = gr.Number(label="Sequence Length", value=2048)
            with gr.Row():
                num_heads_inf = gr.Number(label="Number of Attention Heads", value=96)
                head_dim_inf = gr.Number(label="Head Dimension", value=128)
                num_layers_inf = gr.Number(label="Number of Layers", value=96)
            infer_button = gr.Button("Calculate Inference Memory")
            infer_output = gr.Textbox(label="Inference Memory Usage")
            infer_button.click(calculate_inference_memory, [params_inf, precision_inf, batch_size_inf, seq_length_inf, num_heads_inf, head_dim_inf, num_layers_inf], infer_output)
        
        with gr.Tab("KV Cache Calculation"):
            with gr.Row():
                batch_size_kv = gr.Number(label="Batch Size", value=1)
                seq_length_kv = gr.Number(label="Sequence Length", value=2048)
            with gr.Row():
                num_heads_kv = gr.Number(label="Number of Attention Heads", value=96)
                head_dim_kv = gr.Number(label="Head Dimension", value=128)
                num_layers_kv = gr.Number(label="Number of Layers", value=96)
                precision_kv = gr.Radio(["FP16/BF16", "FP32"], label="Precision", value="FP16/BF16")
            kv_button = gr.Button("Calculate KV Cache Memory")
            kv_output = gr.Textbox(label="KV Cache Memory Usage")
            kv_button.click(calculate_kv_cache, [batch_size_kv, seq_length_kv, num_heads_kv, head_dim_kv, num_layers_kv, precision_kv], kv_output)
    
app.launch()