import gradio as gr
import requests


# 定数定義
API_BASE_URL = "https://huggingface.co/api/models/"
MULTIPLIERS = {
    "I8": 1,
    "U8": 1,
    "F8_E5M2": 1,
    "F8_E4M3": 1,
    "I16": 2,
    "U16": 2,
    "F16": 2,
    "BF16": 2,
    "I32": 4,
    "U32": 4,
    "F32": 4,
    "I64": 8,
    "U64": 8,
    "F64": 8,
}
PRECISION_KEYS = list(MULTIPLIERS.keys())
COMMENT = {
    "I8": "INT8の処理能力が高いGPUを選定してください。",
    "U8": "INT8の処理能力が高いGPUを選定してください。",
    "F8_E5M2": "FP8の処理能力が高いGPUを選定してください。",
    "F8_E4M3": "FP8の処理能力が高いGPUを選定してください。",
    "I16": "INT16の処理能力が高いGPUを選定してください。",
    "U16": "INT16の処理能力が高いGPUを選定してください。",
    "F16": "FP16の処理能力が高いGPUを選定してください。",
    "BF16": "BF16の処理能力が高いGPUを選定してください。",
    "I32": "INT32の処理能力が高いGPUを選定してください。",
    "U32": "INT32の処理能力が高いGPUを選定してください。",
    "F32": "FP32の処理能力が高いGPUを選定してください。",
    "I64": "INT64の処理能力が高いGPUを選定してください。",
    "U64": "INT64の処理能力が高いGPUを選定してください。",
    "F64": "FP64の処理能力が高いGPUを選定してください。",
}

def get_model_api_info(model_name: str) -> dict:
    """
    Hugging Face Hub APIから指定モデルの情報を取得する。

    :param model_name: モデル名（例: "meta-llama/Llama-3.3-70B-Instruct"）
    :return: モデル情報の辞書。取得に失敗した場合はNoneを返す。
    """
    api_url = f"{API_BASE_URL}{model_name}"
    try:
        response = requests.get(api_url)
        response.raise_for_status()
        return response.json()
    except requests.RequestException:
        return None

def sum_precision_sizes(parameters: dict) -> int:
    """
    safetensors内の各精度のパラメータ数に対して、定数の乗数をかけた合計バイト数を算出する。

    :param parameters: 各精度のパラメータ数を含む辞書
    :return: 合計バイト数
    """
    # 各データ型とその乗数の対応表
    total_bytes = 0
    for precision, multiplier in MULTIPLIERS.items():
        count = parameters.get(precision, 0)
        total_bytes += count * multiplier
    return total_bytes

def estimate_gpu_memory(model_name: str) -> str:
    """
    指定したモデル名からAPI情報を取得し、safetensors内の各精度パラメータサイズの合算値から
    GPUメモリ必要量を概算する。

    :param model_name: モデル名
    :return: GPUメモリ必要量などの情報を含むメッセージ文字列
    """
    result_lines = []

    model_info = get_model_api_info(model_name)
    if model_info is None:
        result_lines.append(f"エラー: モデル '{model_name}' の情報が取得できませんでした。")
        return "\n".join(result_lines)

    # リポジトリ全体のファイルサイズ（参考）
    usedStorage = model_info.get("usedStorage")
    usedStorage_gb = usedStorage / (1024 ** 3) if usedStorage else None

    safetensors = model_info.get("safetensors")

    # safetensorsの情報が見つからない場合
    if not safetensors or "parameters" not in safetensors:
        result_lines.append("safetensorsの情報が見つかりません。")
        if usedStorage_gb is not None:
            result_lines.append("")
            result_lines.append(f"参考までに、該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GBです。")
            result_lines.append("")
            result_lines.append("これを全てモデルのデータとして仮定した場合、推論及び、トレーニングに必要な概算のGPUメモリサイズは")
            result_lines.append(f"【推論】約 {usedStorage_gb * 1.5:.2f} GB")
            result_lines.append(f"【トレーニング（LoRA）】約 {usedStorage_gb * 2:.2f} GB")
            result_lines.append(f"【トレーニング（フルパラメータ）】約 {usedStorage_gb * 4:.2f} GB")
            result_lines.append("となります。")
        return "\n".join(result_lines)

    # safetensorsの情報がある場合
    parameters = safetensors["parameters"]
    estimated_bytes = sum_precision_sizes(parameters)
    estimated_gb = estimated_bytes / (1024 ** 3)

    result_lines.append(f"モデル '{model_name}' のsafetensors情報より、各パラメータ数は")

    # 各精度ごとのパラメータ数の出力
    max_precision = None
    max_count = 0
    for precision in PRECISION_KEYS:
        count = parameters.get(precision, 0)
        if count > max_count:
            max_precision = precision
            max_count = count
        if precision in parameters:
            result_lines.append(f"【{precision}】 {parameters[precision]:,}")

    result_lines.append(f"これらを合算するとモデルのデータサイズは約 {estimated_gb:.2f} GB です。")
    result_lines.append("")
    result_lines.append("推論及び、トレーニングに必要な概算のGPUメモリサイズは")
    result_lines.append(f"【推論】約 {estimated_gb * 1.5:.2f} GB")
    result_lines.append(f"【トレーニング（LoRA）】約 {estimated_gb * 2:.2f} GB")
    result_lines.append(f"【トレーニング（フルパラメータ）】約 {estimated_gb * 4:.2f} GB")
    result_lines.append("となります。")

    # GPU選定のコメント
    if max_precision is not None:
        comment_message = COMMENT.get(max_precision, "")
        result_lines.append(comment_message)

    if usedStorage_gb is not None:
        result_lines.append("")
        result_lines.append(f"参考: 該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GB")

    return "\n".join(result_lines)

# Gradio インターフェースの定義
iface = gr.Interface(
    fn=estimate_gpu_memory,
    inputs=gr.Textbox(label="モデル名 (例: meta-llama/Llama-3.3-70B-Instruct)"),
    outputs="text",
    title="Model memory estimator",
    description=("Hugging Face Hub APIから取得したsafetensorsの情報をもとに、すべてのパラメータのデータサイズを合計し、その合計値から必要なGPUメモリサイズを概算（GB単位）で計算します。")
)

iface.launch()