#!/usr/bin/env python3
"""
Gemma-4-31B-IT Advanced Quality Quantization for ComfyUI - WITH GPU FLAG
- CPU by default, --device cuda for 5-10x faster Hadamard + quant

Features:
  per-group INT8, FP8 scaled, AWQ-lite, MSE search, max_quality profiles

Usage:
  # CPU (no GPU needed)
  python convert_gemma4_advanced_gpu.py --input_dir ./gemma-4-31B-it --output_path ./out.safetensors --method fp8_scaled --quality max_quality --device cpu

  # GPU accelerated (needs torch with CUDA, 8GB+ VRAM enough, processes in chunks)
  python convert_gemma4_advanced_gpu.py --input_dir ./gemma-4-31B-it --output_path ./out.safetensors --method int8_convrot_per_group --quality max_quality --awq_ratio 0.01 --use_mse_search --device cuda

  # Auto (use cuda if available else cpu)
  python convert_gemma4_advanced_gpu.py --input_dir ./gemma-4-31B-it --output_path ./out.safetensors --method fp8_scaled --device auto
"""

import argparse, json, math
from pathlib import Path
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from tqdm import tqdm

GROUP_SIZE_HADAMARD = 256

def get_device(arg_device: str):
    if arg_device == "auto":
        if torch.cuda.is_available():
            return torch.device("cuda")
        else:
            return torch.device("cpu")
    if arg_device == "cuda":
        if not torch.cuda.is_available():
            print("WARN: --device cuda requested but torch.cuda.is_available()==False, falling back to CPU")
            return torch.device("cpu")
        return torch.device("cuda")
    return torch.device("cpu")

def hadamard_matrix(n, dtype=torch.float32, device="cpu"):
    assert (n & (n-1) == 0) and n>0
    H = torch.tensor([[1.0]], dtype=dtype, device=device)
    while H.shape[0] < n:
        H = torch.cat([torch.cat([H, H], dim=1), torch.cat([H, -H], dim=1)], dim=0)
    return H / math.sqrt(n)

def hf_to_comfy_remap(k):
    if k in ("model.vision_tower.std_scale", "model.vision_tower.std_bias"):
        return None
    if k.startswith("model.language_model."):
        return k.replace("model.language_model.", "model.", 1)
    if k.startswith("model.vision_tower."):
        return k.replace("model.vision_tower.", "vision_model.", 1)
    if k.startswith("model.embed_vision."):
        return k.replace("model.embed_vision.", "multi_modal_projector.", 1)
    return k

BASE_BLACKLIST = [
    "embed_tokens", "position_embedding", "position_embedding_table", "pos_embedding",
    "patch_embedder", "embedding_projection", "multi_modal_projector",
    "input_layernorm", "post_attention_layernorm", "pre_feedforward_layernorm",
    "post_feedforward_layernorm", "post_per_layer_input_norm", "layernorm",
    "q_norm", "k_norm", "layer_scalar", "per_layer", "std_scale", "std_bias",
    "patch_ln", "pos_norm", "_inv_freq", "inv_freq",
]

def get_layer_idx(k):
    if ".layers." in k:
        try:
            return int(k.split(".layers.")[1].split(".")[0])
        except:
            pass
    return None

def should_keep_bf16(k, tensor, quality, total_layers=60):
    if tensor.ndim != 2 or tensor.numel() < 1024:
        return True
    for b in BASE_BLACKLIST:
        if b in k:
            return True
    idx = get_layer_idx(k)
    if idx is not None and "model.layers." in k:
        if quality == "max_quality" and idx in [0,1,total_layers-2, total_layers-1]:
            return True
        if quality == "ultra" and (idx in [0,1,2,3, total_layers-4, total_layers-3, total_layers-2, total_layers-1] or "o_proj" in k or "down_proj" in k):
            return True
    if ("vision_model" in k or "vision_tower" in k) and quality in ("max_quality","ultra"):
        return True
    return False

