app.py

import gradio as gr
import requests
import json
from transformers import AutoConfig
import math
from typing import Dict, Tuple, Optional


class LLMMemoryCalculator:
    def __init__(self):
        self.precision_bytes = {
            'fp32': 4,
            'fp16': 2,
            'bf16': 2,
            'int8': 1,
            'int4': 0.5
        }

    # -------------------------------------------------
    # 📥 基础工具
    # -------------------------------------------------
    def get_model_config(self, model_id: str) -> Dict:
        """获取模型配置"""
        try:
            config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
            return config
        except Exception as e:
            raise Exception(f"无法获取模型配置: {str(e)}")

    def get_file_size_from_url(self, model_id: str, filename: str) -> int:
        """通过 HEAD 请求获取文件大小（备用）"""
        try:
            url = f"https://huggingface.co/{model_id}/resolve/main/{filename}"
            response = requests.head(url, timeout=10)
            if response.status_code == 200:
                content_length = response.headers.get('Content-Length')
                if content_length:
                    return int(content_length)
            return 0
        except:
            return 0

    # -------------------------------------------------
    # 📦 获取模型权重大小
    # -------------------------------------------------
    def get_model_size_from_hf(self, model_id: str) -> Tuple[float, str]:
        """优先使用 *.index.json 中的 metadata.total_size，回退到文件列表/HEAD"""
        try:
            # 1️⃣ 尝试读取 index.json（safetensors > pytorch）
            for index_name, tag in [
                ("model.safetensors.index.json", "safetensors_index"),
                ("pytorch_model.bin.index.json", "pytorch_index")
            ]:
                url = f"https://huggingface.co/{model_id}/resolve/main/{index_name}"
                resp = requests.get(url, timeout=10)
                if resp.status_code == 200:
                    try:
                        data = resp.json()
                    except ValueError:
                        # 某些仓库 index.json 以文本形式存储，需要手动解析
                        data = json.loads(resp.text)
                    total_bytes = data.get("metadata", {}).get("total_size", 0)
                    if total_bytes > 0:
                        return total_bytes / (1024 ** 3), tag

            # 2️⃣ 调用 Hub API，尝试直接读取 size 字段
            api_url = f"https://huggingface.co/api/models/{model_id}"
            response = requests.get(api_url, timeout=10)
            if response.status_code != 200:
                raise Exception(f"API请求失败: {response.status_code}")
            model_info = response.json()

            # 2a. 查找 siblings 列表中带 size 的 .safetensors 文件
            safetensors_files = [f for f in model_info.get('siblings', [])
                                 if f['rfilename'].endswith('.safetensors') and 'size' in f]
            if safetensors_files:
                total_size = sum(f['size'] for f in safetensors_files)
                return total_size / (1024 ** 3), "safetensors_files"

            # 2b. 使用 HEAD 请求补全未包含 size 的 .safetensors 文件
            safetensors_no_size = [f for f in model_info.get('siblings', [])
                                   if f['rfilename'].endswith('.safetensors')]
            if safetensors_no_size:
                total_size = 0
                for f in safetensors_no_size:
                    total_size += self.get_file_size_from_url(model_id, f['rfilename'])
                if total_size > 0:
                    return total_size / (1024 ** 3), "safetensors_head"

            # 2c. 同理处理 pytorch_model-xxxxx.bin
            pytorch_files = [f for f in model_info.get('siblings', [])
                             if f['rfilename'].endswith('.bin') and 'size' in f]
            if pytorch_files:
                total_size = sum(f['size'] for f in pytorch_files)
                return total_size / (1024 ** 3), "pytorch_files"

            pytorch_no_size = [f for f in model_info.get('siblings', [])
                               if f['rfilename'].endswith('.bin')]
            if pytorch_no_size:
                total_size = 0
                for f in pytorch_no_size:
                    total_size += self.get_file_size_from_url(model_id, f['rfilename'])
                if total_size > 0:
                    return total_size / (1024 ** 3), "pytorch_head"

            # 3️⃣ 如果仍然无法确定大小，走估算逻辑
            raise Exception("未找到权重大小信息")

        except Exception:
            # 估算
            return self.estimate_model_size_from_config(model_id)

    # -------------------------------------------------
    # 📐 估算逻辑（与原始保持一致）
    # -------------------------------------------------
    def estimate_model_size_from_config(self, model_id: str) -> Tuple[float, str]:
        """根据 config.json 估算模型大小（FP16）"""
        try:
            config = self.get_model_config(model_id)

            vocab_size = getattr(config, 'vocab_size', 50000)
            hidden_size = getattr(config, 'hidden_size', getattr(config, 'd_model', 4096))
            num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'num_layers', 32))
            intermediate_size = getattr(config, 'intermediate_size', hidden_size * 4)

            # Embedding
            embedding_params = vocab_size * hidden_size

            # Transformer layer
            attention_params = 4 * hidden_size * hidden_size
            ffn_params = 2 * hidden_size * intermediate_size
            ln_params = 2 * hidden_size
            params_per_layer = attention_params + ffn_params + ln_params

            total_params = embedding_params + num_layers * params_per_layer
            if hasattr(config, 'tie_word_embeddings') and not config.tie_word_embeddings:
                total_params += vocab_size * hidden_size

            model_size_gb = (total_params * 2) / (1024 ** 3)  # 默认 fp16
            return model_size_gb, "estimated"

        except Exception as e:
            raise Exception(f"无法估算模型大小: {str(e)}")

    # -------------------------------------------------
    # 🗄️ KV Cache 计算（原逻辑保持）
    # -------------------------------------------------
    def calculate_kv_cache_size(self, config, context_length: int, batch_size: int = 1) -> Dict[str, float]:
        try:
            num_layers = getattr(config, 'num_hidden_layers', getattr(config, 'num_layers', 32))
            hidden_size = getattr(config, 'hidden_size', getattr(config, 'd_model', 4096))
            num_attention_heads = getattr(config, 'num_attention_heads', getattr(config, 'num_heads', 32))
            num_key_value_heads = getattr(config, 'num_key_value_heads', num_attention_heads)
            is_mla = hasattr(config, 'kv_lora_rank') and config.kv_lora_rank is not None
            head_dim = hidden_size // num_attention_heads

            if is_mla:
                kv_lora_rank = getattr(config, 'kv_lora_rank', 512)
                kv_cache_per_token = kv_lora_rank * 2
                attention_type = "MLA"
            elif num_key_value_heads < num_attention_heads:
                kv_cache_per_token = num_key_value_heads * head_dim * 2
                attention_type = "GQA"
            else:
                kv_cache_per_token = num_attention_heads * head_dim * 2
                attention_type = "MHA"

            total_kv_cache = (kv_cache_per_token * context_length * num_layers * batch_size * 2) / (1024 ** 3)
            return {
                'size_gb': total_kv_cache,
                'attention_type': attention_type,
                'num_kv_heads': num_key_value_heads,
                'num_attention_heads': num_attention_heads,
                'head_dim': head_dim
            }
        except Exception as e:
            raise Exception(f"计算KV Cache失败: {str(e)}")

    # -------------------------------------------------
    # 🧮 综合内存需求计算（保持不变）
    # -------------------------------------------------
    def calculate_memory_requirements(self, model_id: str, gpu_memory_gb: float, num_gpus: int,
                                      context_length: int, utilization_rate: float = 0.9) -> Dict:
        try:
            config = self.get_model_config(model_id)
            model_size_gb, size_source = self.get_model_size_from_hf(model_id)
            kv_info = self.calculate_kv_cache_size(config, context_length)

            available_memory = gpu_memory_gb * num_gpus * utilization_rate
            other_overhead = model_size_gb * 0.1
            total_memory_needed = model_size_gb + kv_info['size_gb'] + other_overhead

            is_feasible = total_memory_needed <= available_memory
            memory_margin = available_memory - total_memory_needed
            memory_per_gpu = total_memory_needed / num_gpus

            return {
                'model_id': model_id,
                'model_size_gb': round(model_size_gb, 2),
                'size_source': size_source,
                'kv_cache_gb': round(kv_info['size_gb'], 2),
                'attention_type': kv_info['attention_type'],
                'other_overhead_gb': round(other_overhead, 2),
                'total_memory_needed_gb': round(total_memory_needed, 2),
                'available_memory_gb': round(available_memory, 2),
                'memory_margin_gb': round(memory_margin, 2),
                'memory_per_gpu_gb': round(memory_per_gpu, 2),
                'is_feasible': is_feasible,
                'utilization_per_gpu': round((memory_per_gpu / gpu_memory_gb) * 100, 1),
                'config_info': {
                    'num_layers': getattr(config, 'num_hidden_layers', getattr(config, 'num_layers', 'N/A')),
                    'hidden_size': getattr(config, 'hidden_size', getattr(config, 'd_model', 'N/A')),
                    'num_attention_heads': kv_info['num_attention_heads'],
                    'num_kv_heads': kv_info['num_kv_heads'],
                    'head_dim': kv_info['head_dim']
                }
            }
        except Exception as e:
            return {'error': str(e)}


