ModelSizer-Vision-v1.py

import argparse
import math
from transformers import AutoConfig

# ---------------------------
# Argument parsing
# ---------------------------
parser = argparse.ArgumentParser(description="VLM Model Sizer")
parser.add_argument("--model_repo", type=str, default="Qwen/Qwen3-VL-32B-Instruct")
parser.add_argument("--kv_dtype", type=str, choices=["fp8", "fp16", "fp32"], default="fp8")
parser.add_argument("--text_tokens", type=int, help="Text tokens (defaults to model max_position_embeddings)")
parser.add_argument("--image_size", type=str, help="HxW in pixels, e.g., 1024x1024; defaults to model maximum")
parser.add_argument("--frames", type=int, default=1, help="Video frames; use 1 for images")
parser.add_argument("--trust_remote_code", action="store_true")
args = parser.parse_args()

model_repo = args.model_repo
kv_cache_dtype = args.kv_dtype
user_text_tokens = args.text_tokens
user_image_size = args.image_size
num_frames = max(1, int(args.frames))
trust_remote_code = args.trust_remote_code

# ---------------------------
# Config loading
# ---------------------------
def load_model_config(repo: str, trust: bool):
    cfg = AutoConfig.from_pretrained(repo, trust_remote_code=trust)
    text_cfg = getattr(cfg, "text_config", cfg)
    vision_cfg = getattr(cfg, "vision_config", None)
    return cfg, text_cfg, vision_cfg

root_cfg, text_cfg, vision_cfg = load_model_config(model_repo, trust_remote_code)

# ---------------------------
# Helpers (units, dtype maps)
# ---------------------------
def bytes_to_gib(n_bytes: int | float) -> float:
    return float(n_bytes) / (1024 ** 3)

def get_bytes_per_element_for_kv(dtype_str: str) -> int:
    if dtype_str == "fp8": return 1
    if dtype_str == "fp16": return 2
    if dtype_str == "fp32": return 4
    return 1

def bytes_per_param_from_dtype(dtype_str: str | None) -> float:
    s = (dtype_str or "").lower()
    table = {
        "fp32": 4.0, "float32": 4.0, "bf32": 4.0, "bfloat32": 4.0, "torch.bfloat32": 4.0,
        "fp16": 2.0, "float16": 2.0, "bf16": 2.0, "bfloat16": 2.0, "torch.bfloat16": 2.0,
        "fp8": 1.0, "e4m3": 1.0, "e5m2": 1.0,
        "mxfp8": 1.0 + 1/32,
        "fp4": 0.5,
        "mxfp4": 0.5 + 1/32,
        "nvfp4": 0.5 + 1/16,
    }
    return table.get(s, 2.0)  # default BF16

# ---------------------------
# Text (LLM) getters
# ---------------------------
def get_text_hidden_size() -> int:
    return int(getattr(text_cfg, "hidden_size", getattr(text_cfg, "d_model", getattr(text_cfg, "n_embd", 0))))

def get_text_layers() -> int:
    return int(getattr(text_cfg, "num_hidden_layers", 0))

def get_text_attention_heads() -> int:
    return int(getattr(text_cfg, "num_attention_heads", getattr(text_cfg, "n_head", 0)))

def get_text_kv_heads() -> int:
    val = (
        getattr(text_cfg, "num_key_value_heads", None)
        or getattr(text_cfg, "num_kv_heads", None)
        or getattr(text_cfg, "n_kv_heads", None)
        or getattr(text_cfg, "n_head_kv", None)
    )
    if val is not None:
        return int(val)
    return get_text_attention_heads()

