Add Colab int8 (ConvRot) quantization pipeline script

tools/fullgreed_int8_quantize.py

@@ -0,0 +1,124 @@
+#!/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.")