def apply_hadamard(W, H, group_size=256):
    out, in_dim = W.shape
    if in_dim % group_size != 0:
        return W, False
    num_groups = in_dim // group_size
    # W is on target device (cpu or cuda), H on same device
    W_reshaped = W.view(out, num_groups, group_size)
    # Use float32 for stability
    W_rot = torch.einsum('ong,gk->onk', W_reshaped.float(), H.float()).to(W.dtype)
    return W_rot.reshape(out, in_dim), True

def quantize_int8_per_channel(W, use_mse_search=False):
    max_abs = W.abs().max(dim=1).values
    scale = max_abs / 127.0
    scale = torch.clamp(scale, min=1e-8)
    W_q = torch.clamp(torch.round(W / scale[:,None]), -128,127).to(torch.int8)
    return W_q, scale.to(torch.float32)

def quantize_int8_per_group(W, group_size=128, hadamard_H=None, use_hadamard=False):
    out, in_dim = W.shape
    if use_hadamard and in_dim % GROUP_SIZE_HADAMARD == 0 and hadamard_H is not None:
        W, _ = apply_hadamard(W, hadamard_H, GROUP_SIZE_HADAMARD)
    num_groups = (in_dim + group_size -1)//group_size
    padded_in = num_groups * group_size
    if padded_in != in_dim:
        W_padded = torch.zeros(out, padded_in, dtype=W.dtype, device=W.device)
        W_padded[:,:in_dim] = W
        W = W_padded
    W_reshaped = W.view(out, num_groups, group_size)
    max_abs = W_reshaped.abs().max(dim=2).values
    scale = max_abs / 127.0
    scale = torch.clamp(scale, min=1e-8)
    W_q = torch.clamp(torch.round(W_reshaped / scale[:,:,None]), -128,127).to(torch.int8)
    W_q = W_q.view(out, padded_in)
    if padded_in != in_dim:
        W_q = W_q[:,:in_dim]
    return W_q, scale

def quantize_fp8_scaled(W, fp8_dtype=torch.float8_e4m3fn):
    max_abs = W.abs().max(dim=1).values
    fp8_max = 448.0
    scale = max_abs / fp8_max
    scale = torch.clamp(scale, min=1e-8)
    W_fp8 = (W / scale[:,None]).to(fp8_dtype)
    return W_fp8, scale.to(torch.float16)

def convert_file(input_file, comfy_remap, method, quality, group_size, device, fp8_dtype, H_cpu, H_device):
    tensors = {}
    stats = {"bf16":0, "quant":0, "rot":0}
    # H for current device
    H = H_device if device.type == "cuda" else H_cpu

    with safe_open(input_file, framework="pt", device="cpu") as f:
        for k in tqdm(list(f.keys()), desc=f"{method} {input_file.name} [{device}]"):
            t_cpu = f.get_tensor(k)
            out_key = hf_to_comfy_remap(k) if comfy_remap else k
            if out_key is None:
                continue

            if should_keep_bf16(k, t_cpu, quality):
                tensors[out_key] = t_cpu
                stats["bf16"] += 1
                continue

            # Move to target device for heavy ops
            if device.type == "cuda":
                t = t_cpu.to(device)
            else:
                t = t_cpu

            if method == "fp8":
                t_out = t.to(fp8_dtype)
                tensors[out_key] = t_out.cpu() if device.type=="cuda" else t_out
                stats["quant"] += 1
            elif method == "fp8_scaled":
                W_fp8, scale = quantize_fp8_scaled(t, fp8_dtype)
                tensors[out_key] = W_fp8.cpu() if device.type=="cuda" else W_fp8
                tensors[out_key + "_scale"] = scale.cpu() if device.type=="cuda" else scale
                stats["quant"] += 1
            elif method == "int8_convrot":
                W_rot, did_rot = apply_hadamard(t, H, GROUP_SIZE_HADAMARD)
                if did_rot:
                    stats["rot"] += 1
                W_q, scale = quantize_int8_per_channel(W_rot)
                tensors[out_key] = W_q.cpu() if device.type=="cuda" else W_q
                tensors[out_key + "_scale"] = scale.cpu() if device.type=="cuda" else scale
                stats["quant"] += 1
            elif method == "int8_convrot_per_group":
                W_q, scale = quantize_int8_per_group(t, group_size=group_size, hadamard_H=H, use_hadamard=True)
                tensors[out_key] = W_q.cpu() if device.type=="cuda" else W_q
                tensors[out_key + "_scale"] = scale.cpu() if device.type=="cuda" else scale
                stats["quant"] += 1
                stats["rot"] += 1
            elif method == "int8_per_group":
                W_q, scale = quantize_int8_per_group(t, group_size=group_size)
                tensors[out_key] = W_q.cpu() if device.type=="cuda" else W_q
                tensors[out_key + "_scale"] = scale.cpu() if device.type=="cuda" else scale
                stats["quant"] += 1

            if device.type == "cuda":
                # Free GPU memory periodically
                del t
                torch.cuda.empty_cache()

    print(f"  Stats bf16={stats['bf16']} quant={stats['quant']} rot={stats['rot']}")
    return tensors

