import torch import requests import hashlib import queue import os import threading import time from safetensors.torch import _tobytes import numpy as np from PIL import Image from diffusers.image_processor import VaeImageProcessor # HF Remote VAE endpoints HF_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/" } 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" def remote_decode_hf(endpoint, tensor, height, width): """ Send latent to HuggingFace Remote VAE endpoint for decoding (Simpler version using requests params) """ try: import requests # Ensure tensor is on CPU and contiguous tensor = tensor.cpu().contiguous() #print(f"[GridTester] 🌐 Remote decode:") #print(f" Tensor shape: {tensor.shape}") #print(f" Tensor dtype: {tensor.dtype}") #print(f" Height: {height}, Width: {width}") headers = { "Content-Type": "tensor/binary", "Accept": "tensor/binary" } # Build parameters - requests will handle repeated 'shape' params # by passing shape as a list in params dict params = { "do_scaling": False, "output_type": "pt", "partial_postprocess": False, "shape": [int(dim) for dim in tensor.shape], # List creates repeated params! "dtype": str(tensor.dtype).split(".")[-1], "height": int(height), "width": int(width) } #print(f"[GridTester] 🌐 Parameters: {params}") # Convert tensor to bytes tensor_data = _tobytes(tensor, "tensor") #print(f"[GridTester] 🌐 Tensor bytes size: {len(tensor_data)}") # requests automatically converts list params to repeated query params! # shape=[1,2,3] becomes ?shape=1&shape=2&shape=3 response = requests.post( endpoint, params=params, # ← requests handles the shape list properly data=tensor_data, headers=headers, timeout=60 ) if not response.ok: error_text = response.text #print(f"[GridTester] ❌ Remote VAE error: {error_text}") #print(f"[GridTester] 🌐 Request URL was: {response.request.url}") raise RuntimeError(f"Remote VAE decode failed: {error_text}") # Parse response output_tensor = response.content response_headers = response.headers #print(f"[GridTester] 🌐 Response headers: {dict(response_headers)}") # The response should have shape and dtype in headers # But format might vary - let's handle both cases shape_header = response_headers.get("shape", "") if shape_header: try: # Try JSON parsing first import json shape = json.loads(shape_header) except: # Fallback: parse comma-separated values shape = [int(x.strip()) for x in shape_header.split(",")] else: raise RuntimeError("No shape header in response") dtype_str = response_headers.get("dtype", "float32") dtype_map = { "float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16, } dtype = dtype_map.get(dtype_str, torch.float32) # Map to numpy dtype for frombuffer numpy_dtype_map = { "float32": np.float32, "float16": np.float16, "bfloat16": np.float32, # NumPy doesn't support bfloat16 } numpy_dtype = numpy_dtype_map.get(dtype_str, np.float32) #print(f"[GridTester] 🌐 Parsed shape: {shape}, dtype: {dtype}, numpy_dtype: {numpy_dtype}") # Convert bytes back to tensor using correct dtype tensor_np = np.frombuffer(output_tensor, dtype=numpy_dtype) result = torch.from_numpy(tensor_np).reshape(shape).to(dtype) #print(f"[GridTester] 🌐 Result tensor shape: {result.shape}") return result except Exception as e: #print(f"[GridTester] ❌ Remote decode error: {e}") import traceback traceback.print_exc() raise class RemoteVAEDecodeWorker: """ Background worker thread for async remote VAE decoding """ def __init__(self, endpoint, img_dir, manifest_path, existing_data, session_name, unique_id): 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 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 = remote_decode_hf(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 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 filename = f"img_{meta['id']}.webp" print("Saving") print(filename) print(self.img_dir) img.save(os.path.join(self.img_dir, filename), 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()