From 76597bdb4234fa671b4c0b7f0997a028a7bb8586 Mon Sep 17 00:00:00 2001 From: nyueki Date: Thu, 1 Jan 2026 07:24:15 -0500 Subject: [PATCH] initial commit --- README.md | 31 ++++++++++ __init__.py | 156 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 2 + 3 files changed, 189 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..00a8800 --- /dev/null +++ b/README.md @@ -0,0 +1,31 @@ +# ComfyUI Remote CLIP Loader + +A small ComfyUI custom node set to offload CLIP encoding to a remote machine. + +Motivation +I created this because I have a PC with 16GB VRAM and some large models must unload and reload the text encoder during generation, which causes inefficient generation times. I decided to run CLIP on my Mac as a remote encoder and keep my PC as the image generator. + +What this does +- Runs a simple TCP worker that accepts prompt encode requests and returns tensors (`cond`, `pooled`). +- Worker code and client loader are in [__init__.py](__init__.py). +- Key symbols: [`RemoteCLIPWorker`](__init__.py), [`RemoteCLIPWorker.start_worker`](__init__.py), [`RemoteCLIPLoader`](__init__.py), [`RemoteCLIPLoader.load_remote`](__init__.py), [`RemoteCLIPProxy`](__init__.py), [`pack_tensors`](__init__.py), [`unpack_tensors`](__init__.py). + +Tested with +- Flux.2 Dev +- Qwen Image +- Wan2.2 + +Requirements +- See [requirements.txt](requirements.txt). At minimum you need PyTorch and numpy. + +Installation +1. Go to your ComfyUI installation's custom_nodes folder: + cd /path/to/ComfyUI/custom_nodes +2. Clone this repository: + git clone +3. Restart ComfyUI. + +Usage +- On the machine that will serve CLIP (your Mac), add the "Remote CLIP Worker" node (see [`RemoteCLIPWorker`](__init__.py)) and run it (set the CLIP model and listen port). +- On your generator PC, add the "Remote CLIP Loader" node (see [`RemoteCLIPLoader`](__init__.py)), point it to the worker's IP and port, and connect it as you would a normal CLIP model. +- The protocol and tensor packing are implemented by [`pack_tensors`](__init__.py) / [`unpack_tensors`](__init__.py). \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..8086c77 --- /dev/null +++ b/__init__.py @@ -0,0 +1,156 @@ +import socket +import struct +import threading +import json +import torch +import time + +def log(msg): + print(f"[RemoteCLIP {time.strftime('%H:%M:%S')}] {msg}", flush=True) + +def recv_exact(sock, size): + buf = b"" + while len(buf) < size: + chunk = sock.recv(size - len(buf)) + if not chunk: + raise ConnectionError("Socket closed") + buf += chunk + return buf + +def send_packet(sock, meta, blob): + meta_bytes = json.dumps(meta).encode("utf-8") + sock.sendall(struct.pack(">Q", len(meta_bytes))) + sock.sendall(meta_bytes) + sock.sendall(blob) + +def recv_packet(sock): + meta_size = struct.unpack(">Q", recv_exact(sock, 8))[0] + meta = json.loads(recv_exact(sock, meta_size).decode("utf-8")) + blob = recv_exact(sock, meta.get("blob_size", 0)) + return meta, blob + +def pack_tensors(tensors): + meta = {} + blobs = [] + for name, t in tensors.items(): + if t is None: + log(f"Warning: tensor '{name}' is None, replacing with zeros") + if name == "cond": + t = torch.zeros((1, 77, 768)) + else: + t = torch.zeros((1, 768)) + t = t.contiguous().cpu() + raw = t.numpy().tobytes() + meta[name] = {"dtype": str(t.dtype), "shape": list(t.shape), "size": len(raw)} + blobs.append(raw) + return meta, b"".join(blobs) + +def unpack_tensors(meta, blob): + out = {} + offset = 0 + for name, info in meta.items(): + size = info["size"] + raw = blob[offset:offset + size] + offset += size + dtype = getattr(torch, info["dtype"].split(".")[-1]) + t = torch.frombuffer(raw, dtype=dtype).clone() + out[name] = t.reshape(info["shape"]) + return out + +class RemoteCLIPProxy: + def __init__(self, ip, port): + self.ip = ip + self.port = port + self.sock = socket.create_connection((ip, port), timeout=60) + log(f"Connected to worker {ip}:{port}") + + def clone(self): + return self + + def add_patches(self, *_, **__): + return self + + def tokenize(self, text, **kwargs): + return {"text": text, "kwargs": kwargs} + + def _infer(self, tokens): + meta = {"cmd": "encode", "text": tokens["text"], "kwargs": tokens["kwargs"], "blob_size": 0} + try: + log("Sending encode request") + send_packet(self.sock, meta, b"") + resp_meta, blob = recv_packet(self.sock) + tensors = unpack_tensors(resp_meta["tensors"], blob) + cond = tensors.get("cond", torch.zeros((1, 77, 768))) + pooled = tensors.get("pooled", torch.zeros((1, 768))) + log("Received embeddings") + return cond, pooled + except Exception as e: + log(f"Remote CLIP inference failed: {e}") + return torch.zeros((1, 77, 768)), torch.zeros((1, 768)) + + def encode_from_tokens(self, tokens, return_pooled=True, return_dict=False): + cond, pooled = self._infer(tokens) + if return_dict: + return {"cond": cond, "pooled_output": pooled} + return cond, pooled + + def encode_from_tokens_scheduled(self, tokens): + cond, pooled = self._infer(tokens) + return [[cond, {"pooled_output": pooled}]] + +class RemoteCLIPWorker: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"clip": ("CLIP",), "listen_port": ("INT", {"default": 8002})}} + + RETURN_TYPES = () + FUNCTION = "start_worker" + OUTPUT_NODE = True + CATEGORY = "Remote CLIP" + + def start_worker(self, clip, listen_port): + server = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server.bind(("0.0.0.0", listen_port)) + server.listen(16) + log(f"Worker listening on port {listen_port}") + + def handle(conn): + addr = conn.getpeername() + log(f"Client connected: {addr}") + with conn: + while True: + try: + meta, _ = recv_packet(conn) + if meta is None or "text" not in meta: + log("Warning: invalid request, skipping") + continue + log(f"Encoding prompt ({len(meta['text'])} chars)") + tokens = clip.tokenize(meta["text"], **meta["kwargs"]) + cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) + tensor_meta, blob = pack_tensors({"cond": cond, "pooled": pooled}) + reply_meta = {"tensors": tensor_meta, "blob_size": len(blob)} + send_packet(conn, reply_meta, blob) + log(f"Sent embeddings ({len(blob)} bytes)") + except Exception as e: + log(f"Client disconnected: {e}") + break + + while True: + conn, _ = server.accept() + threading.Thread(target=handle, args=(conn,), daemon=True).start() + +class RemoteCLIPLoader: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"worker_ip": ("STRING", {"default": "10.0.0.37"}), "port": ("INT", {"default": 8002})}} + + RETURN_TYPES = ("CLIP",) + FUNCTION = "load_remote" + CATEGORY = "Remote CLIP" + + def load_remote(self, worker_ip, port): + return (RemoteCLIPProxy(worker_ip, port),) + +NODE_CLASS_MAPPINGS = {"RemoteCLIPWorker": RemoteCLIPWorker, "RemoteCLIPLoader": RemoteCLIPLoader} +NODE_DISPLAY_NAME_MAPPINGS = {"RemoteCLIPWorker": "Remote CLIP Worker", "RemoteCLIPLoader": "Remote CLIP Loader"} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..f9b2aa9 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +torch +numpy \ No newline at end of file