Jason HokuandClaude Opus 4.6
d5fd56541b
Move torch and safetensors imports inside huggingface_vae_decode() so CivitAI and Distribution functions still work if those packages are missing. Update stale "Uses urllib" comments to reference the gateway. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
282 lines
9.6 KiB
Python
282 lines
9.6 KiB
Python
# ==============================================================================
|
|
# NETWORK UTILS — All outbound network requests are centralized here.
|
|
#
|
|
# External services contacted:
|
|
# 1. CivitAI API (civitai.com/api/v1) — model/LoRA metadata lookup
|
|
# 2. HuggingFace Endpoints (*.huggingface.cloud) — Remote VAE decoding
|
|
# 3. Distribution System (user-configured LAN addresses) — worker coordination
|
|
#
|
|
# No other file in this project makes outbound network requests.
|
|
# All urllib usage is confined to this module.
|
|
# ==============================================================================
|
|
|
|
import json
|
|
import os
|
|
import urllib.request
|
|
import urllib.error
|
|
import urllib.parse
|
|
|
|
|
|
# =============================================================================
|
|
# 1. CIVITAI API — Model/LoRA metadata lookup
|
|
# Target: civitai.com/api/v1 (ONLY)
|
|
# =============================================================================
|
|
|
|
CIVITAI_API_BASE = "https://civitai.com/api/v1"
|
|
CIVITAI_TIMEOUT = 10 # seconds
|
|
CIVITAI_USER_AGENT = "ComfyUI-ConfigBuilder"
|
|
|
|
|
|
def civitai_fetch_by_hash(hash_value):
|
|
"""
|
|
Fetch model version info from CivitAI API using a file hash.
|
|
|
|
Network call: GET https://civitai.com/api/v1/model-versions/by-hash/{hash}
|
|
Timeout: 10 seconds
|
|
Returns: Parsed JSON dict on success, None on any error.
|
|
"""
|
|
api_url = f"{CIVITAI_API_BASE}/model-versions/by-hash/{hash_value}"
|
|
try:
|
|
req = urllib.request.Request(
|
|
api_url,
|
|
headers={"User-Agent": CIVITAI_USER_AGENT}
|
|
)
|
|
with urllib.request.urlopen(req, timeout=CIVITAI_TIMEOUT) as response:
|
|
if response.status == 200:
|
|
return json.loads(response.read().decode("utf-8"))
|
|
else:
|
|
return None
|
|
except urllib.error.HTTPError:
|
|
return None
|
|
except Exception as e:
|
|
print(f"[NetworkUtils] ⚠️ CivitAI lookup failed: {e}")
|
|
return None
|
|
|
|
|
|
# =============================================================================
|
|
# 2. HUGGINGFACE REMOTE VAE — Offload VAE decoding to HF endpoints
|
|
# Target: *.huggingface.cloud (allowlisted endpoints ONLY)
|
|
# =============================================================================
|
|
|
|
HUGGINGFACE_VAE_ENDPOINTS = {
|
|
"SD": "https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud/",
|
|
"SDXL": "https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud/",
|
|
"Flux": "https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud/",
|
|
"HunyuanVideo": "https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud/"
|
|
}
|
|
HUGGINGFACE_VAE_TIMEOUT = 60 # seconds
|
|
|
|
|
|
def huggingface_vae_decode(endpoint_url, tensor, height, width):
|
|
"""
|
|
Send latent tensor to a HuggingFace Remote VAE endpoint for decoding.
|
|
|
|
Network call: POST to an allowlisted *.huggingface.cloud endpoint
|
|
Timeout: 60 seconds
|
|
Validates endpoint_url against HUGGINGFACE_VAE_ENDPOINTS allowlist.
|
|
|
|
Args:
|
|
endpoint_url: Must be one of HUGGINGFACE_VAE_ENDPOINTS values
|
|
tensor: Raw torch tensor (will be serialized internally)
|
|
height: Image height
|
|
width: Image width
|
|
|
|
Returns:
|
|
Tuple of (response_bytes, response_headers_dict) for caller to parse
|
|
|
|
Raises:
|
|
ValueError: If endpoint_url is not in the allowlist
|
|
RuntimeError: If the remote endpoint returns non-200
|
|
"""
|
|
# Lazy imports — only needed for HF VAE, avoids breaking CivitAI/Distribution
|
|
# if torch or safetensors are not installed
|
|
import torch
|
|
from safetensors.torch import _tobytes
|
|
|
|
# Validate endpoint against allowlist
|
|
if endpoint_url not in HUGGINGFACE_VAE_ENDPOINTS.values():
|
|
raise ValueError(
|
|
f"HuggingFace VAE endpoint not in allowlist: {endpoint_url}\n"
|
|
f"Allowed: {list(HUGGINGFACE_VAE_ENDPOINTS.values())}"
|
|
)
|
|
|
|
# Prepare tensor for transport
|
|
tensor = tensor.cpu().contiguous()
|
|
|
|
# Build query parameters
|
|
shape_values = [int(dim) for dim in tensor.shape]
|
|
query_parts = [
|
|
("do_scaling", "False"),
|
|
("output_type", "pt"),
|
|
("partial_postprocess", "False"),
|
|
("dtype", str(tensor.dtype).split(".")[-1]),
|
|
("height", str(int(height))),
|
|
("width", str(int(width))),
|
|
]
|
|
for s in shape_values:
|
|
query_parts.append(("shape", str(s)))
|
|
|
|
query_string = urllib.parse.urlencode(query_parts)
|
|
full_url = f"{endpoint_url}?{query_string}"
|
|
|
|
# Serialize tensor to bytes
|
|
tensor_data = _tobytes(tensor, "tensor")
|
|
|
|
# Execute request
|
|
req = urllib.request.Request(
|
|
full_url,
|
|
data=tensor_data,
|
|
headers={
|
|
"Content-Type": "tensor/binary",
|
|
"Accept": "tensor/binary",
|
|
},
|
|
method="POST",
|
|
)
|
|
|
|
with urllib.request.urlopen(req, timeout=HUGGINGFACE_VAE_TIMEOUT) as response:
|
|
if response.status != 200:
|
|
error_text = response.read().decode("utf-8", errors="replace")
|
|
raise RuntimeError(f"Remote VAE decode failed: {error_text}")
|
|
|
|
output_data = response.read()
|
|
# Capture headers before response closes
|
|
headers = {k: response.headers.get(k) for k in ["shape", "dtype"]}
|
|
|
|
return output_data, headers
|
|
|
|
|
|
# =============================================================================
|
|
# 3. DISTRIBUTION SYSTEM — LAN-only worker coordination
|
|
# Target: User-configured LAN addresses (ComfyUI worker instances)
|
|
# =============================================================================
|
|
|
|
DISTRIBUTION_DEFAULT_TIMEOUT = 10 # seconds
|
|
|
|
|
|
def distribution_get(url, timeout=None):
|
|
"""
|
|
LAN-only: GET request to a user-configured ComfyUI worker/master instance.
|
|
|
|
Args:
|
|
url: Full URL (caller constructs from master_url + endpoint)
|
|
timeout: Request timeout in seconds (default: 10)
|
|
|
|
Returns:
|
|
Tuple of (status_code, response_bytes)
|
|
|
|
Raises:
|
|
urllib.error.HTTPError: On HTTP error responses
|
|
urllib.error.URLError: On connection errors
|
|
"""
|
|
t = timeout if timeout is not None else DISTRIBUTION_DEFAULT_TIMEOUT
|
|
req = urllib.request.Request(url, method="GET")
|
|
with urllib.request.urlopen(req, timeout=t) as resp:
|
|
return resp.status, resp.read()
|
|
|
|
|
|
def distribution_post_json(url, data_dict, timeout=None):
|
|
"""
|
|
LAN-only: POST JSON to a user-configured ComfyUI worker/master instance.
|
|
|
|
Args:
|
|
url: Full URL (caller constructs from master_url + endpoint)
|
|
data_dict: Dict to serialize as JSON body
|
|
timeout: Request timeout in seconds (default: 10)
|
|
|
|
Returns:
|
|
Tuple of (status_code, response_bytes)
|
|
|
|
Raises:
|
|
urllib.error.HTTPError: On HTTP error responses
|
|
urllib.error.URLError: On connection errors
|
|
"""
|
|
t = timeout if timeout is not None else DISTRIBUTION_DEFAULT_TIMEOUT
|
|
payload = json.dumps(data_dict).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
with urllib.request.urlopen(req, timeout=t) as resp:
|
|
return resp.status, resp.read()
|
|
|
|
|
|
def distribution_post_multipart(url, content_type, body, timeout=60):
|
|
"""
|
|
LAN-only: POST multipart/form-data to a user-configured ComfyUI instance.
|
|
|
|
Caller is responsible for building the multipart body and boundary.
|
|
This function handles only the HTTP transport.
|
|
|
|
Args:
|
|
url: Full URL
|
|
content_type: Full Content-Type header value (including boundary)
|
|
body: Pre-built multipart body bytes
|
|
timeout: Request timeout in seconds (default: 60)
|
|
|
|
Returns:
|
|
Tuple of (status_code, response_bytes)
|
|
|
|
Raises:
|
|
urllib.error.HTTPError: On HTTP error responses
|
|
urllib.error.URLError: On connection errors
|
|
"""
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=body,
|
|
headers={"Content-Type": content_type},
|
|
method="POST"
|
|
)
|
|
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
return resp.status, resp.read()
|
|
|
|
|
|
def distribution_stream_download(url, data_dict, target_path, timeout=3600,
|
|
progress_callback=None):
|
|
"""
|
|
LAN-only: POST JSON and stream response body to a file.
|
|
Uses atomic write pattern (downloads to .downloading temp file, then renames).
|
|
|
|
Args:
|
|
url: Full URL
|
|
data_dict: Dict to serialize as JSON body
|
|
target_path: Final destination file path
|
|
timeout: Request timeout in seconds (default: 3600 / 1 hour)
|
|
progress_callback: Optional callable(downloaded_bytes, total_bytes)
|
|
called after each chunk for progress reporting
|
|
|
|
Returns:
|
|
total_bytes downloaded
|
|
|
|
Raises:
|
|
urllib.error.HTTPError: On HTTP error responses
|
|
urllib.error.URLError: On connection errors
|
|
OSError: On file system errors
|
|
"""
|
|
payload = json.dumps(data_dict).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
url,
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST"
|
|
)
|
|
|
|
with urllib.request.urlopen(req, timeout=timeout) as resp:
|
|
total_size = int(resp.headers.get("Content-Length", 0))
|
|
temp_path = target_path + ".downloading"
|
|
downloaded = 0
|
|
|
|
with open(temp_path, "wb") as f:
|
|
while True:
|
|
chunk = resp.read(8 * 1024 * 1024) # 8 MB chunks
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
downloaded += len(chunk)
|
|
if progress_callback:
|
|
progress_callback(downloaded, total_size)
|
|
|
|
os.replace(temp_path, target_path)
|
|
return downloaded
|