Spaces:
Running on Zero
Running on Zero
File size: 4,633 Bytes
4fa0580 edde0c0 4fa0580 edde0c0 4fa0580 edde0c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | """Pinned ComfyUI core with checkpoint loading before GPU requests."""
from pathlib import Path
import subprocess
import sys
import gc
ROOT = Path(__file__).resolve().parent
COMFY_COMMIT = "65787d668397d230bf5839d69a0a7239e2dad378"
MODEL_ID = "Noctaluna/Noct-Q-Uncensored-Qwen-Image-2.1"
MODEL_REVISION = "a81b9af51120a78e285e57906f2250a2a02080e9"
COMPANION_ID = "Comfy-Org/Qwen-Image-2.1"
COMPANION_REVISION = "cb504a4090723e43f17ad01cec0359490e2de613"
CHECKPOINT = "NoctQ_V4_int8_convrot.safetensors"
def prepare_source() -> Path:
"""Fetch one pinned ComfyUI revision; never execute custom nodes."""
target = ROOT / "ComfyUI"
if not target.exists():
target.mkdir()
subprocess.run(["git", "init", str(target)], check=True)
subprocess.run(["git", "-C", str(target), "remote", "add", "origin",
"https://github.com/Comfy-Org/ComfyUI.git"], check=True)
subprocess.run(["git", "-C", str(target), "fetch", "--depth", "1",
"origin", COMFY_COMMIT], check=True)
subprocess.run(["git", "-C", str(target), "checkout", "--detach",
"FETCH_HEAD"], check=True)
actual = subprocess.check_output(
["git", "-C", str(target), "rev-parse", "HEAD"], text=True).strip()
if actual != COMFY_COMMIT:
raise RuntimeError("ComfyUI revision does not match the pinned runtime.")
sys.path.insert(0, str(target))
return target
def load_models():
"""Load the exact INT8 convrot files before the first GPU request."""
# app.py must have imported spaces before calling this function.
import torch
from huggingface_hub import hf_hub_download
prepare_source()
import comfy.options
comfy.options.enable_args_parsing()
saved_argv = sys.argv[:]
# Initialize Comfy's management module on CPU: it otherwise probes a
# numeric CUDA device before ZeroGPU has allocated a worker.
sys.argv = [sys.argv[0], "--cpu", "--disable-dynamic-vram",
"--use-pytorch-cross-attention"]
try:
import comfy.model_management as mm
import comfy.sd as sd
import comfy.utils
import nodes
from comfy_extras.nodes_qwen import TextEncodeQwenImage21
finally:
sys.argv = saved_argv
# Use a device without an index: ZeroGPU assigns it per request.
cuda = torch.device("cuda")
mm.args.cpu = False
mm.args.gpu_only = True
mm.cpu_state = mm.CPUState.GPU
mm.vram_state = mm.VRAMState.HIGH_VRAM
mm.get_torch_device = lambda: torch.device("cuda")
def download(repo: str, filename: str, revision: str) -> str:
return hf_hub_download(repo, filename, revision=revision)
print("Loading Noct Q V4 INT8 convrot diffusion model", flush=True)
model = sd.load_diffusion_model(
download(MODEL_ID, CHECKPOINT, MODEL_REVISION),
model_options={"dtype": torch.bfloat16, "load_device": cuda,
"offload_device": torch.device("cpu")},
disable_dynamic=True,
)
if model is None:
raise RuntimeError("ComfyUI did not recognize the Noct Q checkpoint.")
# Comfy's quantized tensor subclasses keep internal INT8 data and
# rotation/scaling tensors. Keep those on CPU until the real GPU worker
# moves them: ZeroGPU's eager fake-CUDA packing does not preserve them.
model.offload_device = cuda
gc.collect()
print("Loading Qwen 3 VL INT8 convrot text encoder", flush=True)
clip = sd.load_clip(
[download(COMPANION_ID,
"text_encoders/qwen3vl_8b_int8_convrot.safetensors",
COMPANION_REVISION)],
clip_type=sd.CLIPType.QWEN_IMAGE,
model_options={"dtype": torch.bfloat16, "load_device": cuda,
"offload_device": torch.device("cpu"),
"initial_device": torch.device("cpu")},
disable_dynamic=True,
)
clip.patcher.offload_device = cuda
gc.collect()
print("Loading Qwen Image 2.1 BF16 VAE", flush=True)
weights, metadata = comfy.utils.load_torch_file(
download(COMPANION_ID, "vae/qwen_image_2.1_vae_bf16.safetensors",
COMPANION_REVISION), return_metadata=True)
vae = sd.VAE(sd=weights, device=cuda, dtype=torch.bfloat16,
metadata=metadata)
del weights
# Keep all three components together on the ordinary Comfy loading path.
vae.patcher.offload_device = cuda
gc.collect()
print("Noct Q V4, text encoder and VAE loaded; ready for GPU requests.",
flush=True)
return model, clip, vae, nodes, TextEncodeQwenImage21
|