# -------------------------------------------------
# 🌟 Gradio 界面构建（保持原逻辑）
# -------------------------------------------------

def create_gradio_interface():
    calculator = LLMMemoryCalculator()

    def calculate_memory(model_id, gpu_memory, num_gpus, context_length, utilization_rate):
        if not model_id.strip():
            return "请输入模型ID"

        try:
            result = calculator.calculate_memory_requirements(
                model_id.strip(),
                float(gpu_memory),
                int(num_gpus),
                int(context_length),
                float(utilization_rate) / 100
            )

            if 'error' in result:
                return f"❌ 错误: {result['error']}"

            status = "✅ 可以运行" if result['is_feasible'] else "❌ 显存不足"

            output = f"""
## 模型分析结果

**模型**: {result['model_id']}
**状态**: {status}

### 📊 内存分析
- **模型大小**: {result['model_size_gb']} GB ({result['size_source']})
- **KV Cache**: {result['kv_cache_gb']} GB
- **其他开销**: {result['other_overhead_gb']} GB
- **总需求**: {result['total_memory_needed_gb']} GB
- **可用显存**: {result['available_memory_gb']} GB
- **剩余显存**: {result['memory_margin_gb']} GB

### 🔧 模型配置
- **注意力类型**: {result['attention_type']}
- **层数**: {result['config_info']['num_layers']}
- **隐藏维度**: {result['config_info']['hidden_size']}
- **注意力头数**: {result['config_info']['num_attention_heads']}
- **KV头数**: {result['config_info']['num_kv_heads']}
- **头维度**: {result['config_info']['head_dim']}

### 💾 GPU使用情况
- **每GPU内存**: {result['memory_per_gpu_gb']} GB
- **每GPU利用率**: {result['utilization_per_gpu']}%

### 💡 建议
"""
            if result['is_feasible']:
                output += f"✅ 当前配置可以成功运行该模型。剩余 {result['memory_margin_gb']} GB 显存。"
            else:
                needed_extra = abs(result['memory_margin_gb'])
                output += f"❌ 需要额外 {needed_extra} GB 显存才能运行。\n建议：\n- 增加GPU数量\n- 使用更大显存的GPU\n- 减少上下文长度\n- 使用模型量化（如int8/int4）"

            return output
        except Exception as e:
            return f"❌ 计算出错: {str(e)}"

    with gr.Blocks(title="LLM GPU内存计算器", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🚀 LLM GPU内存需求计算器")
        gr.Markdown("输入模型信息和硬件配置，计算是否能够成功运行大语言模型")

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## 📝 输入参数")

                model_id = gr.Textbox(label="🤗 Hugging Face 模型ID",
                                      placeholder="例如: deepseek-ai/DeepSeek-R1-0528-Qwen3-8B",
                                      value="deepseek-ai/DeepSeek-R1-0528-Qwen3-8B")

                with gr.Row():
                    gpu_memory = gr.Number(label="💾 单张GPU显存 (GB)", value=24, minimum=1, maximum=1000)
                    num_gpus = gr.Number(label="🔢 GPU数量", value=1, minimum=1, maximum=64, precision=0)

                with gr.Row():
                    context_length = gr.Number(label="📏 上下文长度", value=16384, minimum=512, maximum=1000000, precision=0)
                    utilization_rate = gr.Slider(label="⚡ 显存利用率 (%)", minimum=50, maximum=95, value=90, step=5)

                calculate_btn = gr.Button("🔍 计算内存需求", variant="primary")

            with gr.Column(scale=2):
                gr.Markdown("## 📊 计算结果")
                output = gr.Markdown("点击计算按钮开始分析...")

        calculate_btn.click(fn=calculate_memory,
                            inputs=[model_id, gpu_memory, num_gpus, context_length, utilization_rate],
                            outputs=output)

        gr.Markdown("""
## 📚 使用示例

**小型模型**: `microsoft/DialoGPT-medium`
**中型模型**: `microsoft/DialoGPT-large`  
**大型模型**: `meta-llama/Llama-2-7b-hf`
**超大模型**: `meta-llama/Llama-2-13b-hf`

注意：某些模型可能需要申请访问权限。
        """)

    return demo


if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch(share=True, debug=True)