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()