Jason Hoku
1f77100ea8
- Deletes the urllib-based remote_decode_hf() function (moved to companion). - Deletes the 'from .network_utils import huggingface_vae_decode...' import. - RemoteVAEDecodeWorker._worker() now calls _companion_decode() instead. The worker class (threading, manifest writing, dashboard events) is unchanged. Only the network atom is forwarded to ComfyUI-USCG-RemoteVAE companion.
302 lines
12 KiB
Python
302 lines
12 KiB
Python
import torch
|
|
import hashlib
|
|
import queue
|
|
import os
|
|
import threading
|
|
import time
|
|
import numpy as np
|
|
from PIL import Image
|
|
try:
|
|
from diffusers.image_processor import VaeImageProcessor
|
|
except ImportError as e:
|
|
VaeImageProcessor = None
|
|
print(f"[GridTester] ⚠️ Could not import diffusers: {e}")
|
|
print("[GridTester] ⚠️ Remote VAE features will be unavailable.")
|
|
print("[GridTester] ⚠️ To fix, run: pip install --upgrade diffusers>=0.25.0")
|
|
print("[GridTester] ⚠️ If another custom node requires an older version, you may need to resolve the conflict manually.")
|
|
|
|
|
|
|
|
INSTALL_INSTRUCTIONS = (
|
|
"Remote VAE requires the ComfyUI-USCG-RemoteVAE companion plugin.\n"
|
|
"Install via Comfy Manager (search 'USCG Remote VAE') or:\n"
|
|
" git clone https://github.com/JasonHoku/ComfyUI-USCG-RemoteVAE\n"
|
|
"into your ComfyUI/custom_nodes/ directory."
|
|
)
|
|
|
|
|
|
def is_remote_vae_available():
|
|
"""Return True if the ComfyUI-USCG-RemoteVAE companion plugin is loaded."""
|
|
try:
|
|
import comfyui_uscg_remote_vae # noqa: F401
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def get_endpoint_names():
|
|
"""Return list of HF endpoint names for UI dropdown.
|
|
Empty list if companion not installed."""
|
|
try:
|
|
from comfyui_uscg_remote_vae import list_endpoints
|
|
return [name for name, _url in list_endpoints()]
|
|
except ImportError:
|
|
return []
|
|
|
|
|
|
def _companion_decode(endpoint, tensor, height, width):
|
|
"""Forward a decode call to the companion plugin.
|
|
|
|
Raises:
|
|
RuntimeError: companion not installed (with install instructions)
|
|
or companion version too old.
|
|
"""
|
|
try:
|
|
from comfyui_uscg_remote_vae import decode, __version__
|
|
except ImportError:
|
|
raise RuntimeError(INSTALL_INSTRUCTIONS)
|
|
if tuple(map(int, __version__.split(".")[:2])) < (0, 1):
|
|
raise RuntimeError(
|
|
f"ComfyUI-USCG-RemoteVAE >= 0.1.0 required (found {__version__}). "
|
|
"Update via Comfy Manager."
|
|
)
|
|
return decode(endpoint, tensor, height, width)
|
|
|
|
|
|
def detect_model_type(model, latent_channels):
|
|
"""
|
|
Auto-detect model type for HF Remote VAE endpoint selection
|
|
"""
|
|
#print(f"[GridTester] 🔍 Detecting model type... latent_channels={latent_channels}")
|
|
|
|
if latent_channels == 16:
|
|
# Could be Flux, SD3, or HunyuanVideo
|
|
if hasattr(model, 'model'):
|
|
model_str = str(type(model.model)).lower()
|
|
if 'xl' in model_str:
|
|
return "SDXL"
|
|
if hasattr(model, 'model') and hasattr(model.model, 'model_type'):
|
|
model_type = str(model.model.model_type).lower()
|
|
#print(f"[GridTester] 🔍 Model type attribute: {model_type}")
|
|
if 'flux' in model_type:
|
|
return "Flux"
|
|
elif 'hunyuan' in model_type or 'video' in model_type:
|
|
return "HunyuanVideo"
|
|
|
|
# Check model name/path
|
|
if hasattr(model, 'model') and hasattr(model.model, 'model_config'):
|
|
config = str(model.model.model_config).lower()
|
|
if 'flux' in config:
|
|
return "Flux"
|
|
|
|
# Default to Flux for 16-channel
|
|
#print(f"[GridTester] 🔍 Defaulting to Flux for 16 channels")
|
|
return "Flux"
|
|
|
|
elif latent_channels == 4:
|
|
# SD or SDXL - check resolution or model size
|
|
# 1. Check model name FIRST (lightweight) ✅
|
|
if hasattr(model, 'model'):
|
|
model_str = str(type(model.model)).lower()
|
|
if 'xl' in model_str or 'sdxl' in model_str:
|
|
return "SDXL"
|
|
|
|
# 2. Check model config (lightweight) ✅
|
|
if hasattr(model, 'model') and hasattr(model.model, 'model_config'):
|
|
config_str = str(model.model.model_config).lower()
|
|
if 'sdxl' in config_str:
|
|
return "SDXL"
|
|
|
|
# 3. ONLY count parameters as LAST RESORT ✅
|
|
with torch.no_grad():
|
|
try:
|
|
param_count = sum(p.numel() for p in model.model.diffusion_model.parameters())
|
|
#print(f"[GridTester] 🔍 Model parameter count: {param_count:,}")
|
|
|
|
if param_count > 1_000_000_000: # > 1B params suggests SDXL
|
|
#print(f"[GridTester] 🔍 Detected SDXL (>1B params)")
|
|
return "SDXL"
|
|
else:
|
|
#print(f"[GridTester] 🔍 Detected SD1.5 (<1B params)")
|
|
return "SD"
|
|
except:
|
|
pass
|
|
|
|
# Fallback: check if model name contains 'xl'
|
|
if hasattr(model, 'model'):
|
|
model_str = str(type(model.model)).lower()
|
|
if 'xl' in model_str or 'sdxl' in model_str:
|
|
#print(f"[GridTester] 🔍 Detected SDXL from model name")
|
|
return "SDXL"
|
|
|
|
#print(f"[GridTester] 🔍 Defaulting to SD for 4 channels")
|
|
return "SD"
|
|
|
|
else:
|
|
#print(f"[GridTester] ⚠️ Unknown latent channels: {latent_channels}, defaulting to SD")
|
|
return "SD"
|
|
|
|
|
|
class RemoteVAEDecodeWorker:
|
|
"""
|
|
Background worker thread for async remote VAE decoding
|
|
"""
|
|
def __init__(self, endpoint, img_dir, manifest_path, existing_data, session_name, unique_id, image_format="webp"):
|
|
self.endpoint = endpoint
|
|
self.img_dir = img_dir
|
|
self.manifest_path = manifest_path
|
|
self.existing_data = existing_data
|
|
self.session_name = session_name
|
|
self.unique_id = unique_id
|
|
fmt = (image_format or "webp").lower().strip()
|
|
if fmt not in ("webp", "png", "jpg", "jpeg"):
|
|
fmt = "webp"
|
|
self.image_format = fmt
|
|
self.queue = queue.Queue()
|
|
self.thread = threading.Thread(target=self._worker, daemon=True)
|
|
self.thread.start()
|
|
self.total_decoded = 0
|
|
|
|
def _worker(self):
|
|
"""Worker thread that processes decode requests"""
|
|
while True:
|
|
item = self.queue.get()
|
|
|
|
if item is None: # Poison pill to stop thread
|
|
self.queue.task_done()
|
|
break
|
|
|
|
try:
|
|
latent_tensor, meta, height, width = item
|
|
|
|
# print(f"[GridTester] 🌐 Processing remote decode for image #{meta['id']}")
|
|
# print(f"[GridTester] 🌐 Input latent shape: {latent_tensor.shape}")
|
|
|
|
# Ensure batch dimension exists
|
|
if latent_tensor.ndim == 3:
|
|
latent_tensor = latent_tensor.unsqueeze(0)
|
|
print(f"[GridTester] 🌐 Added batch dim: {latent_tensor.shape}")
|
|
|
|
# Remote decode - returns [B, C, H, W] tensor
|
|
decoded = _companion_decode(self.endpoint, latent_tensor, height, width)
|
|
|
|
# print(f"[GridTester] 🌐 Decoded shape: {decoded.shape}")
|
|
# print(f"[GridTester] 🌐 Decoded dtype: {decoded.dtype}")
|
|
|
|
# Use VaeImageProcessor to properly postprocess the VAE output
|
|
# This handles denormalization from [-1, 1] to [0, 1] and format conversion
|
|
if VaeImageProcessor is None:
|
|
raise ImportError("diffusers is not installed or outdated. Fix: pip install --upgrade diffusers")
|
|
image_processor = VaeImageProcessor(vae_scale_factor=8)
|
|
image_processor.config.do_resize = False
|
|
|
|
# Get the result and check what format it's in
|
|
result = image_processor.postprocess(decoded, output_type="pt")
|
|
#print(f"[GridTester] 🔧 Postprocessed shape: {result.shape}")
|
|
#print(f"[GridTester] 🔧 Postprocessed dtype: {result.dtype}")
|
|
|
|
# The postprocess with output_type="pt" returns [B, C, H, W] in [0, 1]
|
|
# We need to convert to [H, W, C] for PIL
|
|
image = result[0] # Remove batch: [C, H, W]
|
|
image = image.permute(1, 2, 0) # Convert to [H, W, C]
|
|
image = image.cpu().numpy()
|
|
image_np = (image * 255).round().astype(np.uint8)
|
|
|
|
#print(f"[GridTester] 🔧 Final image shape: {image_np.shape}, dtype: {image_np.dtype}")
|
|
|
|
# Create PIL Image
|
|
img = Image.fromarray(image_np)
|
|
|
|
# Save image in chosen format
|
|
fmt = self.image_format # Already normalised in __init__
|
|
ext = "jpg" if fmt in ("jpg", "jpeg") else fmt
|
|
filename = f"img_{meta['id']}.{ext}"
|
|
filepath = os.path.join(self.img_dir, filename)
|
|
if fmt == "png":
|
|
img.save(filepath, format="PNG")
|
|
elif fmt in ("jpg", "jpeg"):
|
|
img.save(filepath, format="JPEG", quality=95)
|
|
else:
|
|
img.save(filepath, quality=80)
|
|
|
|
meta.update({
|
|
"file": f"/view?filename={filename}&type=output&subfolder=benchmarks/{self.session_name}/images",
|
|
"rejected": False
|
|
})
|
|
|
|
# Update manifest
|
|
import json
|
|
self.existing_data["items"].insert(0, meta)
|
|
if os.path.exists(self.manifest_path):
|
|
try:
|
|
with open(self.manifest_path, "r") as f:
|
|
disk_manifest = json.load(f)
|
|
|
|
# Create a lookup map for items currently in memory
|
|
# This allows us to update self.existing_data in-place
|
|
memory_items_map = {
|
|
i.get("id"): i
|
|
for i in self.existing_data.get("items", [])
|
|
if "id" in i
|
|
}
|
|
|
|
# Check every item on disk. If it exists in memory, copy the tags over.
|
|
for disk_item in disk_manifest.get("items", []):
|
|
d_id = disk_item.get("id")
|
|
if d_id and d_id in memory_items_map:
|
|
local_item = memory_items_map[d_id]
|
|
|
|
# PRESERVE TAGS: Copy these keys from disk to memory
|
|
if "favorited" in disk_item:
|
|
local_item["favorited"] = disk_item["favorited"]
|
|
if "rejected" in disk_item:
|
|
local_item["rejected"] = disk_item["rejected"]
|
|
|
|
except Exception as e:
|
|
print(f"[GridTester] ⚠️ Error syncing with disk manifest: {e}")
|
|
|
|
# 3. Now it is safe to save (Memory is now 'fresh' with disk changes)
|
|
with open(self.manifest_path, "w") as f:
|
|
json.dump(self.existing_data, f, indent=4)
|
|
|
|
# print("Save Manifest at remote vae")
|
|
# Send update to dashboard
|
|
try:
|
|
from server import PromptServer
|
|
if PromptServer:
|
|
# Get meta, use empty dict if not present
|
|
manifest_meta = self.existing_data.get("meta", {})
|
|
|
|
PromptServer.instance.send_sync("ultimate_grid.update", {
|
|
"node": self.unique_id,
|
|
"session_name": self.session_name,
|
|
"new_items": [meta],
|
|
"meta": manifest_meta
|
|
})
|
|
except (ImportError, KeyError):
|
|
# Silently ignore dashboard update errors
|
|
pass
|
|
|
|
self.total_decoded += 1
|
|
# print(f"[GridTester] ✅ Remote VAE decoded #{meta['id']} ({self.total_decoded} total)")
|
|
|
|
except Exception as e:
|
|
print(f"[GridTester] ❌ Remote VAE worker error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
finally:
|
|
self.queue.task_done()
|
|
|
|
|
|
def add_job(self, latent_tensor, meta, height, width):
|
|
"""Add a decode job to the queue"""
|
|
self.queue.put((latent_tensor, meta, height, width))
|
|
|
|
def wait_completion(self):
|
|
"""Wait for all jobs to complete"""
|
|
self.queue.join()
|
|
|
|
def stop(self):
|
|
"""Stop the worker thread"""
|
|
self.queue.put(None)
|
|
self.thread.join() |