#!/usr/bin/env python3
"""
End-to-end INT8 (ConvRot) quantization of John2386/fullgreed for ComfyUI.

RUN THIS ON A CUDA GPU (Colab / RunPod / any NVIDIA box with a terminal).
It will NOT work on a Mac — it needs CUDA + Triton.

Recommended GPU: 24 GB+ VRAM (3090/4090/L4/A100). A 16 GB T4 may OOM on a ~6B model
because on-the-fly quant keeps an extra INT8 copy in memory (see the node's README).

What it does:
  1. installs ComfyUI + ComfyUI-INT8-Fast + Triton
  2. downloads fullgreed_f16.safetensors from your HF repo
  3. runs the OTUNetLoaderW8A8 -> INT8ModelSave quantize workflow headlessly (ConvRot, z-image)
  4. converts the result to ComfyUI's native int8 format (convert_to_comfy.py)
  5. uploads fullgreed_i8_comfy.safetensors back to John2386/fullgreed

Set your HF write token below (or `export HF_TOKEN=...`).
"""
import os, sys, subprocess, time, glob, json

HF_TOKEN = os.environ.get("HF_TOKEN", "PASTE_HF_WRITE_TOKEN_HERE")
HF_REPO  = "John2386/fullgreed"
BASE_FILE = "fullgreed_f16.safetensors"
OUT_NAME  = "fullgreed_i8_comfy.safetensors"

WORK  = "/content" if os.path.isdir("/content") else os.getcwd()
COMFY = os.path.join(WORK, "ComfyUI")
NODE  = os.path.join(COMFY, "custom_nodes", "ComfyUI-INT8-Fast")

def run(cmd):
    print("\n+ " + cmd, flush=True)
    subprocess.run(cmd, shell=True, check=True)

# ---- 0) GPU sanity + Triton version (T4/20-series are sm75: need triton==3.2.0) ----
import torch
assert torch.cuda.is_available(), "No CUDA GPU visible — set Runtime > Change runtime type > GPU"
cap = torch.cuda.get_device_capability(0)
name = torch.cuda.get_device_name(0)
vram = torch.cuda.get_device_properties(0).total_memory / 1e9
print(f"GPU: {name}  sm{cap[0]}{cap[1]}  {vram:.1f} GB VRAM", flush=True)
if vram < 20:
    print("WARNING: <20 GB VRAM — a ~6B model may OOM during on-the-fly quantization.", flush=True)
TRITON = "triton==3.2.0" if cap == (7, 5) else "triton"   # sm75 support dropped in triton 3.3

# ---- 1) ComfyUI + node + deps -------------------------------------------------
if not os.path.isdir(COMFY):
    run(f"git clone --depth 1 https://github.com/comfyanonymous/ComfyUI '{COMFY}'")
run(f"pip install -q -r '{COMFY}/requirements.txt'")
run(f"pip install -q {TRITON} safetensors huggingface_hub requests")
if not os.path.isdir(NODE):
    run(f"git clone --depth 1 https://github.com/BobJohnson24/ComfyUI-INT8-Fast '{NODE}'")
if os.path.exists(f"{NODE}/requirements.txt"):
    run(f"pip install -q -r '{NODE}/requirements.txt'")

# ---- 2) download the fp16 base model -----------------------------------------
from huggingface_hub import hf_hub_download, upload_file
dst = os.path.join(COMFY, "models", "diffusion_models")
os.makedirs(dst, exist_ok=True)
print("downloading base model from HF...", flush=True)
hf_hub_download(HF_REPO, BASE_FILE, local_dir=dst)

# ---- 3) start ComfyUI headless + queue the quantize prompt --------------------
import requests
print("starting ComfyUI (headless)...", flush=True)
server = subprocess.Popen(
    [sys.executable, "main.py", "--listen", "127.0.0.1", "--port", "8188",
     "--disable-auto-launch"], cwd=COMFY)
for _ in range(180):                       # wait up to ~6 min for startup
    try:
        if requests.get("http://127.0.0.1:8188/system_stats", timeout=2).ok:
            break
    except Exception:
        pass
    time.sleep(2)
else:
    raise SystemExit("ComfyUI did not come up — check the log above for node/Triton errors.")

# API-format graph: loader (on-the-fly int8 + convrot, z-image) -> save
prompt = {
    "1": {"class_type": "OTUNetLoaderW8A8", "inputs": {
        "unet_name": BASE_FILE,
        "weight_dtype": "default",
        "model_type": "z-image",
        "on_the_fly_quantization": True,
        "enable_convrot": True,
        "lora_mode": "None"}},
    "2": {"class_type": "INT8ModelSave", "inputs": {
        "model": ["1", 0],
        "filename_prefix": "int8_models/fullgreed_i8"}},
}
r = requests.post("http://127.0.0.1:8188/prompt", json={"prompt": prompt})
if not r.ok:
    raise SystemExit(f"/prompt rejected: {r.status_code} {r.text}")
pid = r.json()["prompt_id"]
print("queued prompt", pid, "- quantizing (this is the slow part)...", flush=True)

while True:                                # poll until the job leaves the queue
    h = requests.get(f"http://127.0.0.1:8188/history/{pid}", timeout=10).json()
    if pid in h:
        st = h[pid].get("status", {})
        print("job finished:", st.get("status_str", st), flush=True)
        if st.get("status_str") == "error":
            raise SystemExit("Quantization errored — see ComfyUI log above.")
        break
    time.sleep(5)

# ---- 4) convert I8Fast -> native ComfyUI int8 --------------------------------
cands = sorted(glob.glob(os.path.join(COMFY, "output", "int8_models", "fullgreed_i8*.safetensors")))
if not cands:
    raise SystemExit("No int8 output file was produced.")
i8fast = cands[-1]
out_path = os.path.join(COMFY, "output", OUT_NAME)
print("produced:", i8fast)
run(f"python '{NODE}/convert_to_comfy.py' '{i8fast}' '{out_path}'")

# ---- 5) upload back to HF -----------------------------------------------------
print("uploading to HF...", flush=True)
url = upload_file(path_or_fileobj=out_path, path_in_repo=OUT_NAME,
                  repo_id=HF_REPO, repo_type="model", token=HF_TOKEN,
                  commit_message="Add INT8 (ConvRot) ComfyUI quant of fullgreed")
print("\nDONE ->", url)
print("Load it in ComfyUI with the native INT8 loader (or the INT8-Fast loader with "
      "on_the_fly_quantization OFF), plus Qwen3-4B text encoder + Flux 16ch VAE.")