Update protocol version and improve connection handling
Updated protocol version and added error handling for packet size limits. Enhanced connection handling with new timeout and retry parameters.
This commit is contained in:
nyueki
+329
-326
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import math
|
||||
import json
|
||||
import time
|
||||
import hmac
|
||||
@@ -11,7 +12,7 @@ import folder_paths
|
||||
import comfy.sd
|
||||
import comfy.utils
|
||||
DEFAULT_PORT = 8181
|
||||
PROTOCOL_VERSION = 2
|
||||
PROTOCOL_VERSION = 3
|
||||
HEADER_LEN_FMT = ">Q"
|
||||
HEADER_LEN_SIZE = struct.calcsize(HEADER_LEN_FMT)
|
||||
SOCKET_TIMEOUT = 120
|
||||
@@ -22,7 +23,6 @@ RECV_CHUNK = 1 << 20
|
||||
MAX_HEADER_BYTES = 8 * 1024 * 1024
|
||||
MAX_BLOB_BYTES = 512 * 1024 * 1024
|
||||
PATCHED_CLIP_CACHE = 4
|
||||
EMBED_CACHE = 64
|
||||
ALLOWED_DTYPES = {
|
||||
"torch.float32": torch.float32,
|
||||
"torch.float16": torch.float16,
|
||||
@@ -39,7 +39,7 @@ def resolve_transport_dtype(mode, ip):
|
||||
if mode == "fp16":
|
||||
return "torch.float16"
|
||||
if mode == "fp32":
|
||||
return None
|
||||
return "torch.float32"
|
||||
is_local = ip in _LOCAL_HOSTS
|
||||
return None if is_local else "torch.float16"
|
||||
def _recv_exact(sock, size):
|
||||
@@ -54,6 +54,8 @@ def _recv_exact(sock, size):
|
||||
return b"".join(chunks)
|
||||
def send_packet(sock, header, blob=b""):
|
||||
header_bytes = json.dumps(header).encode("utf-8")
|
||||
if len(header_bytes) > MAX_HEADER_BYTES or len(blob) > MAX_BLOB_BYTES:
|
||||
raise ValueError("Packet exceeds protocol size limits")
|
||||
prefix = struct.pack(HEADER_LEN_FMT, len(header_bytes))
|
||||
sock.sendall(prefix + header_bytes)
|
||||
if blob:
|
||||
@@ -62,10 +64,13 @@ def recv_header(sock):
|
||||
header_len = struct.unpack(HEADER_LEN_FMT, _recv_exact(sock, HEADER_LEN_SIZE))[0]
|
||||
if header_len > MAX_HEADER_BYTES:
|
||||
raise ValueError(f"Header too large: {header_len} bytes")
|
||||
return json.loads(_recv_exact(sock, header_len).decode("utf-8"))
|
||||
header = json.loads(_recv_exact(sock, header_len).decode("utf-8"))
|
||||
if not isinstance(header, dict):
|
||||
raise ValueError("Packet header must be an object")
|
||||
return header
|
||||
def recv_blob(sock, header):
|
||||
blob_size = header.get("blob_size", 0)
|
||||
if blob_size > MAX_BLOB_BYTES:
|
||||
if type(blob_size) is not int or not 0 <= blob_size <= MAX_BLOB_BYTES:
|
||||
raise ValueError(f"Blob too large: {blob_size} bytes")
|
||||
return _recv_exact(sock, blob_size) if blob_size else b""
|
||||
def recv_packet(sock):
|
||||
@@ -109,57 +114,95 @@ def unpack_tensors(meta, blob):
|
||||
if dtype_name not in ALLOWED_DTYPES:
|
||||
raise ValueError(f"Refusing to decode disallowed dtype: {dtype_name}")
|
||||
dtype = ALLOWED_DTYPES[dtype_name]
|
||||
start = info["offset"]
|
||||
raw = bytearray(view[start:start + info["size"]])
|
||||
t = torch.frombuffer(raw, dtype=dtype).reshape(info["shape"]).clone()
|
||||
start, size, shape = info["offset"], info["size"], info["shape"]
|
||||
if (type(start) is not int or type(size) is not int or
|
||||
start < 0 or size < 0 or start + size > len(blob)):
|
||||
raise ValueError("Invalid tensor byte range")
|
||||
if not isinstance(shape, list) or any(type(d) is not int or d < 0 for d in shape):
|
||||
raise ValueError("Invalid tensor shape")
|
||||
count = math.prod(shape)
|
||||
if count * torch.empty((), dtype=dtype).element_size() != size:
|
||||
raise ValueError("Tensor shape does not match byte size")
|
||||
if count == 0:
|
||||
t = torch.empty(shape, dtype=dtype)
|
||||
else:
|
||||
raw = bytearray(view[start:start + size])
|
||||
t = torch.frombuffer(raw, dtype=dtype).reshape(shape).clone()
|
||||
orig = info.get("orig_dtype")
|
||||
if orig in ALLOWED_DTYPES and ALLOWED_DTYPES[orig] != dtype:
|
||||
t = t.to(ALLOWED_DTYPES[orig])
|
||||
out[name] = t
|
||||
return out
|
||||
def extract_tensors(obj, prefix, out):
|
||||
# Tag every container so user dictionaries cannot mimic tensor references.
|
||||
if isinstance(obj, torch.Tensor):
|
||||
out[prefix] = obj
|
||||
return {"__tensor__": prefix}
|
||||
key = f"tensor_{len(out)}"
|
||||
out[key] = obj
|
||||
return {"kind": "tensor", "value": key}
|
||||
if isinstance(obj, dict):
|
||||
return {k: extract_tensors(v, f"{prefix}.{k}", out) for k, v in obj.items()}
|
||||
return {"kind": "dict", "value": [
|
||||
[extract_tensors(k, prefix, out), extract_tensors(v, prefix, out)]
|
||||
for k, v in obj.items()]}
|
||||
if isinstance(obj, (list, tuple)):
|
||||
return [extract_tensors(v, f"{prefix}.{i}", out) for i, v in enumerate(obj)]
|
||||
return obj
|
||||
return {"kind": "tuple" if isinstance(obj, tuple) else "list",
|
||||
"value": [extract_tensors(v, prefix, out) for v in obj]}
|
||||
return {"kind": "value", "value": obj}
|
||||
|
||||
|
||||
def restore_tensors(obj, tensors):
|
||||
if isinstance(obj, dict):
|
||||
if len(obj) == 1 and "__tensor__" in obj:
|
||||
return tensors[obj["__tensor__"]]
|
||||
return {k: restore_tensors(v, tensors) for k, v in obj.items()}
|
||||
if isinstance(obj, list):
|
||||
return [restore_tensors(v, tensors) for v in obj]
|
||||
return obj
|
||||
kind, value = obj["kind"], obj["value"]
|
||||
if kind == "tensor":
|
||||
return tensors[value]
|
||||
if kind == "dict":
|
||||
return {restore_tensors(k, tensors): restore_tensors(v, tensors) for k, v in value}
|
||||
if kind in ("list", "tuple"):
|
||||
values = [restore_tensors(v, tensors) for v in value]
|
||||
return tuple(values) if kind == "tuple" else values
|
||||
if kind != "value":
|
||||
raise ValueError(f"Invalid structure kind: {kind}")
|
||||
return value
|
||||
|
||||
|
||||
def pack_object(obj, transport_dtype=None):
|
||||
tensors = {}
|
||||
structure = extract_tensors(obj, "", tensors)
|
||||
meta, blob = pack_tensors(tensors, transport_dtype)
|
||||
return {"struct": structure, "tensors": meta, "blob_size": len(blob)}, blob
|
||||
|
||||
|
||||
def unpack_object(header, blob):
|
||||
return restore_tensors(header["struct"], unpack_tensors(header["tensors"], blob))
|
||||
|
||||
|
||||
class _Connection:
|
||||
def __init__(self, ip, port, auth_token=""):
|
||||
self.ip = ip
|
||||
self.port = port
|
||||
self.auth_token = auth_token
|
||||
def __init__(self, ip, port, auth_token="", connect_timeout=SOCKET_TIMEOUT,
|
||||
request_attempts=CONNECT_RETRIES, connect_backoff=CONNECT_BACKOFF,
|
||||
encode_timeout=SOCKET_TIMEOUT):
|
||||
self.ip, self.port, self.auth_token = ip, port, auth_token
|
||||
self.connect_timeout = float(connect_timeout)
|
||||
self.request_attempts = int(request_attempts)
|
||||
self.connect_backoff = float(connect_backoff)
|
||||
self.encode_timeout = float(encode_timeout)
|
||||
for value in (self.connect_timeout, self.encode_timeout):
|
||||
if not math.isfinite(value) or value <= 0:
|
||||
raise ValueError("Timeouts must be finite and positive")
|
||||
if self.request_attempts < 1:
|
||||
raise ValueError("Request attempts must be at least 1")
|
||||
if not math.isfinite(self.connect_backoff) or self.connect_backoff < 0:
|
||||
raise ValueError("Retry delay must be finite and nonnegative")
|
||||
self._sock = None
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def _connect(self):
|
||||
last_err = None
|
||||
backoff = CONNECT_BACKOFF
|
||||
for attempt in range(1, CONNECT_RETRIES + 1):
|
||||
try:
|
||||
sock = socket.create_connection((self.ip, self.port), timeout=SOCKET_TIMEOUT)
|
||||
_set_socket_opts(sock)
|
||||
self._sock = sock
|
||||
log(f"Connected to worker {self.ip}:{self.port}")
|
||||
return
|
||||
except OSError as e:
|
||||
last_err = e
|
||||
log(f"Connect attempt {attempt}/{CONNECT_RETRIES} failed: {e}")
|
||||
if attempt < CONNECT_RETRIES:
|
||||
time.sleep(backoff)
|
||||
backoff *= 2
|
||||
raise ConnectionError(
|
||||
f"Could not connect to worker {self.ip}:{self.port}: {last_err}"
|
||||
)
|
||||
sock = socket.create_connection((self.ip, self.port), timeout=self.connect_timeout)
|
||||
try:
|
||||
_set_socket_opts(sock)
|
||||
except Exception:
|
||||
sock.close()
|
||||
raise
|
||||
self._sock = sock
|
||||
log(f"Connected to worker {self.ip}:{self.port}")
|
||||
|
||||
def _close(self):
|
||||
if self._sock is not None:
|
||||
try:
|
||||
@@ -167,351 +210,311 @@ class _Connection:
|
||||
except OSError:
|
||||
pass
|
||||
self._sock = None
|
||||
def request(self, header, blob=b"", response_timeout=None, retries=2):
|
||||
if self.auth_token:
|
||||
header = {**header, "auth": self.auth_token}
|
||||
|
||||
def request(self, header, blob=b"", response_timeout=None, replay_safe=True):
|
||||
header = {**header, "auth": self.auth_token}
|
||||
timeout = self.encode_timeout if response_timeout is None else response_timeout
|
||||
if not math.isfinite(timeout) or timeout <= 0:
|
||||
raise ValueError("Response timeout must be finite and positive")
|
||||
with self._lock:
|
||||
for attempt in range(1, retries + 1):
|
||||
backoff = self.connect_backoff
|
||||
for attempt in range(1, self.request_attempts + 1):
|
||||
sending = False
|
||||
try:
|
||||
if self._sock is None:
|
||||
self._connect()
|
||||
if response_timeout is not None:
|
||||
self._sock.settimeout(response_timeout)
|
||||
self._sock.settimeout(timeout)
|
||||
sending = True
|
||||
send_packet(self._sock, header, blob)
|
||||
resp = recv_packet(self._sock)
|
||||
if response_timeout is not None:
|
||||
self._sock.settimeout(SOCKET_TIMEOUT)
|
||||
return resp
|
||||
except (ConnectionError, OSError, struct.error, ValueError) as e:
|
||||
log(f"Request failed ({e}); reconnecting (attempt {attempt}/{retries})")
|
||||
result = recv_packet(self._sock)
|
||||
self._sock.settimeout(self.encode_timeout)
|
||||
return result
|
||||
except (OSError, ConnectionError, struct.error, ValueError):
|
||||
self._close()
|
||||
if attempt == retries:
|
||||
# Generation may have executed even when its response was lost.
|
||||
if attempt == self.request_attempts or (sending and not replay_safe):
|
||||
raise
|
||||
log(f"Request attempt {attempt}/{self.request_attempts} failed; retrying")
|
||||
time.sleep(backoff)
|
||||
backoff *= 2
|
||||
|
||||
|
||||
class RemoteCLIPProxy:
|
||||
def __init__(self, ip, port, auth_token="", transport_mode="auto",
|
||||
connection=None, lora_stack=None):
|
||||
self.ip = ip
|
||||
self.port = port
|
||||
connection=None, lora_stack=None, connect_timeout=SOCKET_TIMEOUT,
|
||||
request_attempts=CONNECT_RETRIES, connect_backoff=CONNECT_BACKOFF,
|
||||
encode_timeout=SOCKET_TIMEOUT, generate_timeout=GENERATE_TIMEOUT,
|
||||
layer_idx=None):
|
||||
self.ip, self.port = ip, port
|
||||
self.transport_mode = transport_mode
|
||||
self.transport_dtype_name = resolve_transport_dtype(transport_mode, ip)
|
||||
self.transport_dtype = (
|
||||
ALLOWED_DTYPES[self.transport_dtype_name]
|
||||
if self.transport_dtype_name is not None
|
||||
else None
|
||||
)
|
||||
self._conn = connection or _Connection(ip, port, auth_token)
|
||||
self.transport_dtype = ALLOWED_DTYPES.get(self.transport_dtype_name)
|
||||
self.generate_timeout = float(generate_timeout)
|
||||
if not math.isfinite(self.generate_timeout) or self.generate_timeout <= 0:
|
||||
raise ValueError("Generation timeout must be finite and positive")
|
||||
self._conn = connection or _Connection(
|
||||
ip, port, auth_token, connect_timeout, request_attempts,
|
||||
connect_backoff, encode_timeout)
|
||||
self.lora_stack = list(lora_stack or [])
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
def clone(self, disable_dynamic=False):
|
||||
return RemoteCLIPProxy(
|
||||
self.ip, self.port,
|
||||
transport_mode=self.transport_mode,
|
||||
connection=self._conn,
|
||||
lora_stack=self.lora_stack,
|
||||
)
|
||||
self.ip, self.port, transport_mode=self.transport_mode,
|
||||
connection=self._conn, lora_stack=self.lora_stack,
|
||||
generate_timeout=self.generate_timeout, layer_idx=self.layer_idx)
|
||||
|
||||
def clip_layer(self, layer_idx):
|
||||
self.layer_idx = layer_idx
|
||||
|
||||
def with_lora(self, lora_name, strength_clip):
|
||||
c = self.clone()
|
||||
c.lora_stack = self.lora_stack + [[lora_name, float(strength_clip)]]
|
||||
return c
|
||||
|
||||
def add_patches(self, *_args, **_kwargs):
|
||||
raise NotImplementedError(
|
||||
"Apply LoRAs to a remote CLIP with the 'LoraLoaderCLIPOnly' node, "
|
||||
"which forwards them to the worker."
|
||||
)
|
||||
"Apply LoRAs to a remote CLIP with the 'LoraLoaderCLIPOnly' node.")
|
||||
|
||||
def _rpc(self, cmd, payload, generation=False):
|
||||
# Token tensors retain their original numerical values. Transport precision
|
||||
# applies to image/tokenizer inputs and embedding responses, not token IDs.
|
||||
header, blob = pack_object(payload)
|
||||
header.update(cmd=cmd, proto=PROTOCOL_VERSION,
|
||||
lora_stack=self.lora_stack, layer_idx=self.layer_idx,
|
||||
transport_dtype=self.transport_dtype_name)
|
||||
response, data = self._conn.request(
|
||||
header, blob,
|
||||
response_timeout=self.generate_timeout if generation else None,
|
||||
replay_safe=not generation)
|
||||
if response.get("error"):
|
||||
raise RuntimeError(f"Remote CLIP worker error: {response['error']}")
|
||||
return unpack_object(response, data)
|
||||
|
||||
def tokenize(self, text, return_word_ids=False, **kwargs):
|
||||
if return_word_ids:
|
||||
kwargs["return_word_ids"] = True
|
||||
return {"text": text, "kwargs": kwargs, "lora_stack": self.lora_stack}
|
||||
def _encode(self, tokens):
|
||||
kwargs = dict(tokens.get("kwargs", {}))
|
||||
tensor_inputs = {}
|
||||
kwargs = extract_tensors(kwargs, "kwargs", tensor_inputs)
|
||||
if tensor_inputs:
|
||||
in_meta, in_blob = pack_tensors(tensor_inputs, self.transport_dtype)
|
||||
else:
|
||||
in_meta, in_blob = {}, b""
|
||||
header = {
|
||||
"cmd": "encode",
|
||||
"proto": PROTOCOL_VERSION,
|
||||
"text": tokens["text"],
|
||||
"kwargs": kwargs,
|
||||
"lora_stack": tokens.get("lora_stack", self.lora_stack),
|
||||
"tensor_inputs": in_meta,
|
||||
"transport_dtype": self.transport_dtype_name,
|
||||
"blob_size": len(in_blob),
|
||||
}
|
||||
log("Sending encode request")
|
||||
resp, blob = self._conn.request(header, in_blob)
|
||||
if resp.get("error"):
|
||||
raise RuntimeError(f"Remote CLIP worker error: {resp['error']}")
|
||||
tensors = unpack_tensors(resp["tensors"], blob)
|
||||
out = restore_tensors(resp["struct"], tensors)
|
||||
log("Received embeddings")
|
||||
return out
|
||||
# Return actual token dictionaries: callers such as SDXL edit l/g streams.
|
||||
return self._rpc("tokenize", {"text": text, "kwargs": {
|
||||
**kwargs, "return_word_ids": return_word_ids}})
|
||||
|
||||
def encode_from_tokens(self, tokens, return_pooled=False, return_dict=False):
|
||||
out = self._encode(tokens)
|
||||
out = self._rpc("encode", {"tokens": tokens, "return_pooled": return_pooled})
|
||||
if return_dict:
|
||||
return out
|
||||
cond = out["cond"]
|
||||
if return_pooled:
|
||||
return cond, out.get("pooled_output")
|
||||
return cond
|
||||
def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict: dict = None, show_pbar=True):
|
||||
return_pooled = "unprojected" if unprojected else True
|
||||
pooled_dict = self.encode_from_tokens(tokens, return_pooled=return_pooled, return_dict=True)
|
||||
cond = pooled_dict.pop("cond")
|
||||
if add_dict:
|
||||
pooled_dict.update(add_dict)
|
||||
return [[cond, pooled_dict]]
|
||||
return out["cond"], out.get("pooled_output")
|
||||
return out["cond"]
|
||||
|
||||
def encode_from_tokens_scheduled(self, tokens, unprojected=False, add_dict=None, show_pbar=True):
|
||||
return self._rpc("encode_scheduled", {
|
||||
"tokens": tokens, "unprojected": unprojected,
|
||||
"add_dict": {} if add_dict is None else add_dict, "show_pbar": show_pbar})
|
||||
|
||||
def generate(self, tokens, **gen_kwargs):
|
||||
text = tokens["text"]
|
||||
kwargs = dict(tokens.get("kwargs", {}))
|
||||
lora_stack = tokens.get("lora_stack", self.lora_stack)
|
||||
tensor_inputs = {}
|
||||
kwargs = extract_tensors(kwargs, "kwargs", tensor_inputs)
|
||||
if tensor_inputs:
|
||||
meta, blob = pack_tensors(tensor_inputs, self.transport_dtype)
|
||||
else:
|
||||
meta, blob = {}, b""
|
||||
header = {
|
||||
"cmd": "generate",
|
||||
"proto": PROTOCOL_VERSION,
|
||||
"text": text,
|
||||
"kwargs": kwargs,
|
||||
"gen_kwargs": gen_kwargs,
|
||||
"lora_stack": lora_stack,
|
||||
"tensor_inputs": meta,
|
||||
"transport_dtype": self.transport_dtype_name,
|
||||
"blob_size": len(blob),
|
||||
}
|
||||
log("Sending generate request")
|
||||
resp, _ = self._conn.request(
|
||||
header, blob, response_timeout=GENERATE_TIMEOUT, retries=1
|
||||
)
|
||||
if resp.get("error"):
|
||||
raise RuntimeError(f"Remote CLIP worker error: {resp['error']}")
|
||||
log("Received generated text")
|
||||
return resp["text"]
|
||||
def decode(self, generated):
|
||||
return self._rpc("generate", {"tokens": tokens, "gen_kwargs": gen_kwargs}, generation=True)
|
||||
|
||||
def decode(self, generated, skip_special_tokens=True):
|
||||
return generated
|
||||
|
||||
|
||||
class _Worker:
|
||||
def __init__(self, clip, auth_token=""):
|
||||
self.base_clip = clip
|
||||
self.auth_token = auth_token
|
||||
self._cache_lock = threading.Lock()
|
||||
self._infer_lock = threading.Lock()
|
||||
self._patched_clips = OrderedDict()
|
||||
self._embed_cache = OrderedDict()
|
||||
@staticmethod
|
||||
def _sig(obj):
|
||||
return json.dumps(obj, sort_keys=True)
|
||||
def _get_clip(self, lora_stack):
|
||||
if not lora_stack:
|
||||
return self.base_clip
|
||||
sig = self._sig(lora_stack)
|
||||
clip = self._patched_clips.get(sig)
|
||||
if clip is not None:
|
||||
self._infer_lock = threading.Lock()
|
||||
self._patched_clips = OrderedDict()
|
||||
|
||||
def _get_clip(self, lora_stack, layer_idx):
|
||||
sig = json.dumps([lora_stack, layer_idx], sort_keys=True)
|
||||
if sig in self._patched_clips:
|
||||
self._patched_clips.move_to_end(sig)
|
||||
return clip
|
||||
clip = self.base_clip
|
||||
return self._patched_clips[sig]
|
||||
clip = self.base_clip.clone()
|
||||
if layer_idx is not None:
|
||||
clip.clip_layer(layer_idx)
|
||||
for lora_name, strength in lora_stack:
|
||||
if strength == 0:
|
||||
continue
|
||||
lora_path = folder_paths.get_full_path("loras", lora_name)
|
||||
if lora_path is None:
|
||||
path = folder_paths.get_full_path("loras", lora_name)
|
||||
if path is None:
|
||||
raise FileNotFoundError(f"LoRA not found on worker: {lora_name}")
|
||||
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
||||
lora = comfy.utils.load_torch_file(path, safe_load=True)
|
||||
_, clip = comfy.sd.load_lora_for_models(None, clip, lora, 0, strength)
|
||||
self._patched_clips[sig] = clip
|
||||
while len(self._patched_clips) > PATCHED_CLIP_CACHE:
|
||||
self._patched_clips.popitem(last=False)
|
||||
return clip
|
||||
def encode(self, text, kwargs, lora_stack):
|
||||
|
||||
def execute(self, cmd, payload, lora_stack, layer_idx):
|
||||
with self._infer_lock, torch.inference_mode():
|
||||
clip = self._get_clip(lora_stack, layer_idx)
|
||||
if cmd == "tokenize":
|
||||
return clip.tokenize(payload["text"], **payload["kwargs"])
|
||||
if cmd == "generate":
|
||||
return clip.decode(clip.generate(payload["tokens"], **payload["gen_kwargs"]))
|
||||
if cmd == "encode":
|
||||
out = clip.encode_from_tokens(payload["tokens"],
|
||||
return_pooled=payload["return_pooled"], return_dict=True)
|
||||
elif cmd == "encode_scheduled":
|
||||
out = clip.encode_from_tokens_scheduled(**payload)
|
||||
else:
|
||||
raise ValueError(f"Unknown command: {cmd}")
|
||||
return out
|
||||
|
||||
|
||||
class _Server:
|
||||
def __init__(self, clip, host, port, token):
|
||||
self.worker = _Worker(clip, token)
|
||||
self.token = token
|
||||
self.stopped = threading.Event()
|
||||
self._clients = set()
|
||||
self._clients_lock = threading.Lock()
|
||||
self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
try:
|
||||
cache_key = self._sig([text, kwargs, lora_stack])
|
||||
except TypeError:
|
||||
cache_key = None
|
||||
if cache_key is not None:
|
||||
with self._cache_lock:
|
||||
cached = self._embed_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
self._embed_cache.move_to_end(cache_key)
|
||||
return cached
|
||||
with self._infer_lock:
|
||||
if cache_key is not None:
|
||||
with self._cache_lock:
|
||||
cached = self._embed_cache.get(cache_key)
|
||||
if cached is not None:
|
||||
self._embed_cache.move_to_end(cache_key)
|
||||
return cached
|
||||
clip = self._get_clip(lora_stack)
|
||||
with torch.inference_mode():
|
||||
tokens = clip.tokenize(text, **kwargs)
|
||||
out = clip.encode_from_tokens(tokens, return_dict=True)
|
||||
cond = out.get("cond")
|
||||
if cond is None:
|
||||
raise RuntimeError("encode produced no cond tensor")
|
||||
result = {}
|
||||
for k, v in out.items():
|
||||
result[k] = v.detach().cpu().clone() if isinstance(v, torch.Tensor) else v
|
||||
if result.get("pooled_output") is None:
|
||||
result["pooled_output"] = torch.zeros((cond.shape[0], cond.shape[-1]))
|
||||
if cache_key is not None:
|
||||
with self._cache_lock:
|
||||
self._embed_cache[cache_key] = result
|
||||
while len(self._embed_cache) > EMBED_CACHE:
|
||||
self._embed_cache.popitem(last=False)
|
||||
return result
|
||||
def generate(self, text, kwargs, gen_kwargs, lora_stack):
|
||||
with self._infer_lock:
|
||||
clip = self._get_clip(lora_stack)
|
||||
with torch.inference_mode():
|
||||
tokens = clip.tokenize(text, **kwargs)
|
||||
generated_ids = clip.generate(tokens, **gen_kwargs)
|
||||
generated_text = clip.decode(generated_ids)
|
||||
return generated_text
|
||||
self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
self.socket.bind((host, port))
|
||||
self.socket.listen(16)
|
||||
self.socket.settimeout(0.5)
|
||||
except Exception:
|
||||
self.socket.close()
|
||||
raise
|
||||
self.thread = threading.Thread(target=self._accept_loop, daemon=True)
|
||||
self.thread.start()
|
||||
|
||||
@staticmethod
|
||||
def _close_socket(sock):
|
||||
try:
|
||||
sock.shutdown(socket.SHUT_RDWR)
|
||||
except OSError:
|
||||
pass
|
||||
sock.close()
|
||||
|
||||
def close(self):
|
||||
self.stopped.set()
|
||||
self._close_socket(self.socket)
|
||||
with self._clients_lock:
|
||||
clients = list(self._clients)
|
||||
for conn in clients:
|
||||
self._close_socket(conn)
|
||||
self.thread.join(timeout=1)
|
||||
|
||||
def _accept_loop(self):
|
||||
while not self.stopped.is_set():
|
||||
try:
|
||||
conn, addr = self.socket.accept()
|
||||
except socket.timeout:
|
||||
continue
|
||||
except OSError:
|
||||
break
|
||||
with self._clients_lock:
|
||||
if self.stopped.is_set():
|
||||
conn.close()
|
||||
break
|
||||
self._clients.add(conn)
|
||||
threading.Thread(target=self._handle, args=(conn, addr), daemon=True).start()
|
||||
|
||||
def _handle(self, conn, addr):
|
||||
try:
|
||||
_set_socket_opts(conn)
|
||||
while not self.stopped.is_set():
|
||||
header = recv_header(conn)
|
||||
auth = header.get("auth", "")
|
||||
if self.token and (not isinstance(auth, str) or not hmac.compare_digest(
|
||||
auth.encode("utf-8"), self.token.encode("utf-8"))):
|
||||
send_packet(conn, {"error": "unauthorized", "blob_size": 0})
|
||||
break
|
||||
if header.get("proto") != PROTOCOL_VERSION:
|
||||
send_packet(conn, {"error": f"Protocol mismatch: install v{PROTOCOL_VERSION} on both ends", "blob_size": 0})
|
||||
break
|
||||
blob = recv_blob(conn, header)
|
||||
try:
|
||||
payload = unpack_object(header, blob)
|
||||
if self.stopped.is_set():
|
||||
break
|
||||
out = self.worker.execute(header["cmd"], payload,
|
||||
header.get("lora_stack", []), header.get("layer_idx"))
|
||||
if self.stopped.is_set():
|
||||
break
|
||||
transport = header.get("transport_dtype")
|
||||
if transport not in (None, "torch.float16", "torch.float32"):
|
||||
raise ValueError("Invalid transport precision")
|
||||
dtype = ALLOWED_DTYPES.get(transport) if header["cmd"] in ("encode", "encode_scheduled") else None
|
||||
response, data = pack_object(out, dtype)
|
||||
send_packet(conn, response, data)
|
||||
except Exception as e:
|
||||
if self.stopped.is_set():
|
||||
break
|
||||
log(f"Worker request failed: {e}")
|
||||
send_packet(conn, {"error": str(e), "blob_size": 0})
|
||||
except (ConnectionError, OSError, ValueError, TypeError, KeyError, struct.error) as e:
|
||||
log(f"Client {addr} disconnected: {e}")
|
||||
finally:
|
||||
with self._clients_lock:
|
||||
self._clients.discard(conn)
|
||||
self._close_socket(conn)
|
||||
|
||||
|
||||
class SendRemoteCLIP:
|
||||
_servers = {}
|
||||
_servers = {}
|
||||
_servers_lock = threading.Lock()
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"CLIP": ("CLIP",),
|
||||
"listen_port": ("INT", {"default": DEFAULT_PORT, "min": 1, "max": 65535}),
|
||||
},
|
||||
"optional": {
|
||||
"bind_host": ("STRING", {"default": "0.0.0.0"}),
|
||||
"auth_token": ("STRING", {"default": ""}),
|
||||
},
|
||||
}
|
||||
return {"required": {
|
||||
"CLIP": ("CLIP",),
|
||||
"listen_port": ("INT", {"default": DEFAULT_PORT, "min": 1, "max": 65535}),
|
||||
}, "optional": {
|
||||
"bind_host": ("STRING", {"default": "0.0.0.0"}),
|
||||
"auth_token": ("STRING", {"default": ""}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "start_worker"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "Remote CLIP"
|
||||
|
||||
def start_worker(self, CLIP, listen_port, bind_host="0.0.0.0", auth_token=""):
|
||||
token = auth_token or os.environ.get("REMOTE_CLIP_TOKEN", "")
|
||||
existing = SendRemoteCLIP._servers.pop(listen_port, None)
|
||||
if existing is not None:
|
||||
try:
|
||||
existing.close()
|
||||
except OSError:
|
||||
pass
|
||||
worker = _Worker(CLIP, token)
|
||||
server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
server.bind((bind_host, listen_port))
|
||||
server.listen(16)
|
||||
SendRemoteCLIP._servers[listen_port] = server
|
||||
with self._servers_lock:
|
||||
previous = self._servers.pop(listen_port, None)
|
||||
if previous is not None:
|
||||
previous.close()
|
||||
self._servers[listen_port] = _Server(CLIP, bind_host, listen_port, token)
|
||||
if bind_host == "0.0.0.0" and not token:
|
||||
log("WARNING: worker bound to 0.0.0.0 with no auth token. Anyone on "
|
||||
"the network can use this CLIP. Set auth_token or bind to a "
|
||||
"specific host.")
|
||||
log("WARNING: worker is accessible on all interfaces without authentication")
|
||||
log(f"Worker listening on {bind_host}:{listen_port}")
|
||||
def handle(conn, addr):
|
||||
log(f"Client connected: {addr}")
|
||||
try:
|
||||
_set_socket_opts(conn)
|
||||
while True:
|
||||
header = recv_header(conn)
|
||||
if token and not hmac.compare_digest(header.get("auth", ""), token):
|
||||
send_packet(conn, {"error": "unauthorized", "blob_size": 0})
|
||||
log(f"Rejected unauthorized client {addr}")
|
||||
break
|
||||
client_proto = header.get("proto", 0)
|
||||
if client_proto != PROTOCOL_VERSION:
|
||||
send_packet(conn, {
|
||||
"error": (f"protocol mismatch: worker speaks v{PROTOCOL_VERSION}, "
|
||||
f"client sent v{client_proto}"),
|
||||
"blob_size": 0,
|
||||
})
|
||||
log(f"Rejected client {addr}: protocol v{client_proto} "
|
||||
f"!= v{PROTOCOL_VERSION}")
|
||||
break
|
||||
blob_data = recv_blob(conn, header)
|
||||
cmd = header.get("cmd")
|
||||
if cmd == "encode" and "text" in header:
|
||||
try:
|
||||
transport = header.get("transport_dtype")
|
||||
transport_dtype = ALLOWED_DTYPES.get(transport) if transport else None
|
||||
input_tensors = unpack_tensors(
|
||||
header.get("tensor_inputs", {}), blob_data
|
||||
)
|
||||
enc_kwargs = restore_tensors(
|
||||
header.get("kwargs", {}), input_tensors
|
||||
)
|
||||
out = worker.encode(
|
||||
header["text"],
|
||||
enc_kwargs,
|
||||
header.get("lora_stack", []),
|
||||
)
|
||||
tensors = {}
|
||||
out_struct = extract_tensors(out, "out", tensors)
|
||||
meta, blob = pack_tensors(tensors, transport_dtype)
|
||||
send_packet(
|
||||
conn,
|
||||
{"struct": out_struct, "tensors": meta, "blob_size": len(blob)},
|
||||
blob,
|
||||
)
|
||||
log(f"Sent embeddings ({len(blob)} bytes)")
|
||||
except Exception as e:
|
||||
log(f"Encode failed: {e}")
|
||||
send_packet(conn, {"error": str(e), "blob_size": 0})
|
||||
elif cmd == "generate" and "text" in header:
|
||||
try:
|
||||
input_tensors = unpack_tensors(
|
||||
header.get("tensor_inputs", {}), blob_data
|
||||
)
|
||||
kwargs = restore_tensors(
|
||||
header.get("kwargs", {}), input_tensors
|
||||
)
|
||||
generated_text = worker.generate(
|
||||
header["text"],
|
||||
kwargs,
|
||||
header.get("gen_kwargs", {}),
|
||||
header.get("lora_stack", []),
|
||||
)
|
||||
send_packet(conn, {"text": generated_text, "blob_size": 0})
|
||||
log("Sent generated text")
|
||||
except Exception as e:
|
||||
log(f"Generate failed: {e}")
|
||||
send_packet(conn, {"error": str(e), "blob_size": 0})
|
||||
else:
|
||||
send_packet(conn, {"error": "bad request", "blob_size": 0})
|
||||
except (ConnectionError, OSError) as e:
|
||||
log(f"Client {addr} disconnected: {e}")
|
||||
finally:
|
||||
try:
|
||||
conn.close()
|
||||
except OSError:
|
||||
pass
|
||||
def accept_loop():
|
||||
while True:
|
||||
try:
|
||||
conn, addr = server.accept()
|
||||
except OSError:
|
||||
log("Server socket closed; stopping accept loop")
|
||||
break
|
||||
threading.Thread(target=handle, args=(conn, addr), daemon=True).start()
|
||||
threading.Thread(target=accept_loop, daemon=True).start()
|
||||
return {"ui": {"text": [f"Remote CLIP worker on {bind_host}:{listen_port}"]}}
|
||||
|
||||
|
||||
class LoadRemoteCLIP:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"worker_ip": ("STRING", {"default": "127.0.0.1"}),
|
||||
"port": ("INT", {"default": DEFAULT_PORT, "min": 1, "max": 65535}),
|
||||
},
|
||||
"optional": {
|
||||
"auth_token": ("STRING", {"default": ""}),
|
||||
"transport_precision": (["auto", "fp16", "fp32"], {"default": "auto"}),
|
||||
},
|
||||
}
|
||||
return {"required": {
|
||||
"worker_ip": ("STRING", {"default": "127.0.0.1"}),
|
||||
"port": ("INT", {"default": DEFAULT_PORT, "min": 1, "max": 65535}),
|
||||
}, "optional": {
|
||||
"auth_token": ("STRING", {"default": ""}),
|
||||
"transport_precision": (["auto", "fp16", "fp32"], {"default": "auto"}),
|
||||
"connect_timeout": ("FLOAT", {"default": SOCKET_TIMEOUT, "min": 0.1, "max": 3600.0, "step": 0.1, "tooltip": "Seconds per connection attempt."}),
|
||||
"request_attempts": ("INT", {"default": CONNECT_RETRIES, "min": 1, "max": 10, "tooltip": "Total attempts per operation, including the first. Connection failures and safe request retries share this single limit. Generation is never replayed after sending starts."}),
|
||||
"connect_backoff": ("FLOAT", {"default": CONNECT_BACKOFF, "min": 0.0, "max": 60.0, "step": 0.1, "tooltip": "Initial retry delay in seconds; doubles after each failure."}),
|
||||
"encode_timeout": ("FLOAT", {"default": SOCKET_TIMEOUT, "min": 0.1, "max": 86400.0, "step": 0.1, "tooltip": "Socket timeout for tokenization and encoding; not a total operation deadline."}),
|
||||
"generate_timeout": ("FLOAT", {"default": GENERATE_TIMEOUT, "min": 0.1, "max": 86400.0, "step": 0.1, "tooltip": "Generation socket timeout; not a total operation deadline."}),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("CLIP",)
|
||||
FUNCTION = "load_remote"
|
||||
CATEGORY = "Remote CLIP"
|
||||
def load_remote(self, worker_ip, port, auth_token="", transport_precision="auto"):
|
||||
|
||||
def load_remote(self, worker_ip, port, auth_token="", transport_precision="auto",
|
||||
connect_timeout=SOCKET_TIMEOUT, request_attempts=CONNECT_RETRIES,
|
||||
connect_backoff=CONNECT_BACKOFF, encode_timeout=SOCKET_TIMEOUT,
|
||||
generate_timeout=GENERATE_TIMEOUT):
|
||||
token = auth_token or os.environ.get("REMOTE_CLIP_TOKEN", "")
|
||||
return (RemoteCLIPProxy(worker_ip, port, token, transport_precision),)
|
||||
return (RemoteCLIPProxy(worker_ip, port, token, transport_precision,
|
||||
connect_timeout=connect_timeout, request_attempts=request_attempts,
|
||||
connect_backoff=connect_backoff, encode_timeout=encode_timeout,
|
||||
generate_timeout=generate_timeout),)
|
||||
|
||||
class LoraLoaderCLIPOnly:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
|
||||
Reference in New Issue
Block a user