def get_text_head_dim() -> int:
    hd = (
        getattr(text_cfg, "head_dim", None)
        or getattr(text_cfg, "attention_head_size", None)
        or getattr(text_cfg, "dim_head", None)
        or getattr(text_cfg, "headdim", None)
    )
    if hd is not None:
        return int(hd)
    h = get_text_hidden_size()
    n = get_text_attention_heads()
    return int(h // n) if h and n else 0

def get_vocab_size() -> int:
    return int(getattr(text_cfg, "vocab_size", 0))

def uses_tied_embeddings() -> bool:
    return bool(getattr(root_cfg, "tie_word_embeddings", True))

def is_gated_mlp_text() -> bool:
    act = (getattr(text_cfg, "hidden_act", "") or getattr(text_cfg, "activation_function", "") or "").lower()
    if "glu" in act:
        return True
    model_type = (getattr(text_cfg, "model_type", "") or getattr(root_cfg, "model_type", "")).lower()
    return any(k in model_type for k in ["llama", "mistral", "qwen", "phi", "gemma"])

def get_text_intermediate_size() -> int:
    return int(getattr(text_cfg, "intermediate_size", 0))

def get_text_tokens() -> int:
    if user_text_tokens is not None:
        return int(user_text_tokens)
    max_pos = getattr(text_cfg, "max_position_embeddings", None)
    return int(max_pos) if max_pos is not None else 0

# ---------------------------
# Vision getters
# ---------------------------
def has_vision() -> bool:
    return vision_cfg is not None

def get_vision_hidden_size() -> int:
    return int(getattr(vision_cfg, "hidden_size", 0)) if has_vision() else 0

def get_vision_depth() -> int:
    return int(getattr(vision_cfg, "depth", 0)) if has_vision() else 0

def get_vision_heads() -> int:
    return int(getattr(vision_cfg, "num_heads", 0)) if has_vision() else 0

def get_vision_intermediate_size() -> int:
    return int(getattr(vision_cfg, "intermediate_size", 0)) if has_vision() else 0

def get_vision_in_channels() -> int:
    return int(getattr(vision_cfg, "in_channels", 3)) if has_vision() else 3

def get_vision_patch_size() -> int:
    return int(getattr(vision_cfg, "patch_size", 16)) if has_vision() else 16

def get_vision_temporal_patch_size() -> int:
    return int(getattr(vision_cfg, "temporal_patch_size", 1)) if has_vision() else 1

def get_vision_spatial_merge_size() -> int:
    return int(getattr(vision_cfg, "spatial_merge_size", 1)) if has_vision() else 1

def get_vision_num_pos_embeddings() -> int:
    return int(getattr(vision_cfg, "num_position_embeddings", 0)) if has_vision() else 0

def get_vision_out_hidden_size() -> int:
    return int(getattr(vision_cfg, "out_hidden_size", get_text_hidden_size())) if has_vision() else get_text_hidden_size()

# ---------------------------
# Image size selection
# ---------------------------
def parse_user_image_size(hxw: str | None) -> tuple[int, int]:
    if not hxw:
        return None
    parts = hxw.lower().replace("×", "x").split("x")
    if len(parts) != 2:
        return None
    try:
        h = int(parts[0]); w = int(parts[1])
        return (h, w)
    except ValueError:
        return None

def get_model_max_image_size() -> tuple[int, int]:
    """
    Returns a square image size that maximizes visual tokens without exceeding the model's
    positional capacity after patching, temporal grouping, and spatial merging.
    Priority:
      1) explicit vision_config.image_size (int or [H,W]) if present
      2) derive from num_position_embeddings, patch_size, temporal_patch_size, spatial_merge_size
      3) fallback to 1024x1024
    """
    if not has_vision():
        return (1024, 1024)

    # 1) explicit field
    explicit = getattr(vision_cfg, "image_size", None) or getattr(root_cfg, "image_size", None)
    if explicit is not None:
        if isinstance(explicit, int):
            return (int(explicit), int(explicit))
        if isinstance(explicit, (list, tuple)) and len(explicit) >= 2:
            return (int(explicit[0]), int(explicit[1]))

    # 2) derive from positional capacity
    pos_capacity = get_vision_num_pos_embeddings()
    if pos_capacity > 0:
        patch = max(1, get_vision_patch_size())
        merge = max(1, get_vision_spatial_merge_size())
        temporal = max(1, get_vision_temporal_patch_size())
        time_groups = math.ceil(num_frames / temporal)

        # tokens allowed per frame before spatial merge division
        max_tokens_per_frame = max(1, (pos_capacity * (merge ** 2)) // time_groups)
        grid_side = int(math.floor(math.sqrt(max_tokens_per_frame)))
        max_side_px = grid_side * patch
        if max_side_px > 0:
            return (max_side_px, max_side_px)

    # 3) conservative default
    return (1024, 1024)

def get_effective_image_size() -> tuple[int, int]:
    """
    If the user supplies --image_size, use it.
    Otherwise, default to the model's maximum image size.
    """
    parsed = parse_user_image_size(user_image_size)
    return parsed if parsed else get_model_max_image_size()

# ---------------------------
# Visual token estimation
# ---------------------------
def estimate_visual_token_count() -> int:
    if not has_vision():
        return 0

    img_h, img_w = get_effective_image_size()
    patch = max(1, get_vision_patch_size())
    temporal = max(1, get_vision_temporal_patch_size())
    merge = max(1, get_vision_spatial_merge_size())

    grid_h = math.ceil(img_h / patch)
    grid_w = math.ceil(img_w / patch)
    tokens_spatial = grid_h * grid_w

    time_groups = math.ceil(num_frames / temporal)
    total_tokens = tokens_spatial * time_groups

    total_tokens = math.ceil(total_tokens / (merge ** 2))

    # cap to positional capacity if present
    pos_capacity = get_vision_num_pos_embeddings()
    if pos_capacity > 0:
        total_tokens = min(total_tokens, pos_capacity)

    return total_tokens

# ---------------------------
# Parameter counting (text)
# ---------------------------
def compute_text_attention_params() -> int:
    d = get_text_hidden_size()
    H = get_text_attention_heads()
    Hkv = get_text_kv_heads()
    d_k = get_text_head_dim()
    if not (d and H and Hkv and d_k):
        return 0
    q_and_o = 2 * d * (H * d_k)
    k_and_v = 2 * d * (Hkv * d_k)
    return q_and_o + k_and_v

def compute_text_mlp_params_per_layer() -> int:
    d = get_text_hidden_size()
    d_ff = get_text_intermediate_size()
    if not (d and d_ff):
        return 0
    return (3 if is_gated_mlp_text() else 2) * d * d_ff

def compute_text_embedding_params() -> int:
    V = get_vocab_size()
    d = get_text_hidden_size()
    if not (V and d): return 0
    base = V * d
    return base if uses_tied_embeddings() else 2 * base

def compute_text_total_params() -> int:
    per_layer = compute_text_attention_params() + compute_text_mlp_params_per_layer()
    return get_text_layers() * per_layer + compute_text_embedding_params()

# ---------------------------
# Parameter counting (vision)
# ---------------------------
def compute_vision_patch_embed_params() -> int:
    if not has_vision(): return 0
    c = get_vision_in_channels()
    d_v = get_vision_hidden_size()
    ps = get_vision_patch_size()
    return (ps * ps * c) * d_v

def compute_vision_positional_params() -> int:
    if not has_vision(): return 0
    d_v = get_vision_hidden_size()
    n_pos = get_vision_num_pos_embeddings()
    return d_v * n_pos if n_pos > 0 else 0

def compute_vision_block_params_per_layer() -> int:
    if not has_vision(): return 0
    d_v = get_vision_hidden_size()
    d_ff_v = get_vision_intermediate_size()
    attn = 4 * d_v * d_v
    mlp = 2 * d_v * d_ff_v
    return attn + mlp

def compute_vision_projector_params() -> int:
    if not has_vision(): return 0
    d_v = get_vision_hidden_size()
    d_out = get_vision_out_hidden_size()
    return d_v * d_out

def compute_vision_total_params() -> int:
    if not has_vision(): return 0
    patch = compute_vision_patch_embed_params()
    pos = compute_vision_positional_params()
    blocks = get_vision_depth() * compute_vision_block_params_per_layer()
    proj = compute_vision_projector_params()
    return patch + pos + blocks + proj

# ---------------------------
# KV cache (text + visual tokens)
# ---------------------------
def compute_kv_cache_size_bytes() -> int:
    text_tokens = get_text_tokens()
    visual_tokens = estimate_visual_token_count()
    total_tokens = text_tokens + visual_tokens

    num_layers = get_text_layers()
    n_kv_heads = get_text_kv_heads()
    head_dim = get_text_head_dim()
    bytes_per_elem = get_bytes_per_element_for_kv(kv_cache_dtype)

    if not (total_tokens and num_layers and n_kv_heads and head_dim and bytes_per_elem):
        return 0

    per_token_per_layer_bytes = 2 * n_kv_heads * head_dim * bytes_per_elem
    return int(total_tokens * num_layers * per_token_per_layer_bytes)

# ---------------------------
# Weight sizes
# ---------------------------
def get_text_bytes_per_param() -> float:
    dtype = (getattr(text_cfg, "dtype", None) or getattr(root_cfg, "dtype", None))
    dtype_str = str(dtype) if dtype else ""
    qcfg = getattr(root_cfg, "quantization_config", {}) or {}
    quant = (qcfg.get("quant_method") or dtype_str or "").lower()
    return bytes_per_param_from_dtype(quant)

def get_vision_bytes_per_param() -> float:
    dtype = getattr(vision_cfg, "dtype", None) if has_vision() else None
    return bytes_per_param_from_dtype(dtype or "bf16")

def compute_weights_size_bytes() -> tuple[int, int, int]:
    text_params = compute_text_total_params()
    vision_params = compute_vision_total_params()
    text_bpp = get_text_bytes_per_param()
    vision_bpp = get_vision_bytes_per_param()
    text_bytes = int(text_params * text_bpp)
    vision_bytes = int(vision_params * vision_bpp)
    return text_bytes, vision_bytes, text_bytes + vision_bytes
 
# ---------------------------
# Compute and report
# ---------------------------
text_params = compute_text_total_params()
vision_params = compute_vision_total_params()
text_bytes, vision_bytes, total_weight_bytes = compute_weights_size_bytes()
kv_cache_bytes = compute_kv_cache_size_bytes()

eff_h, eff_w = get_effective_image_size()
visual_tokens = estimate_visual_token_count()
text_tokens = get_text_tokens()

print(
    f"Model: {model_repo}\n"
    f"Vision enabled: {has_vision()}\n"
    f"Effective image size: {eff_h}x{eff_w}, frames: {num_frames}\n"
    f"Estimated visual tokens: {visual_tokens}\n"
    f"Text tokens: {text_tokens}\n"
    f"KV dtype: {kv_cache_dtype}\n"
    f"---- Text (LLM) ----\n"
    f"Hidden size: {get_text_hidden_size()}, Layers: {get_text_layers()}, Heads: {get_text_attention_heads()}, KV heads: {get_text_kv_heads()}, Head dim: {get_text_head_dim()}\n"
    f"Intermediate size: {get_text_intermediate_size()}, Gated MLP: {is_gated_mlp_text()}\n"
    f"Vocab size: {get_vocab_size()}, Tied embeddings: {uses_tied_embeddings()}\n"
    f"Text parameters: {text_params/1e9:.2f} Billion\n"
    f"Bytes/param (text weights): {get_text_bytes_per_param():.5f}\n"
    f"Text weight size: {bytes_to_gib(text_bytes):.2f} GiB\n"
    f"---- Vision tower ----\n"
    f"Vision hidden size: {get_vision_hidden_size()}, Depth: {get_vision_depth()}, Heads: {get_vision_heads()}\n"
    f"Vision intermediate size: {get_vision_intermediate_size()}, Patch: {get_vision_patch_size()}, Temporal patch: {get_vision_temporal_patch_size()}, Merge: {get_vision_spatial_merge_size()}\n"
    f"Vision out_hidden_size (projector): {get_vision_out_hidden_size()}, Pos embeddings: {get_vision_num_pos_embeddings()}\n"
    f"Vision parameters: {vision_params/1e9:.3f} Billion\n"
    f"Bytes/param (vision weights): {get_vision_bytes_per_param():.5f}\n"
    f"Vision weight size: {bytes_to_gib(vision_bytes):.3f} GiB\n"
    f"---- Totals ----\n"
    f"Total parameters (text+vision): {(text_params+vision_params)/1e9:.2f} Billion\n"
    f"Total weight size: {bytes_to_gib(total_weight_bytes):.2f} GiB\n"
    f"Estimated KV cache size (text+vision tokens): {bytes_to_gib(kv_cache_bytes):.2f} GiB"
)