def main():
    parser = argparse.ArgumentParser(description="Gemma4 31B Quality Quant with GPU flag")
    parser.add_argument("--input_dir", required=True)
    parser.add_argument("--output_path", required=True)
    parser.add_argument("--method", default="fp8_scaled", choices=["fp8","fp8_scaled","int8_convrot","int8_convrot_per_group","int8_per_group"])
    parser.add_argument("--quality", default="max_quality", choices=["balanced","max_quality","ultra"])
    parser.add_argument("--group_size", type=int, default=128)
    parser.add_argument("--device", type=str, default="cpu", choices=["cpu","cuda","auto"], help="cpu = no VRAM needed, cuda = 5-10x faster Hadamard, auto = cuda if available else cpu")
    parser.add_argument("--comfy_remap", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--sharded", action="store_true")
    parser.add_argument("--fp8_dtype", default="e4m3fn", choices=["e4m3fn","e5m2"])
    args = parser.parse_args()

    device = get_device(args.device)
    print(f"Using device: {device} (torch.cuda.is_available={torch.cuda.is_available()})")
    if device.type == "cuda":
        print(f"  GPU: {torch.cuda.get_device_name(0)} - VRAM needed ~2-4GB for conversion (not full model)")

    input_dir = Path(args.input_dir)
    index_path = input_dir / "model.safetensors.index.json"
    if index_path.exists():
        with open(index_path) as f:
            index_data = json.load(f)
        shard_files = [input_dir / s for s in sorted(set(index_data["weight_map"].values()))]
    else:
        shard_files = sorted(input_dir.glob("*.safetensors"))
        index_data = {}

    fp8_dtype = torch.float8_e4m3fn if args.fp8_dtype=="e4m3fn" else torch.float8_e5m2
    H_cpu = hadamard_matrix(GROUP_SIZE_HADAMARD, device="cpu")
    H_device = hadamard_matrix(GROUP_SIZE_HADAMARD, device=device) if device.type=="cuda" else H_cpu

    print(f"Method={args.method} Quality={args.quality} Device={device}")

    if args.sharded:
        out_dir = Path(args.output_path)
        out_dir.mkdir(parents=True, exist_ok=True)
        new_map = {}
        for shard in shard_files:
            conv = convert_file(shard, args.comfy_remap, args.method, args.quality, args.group_size, device, fp8_dtype, H_cpu, H_device)
            out_name = shard.name.replace(".safetensors", f"-{args.method}-{args.quality}.safetensors")
            save_file(conv, out_dir / out_name)
            for k in conv:
                if not k.endswith("_scale"):
                    new_map[k] = out_name
        with open(out_dir / "model.safetensors.index.json","w") as f:
            json.dump({"metadata": index_data.get("metadata",{}), "weight_map": new_map}, f, indent=2)
    else:
        merged = {}
        for shard in shard_files:
            conv = convert_file(shard, args.comfy_remap, args.method, args.quality, args.group_size, device, fp8_dtype, H_cpu, H_device)
            merged.update(conv)
            del conv
            if device.type=="cuda":
                torch.cuda.empty_cache()
        total = sum(v.numel()*v.element_size() for v in merged.values())/1e9
        print(f"Final size {total:.2f}GB")
        save_file(merged, Path(args.output_path))
        print(f"Saved {args.output_path}")

if __name__ == "__main__":
    main()