| import base64 |
| import hashlib |
| import os |
| import sys |
| import threading |
| from collections import OrderedDict |
| from concurrent.futures import ThreadPoolExecutor, as_completed |
| from io import BytesIO |
| from typing import Callable, Dict, List, Optional, Tuple |
|
|
| import time |
|
|
| import numpy as np |
| import torch |
| from PIL import Image, ImageDraw, ImageFont |
|
|
| import comfy.model_management |
|
|
| MODULE_DIR = os.path.dirname(os.path.abspath(__file__)) |
| if MODULE_DIR not in sys.path: |
| sys.path.insert(0, MODULE_DIR) |
|
|
| from logger import logger |
|
|
|
|
| class ImageCodec: |
| """负责 Tensor/Base64 转换、缓存与实时预览构建。""" |
|
|
| def __init__( |
| self, |
| logger_instance=logger, |
| ensure_not_interrupted: Optional[Callable[[], None]] = None, |
| cache_size: int = 16, |
| ): |
| self.logger = logger_instance |
| self.ensure_not_interrupted = ensure_not_interrupted |
| self._image_cache: "OrderedDict[str, str]" = OrderedDict() |
| self._cache_lock = threading.Lock() |
| self._cache_size = max(1, cache_size) |
|
|
| def _maybe_interrupt(self): |
| if self.ensure_not_interrupted: |
| self.ensure_not_interrupted() |
|
|
| def _tensor_cache_key( |
| self, |
| tensor: Optional[torch.Tensor] = None, |
| np_data: Optional[np.ndarray] = None, |
| ) -> Optional[str]: |
| if tensor is None and np_data is None: |
| return None |
| try: |
| target = np_data |
| if target is None: |
| target = tensor.detach().cpu().numpy() |
| return hashlib.sha1(target.tobytes()).hexdigest() |
| except Exception: |
| return None |
|
|
| def _get_cached_image_b64(self, cache_key: Optional[str]) -> Optional[str]: |
| if not cache_key: |
| return None |
| with self._cache_lock: |
| value = self._image_cache.get(cache_key) |
| if value is not None: |
| self._image_cache.move_to_end(cache_key) |
| return value |
|
|
| def _set_cached_image_b64(self, cache_key: Optional[str], value: str) -> None: |
| if not cache_key or not value: |
| return |
| with self._cache_lock: |
| self._image_cache[cache_key] = value |
| self._image_cache.move_to_end(cache_key) |
| while len(self._image_cache) > self._cache_size: |
| self._image_cache.popitem(last=False) |
|
|
| def extract_numpy_images(self, tensor: torch.Tensor) -> List[np.ndarray]: |
| images: List[np.ndarray] = [] |
| if tensor is None: |
| return images |
| try: |
| np_data = tensor.detach().cpu().numpy() |
| except Exception as exc: |
| self.logger.error(f"输入图像转换失败: {exc}") |
| return images |
|
|
| if np_data.ndim == 3: |
| np_data = np_data[np.newaxis, ...] |
| np_data = np.clip(np_data, 0.0, 1.0) |
|
|
| for sample in np_data: |
| if sample.ndim == 2: |
| sample = np.expand_dims(sample, axis=-1) |
| if sample.shape[-1] == 1: |
| sample = np.repeat(sample, 3, axis=-1) |
| images.append(np.ascontiguousarray(sample)) |
| return images |
|
|
| def tensor_to_base64( |
| self, |
| tensor: Optional[torch.Tensor] = None, |
| np_image: Optional[np.ndarray] = None, |
| ) -> str: |
| if np_image is None: |
| if tensor is None: |
| raise ValueError("必须提供 tensor 或 numpy 图像数据用于编码") |
| samples = self.extract_numpy_images(tensor) |
| if not samples: |
| raise ValueError("无法从 tensor 中提取有效图像数据") |
| np_image = samples[0] |
|
|
| img_array = np.clip(np_image, 0.0, 1.0) |
| img_uint8 = (img_array * 255).astype(np.uint8) |
| img = Image.fromarray(img_uint8) |
| buffered = BytesIO() |
| img.save(buffered, format="PNG") |
| return base64.b64encode(buffered.getvalue()).decode() |
|
|
| def prepare_input_images(self, tensors: List[torch.Tensor]) -> List[str]: |
| if not tensors: |
| return [] |
| encoded_images: List[str] = [] |
| for tensor in tensors: |
| if tensor is None: |
| continue |
| for sample in self.extract_numpy_images(tensor): |
| cache_key = self._tensor_cache_key(np_data=sample) |
| cached_value = self._get_cached_image_b64(cache_key) |
| if cached_value is None: |
| base64_value = self.tensor_to_base64(np_image=sample) |
| self._set_cached_image_b64(cache_key, base64_value) |
| else: |
| base64_value = cached_value |
| encoded_images.append(base64_value) |
| return encoded_images |
|
|
| def base64_to_tensor_single(self, b64_str: str) -> np.ndarray: |
| try: |
| img_data = base64.b64decode(b64_str) |
| img = Image.open(BytesIO(img_data)).convert('RGB') |
| img_array = np.array(img).astype(np.float32) / 255.0 |
| return img_array |
| except Exception as exc: |
| self.logger.error(f"图片解码失败: {exc}") |
| return np.zeros((64, 64, 3), dtype=np.float32) |
|
|
| def base64_to_tensor_parallel( |
| self, |
| base64_strings: List[str], |
| log_prefix: Optional[str] = None, |
| max_workers: Optional[int] = None, |
| ) -> torch.Tensor: |
| if not isinstance(base64_strings, list) or len(base64_strings) == 0: |
| return torch.zeros((1, 64, 64, 3), dtype=torch.float32) |
|
|
| decode_start = time.time() |
| images = [] |
| worker_cap = max_workers if max_workers is not None else max(4, os.cpu_count() or 1) |
| worker_cap = max(1, worker_cap) |
| effective_workers = min(worker_cap, len(base64_strings)) |
|
|
| self._maybe_interrupt() |
| executor = ThreadPoolExecutor(max_workers=effective_workers) |
| try: |
| future_to_index = {executor.submit(self.base64_to_tensor_single, b64): i |
| for i, b64 in enumerate(base64_strings)} |
|
|
| results = [None] * len(base64_strings) |
| try: |
| for future in as_completed(future_to_index): |
| index = future_to_index[future] |
| try: |
| self._maybe_interrupt() |
| results[index] = future.result() |
| except comfy.model_management.InterruptProcessingException: |
| for pending in future_to_index: |
| pending.cancel() |
| raise |
| except Exception as exc: |
| self.logger.error(f"图片{index+1}解码异常: {exc}") |
| results[index] = np.zeros((64, 64, 3), dtype=np.float32) |
|
|
| images = [r for r in results if r is not None] |
| except comfy.model_management.InterruptProcessingException: |
| executor.shutdown(wait=False, cancel_futures=True) |
| raise |
| finally: |
| if not executor._shutdown: |
| executor.shutdown(wait=False, cancel_futures=True) |
|
|
| decode_time = time.time() - decode_start |
| prefix = log_prefix or "" |
| prefix = f"{prefix} " if prefix else "" |
| self.logger.success(f"{prefix}并发解码 {len(images)} 张图片完成，耗时: {decode_time:.2f}s") |
|
|
| return torch.from_numpy(np.stack(images)) |
|
|
| def build_preview_tuple( |
| self, |
| tensor: Optional[torch.Tensor], |
| batch_index: int, |
| max_size: int = 512, |
| ) -> Optional[Tuple[str, Image.Image, int]]: |
| if tensor is None or tensor.shape[0] == 0: |
| return None |
|
|
| try: |
| preview_tensor = tensor[0].detach().cpu() |
| preview_tensor = torch.clamp(preview_tensor, 0.0, 1.0) |
| preview_array = (preview_tensor.numpy() * 255).astype(np.uint8) |
|
|
| if preview_array.ndim == 3 and preview_array.shape[2] == 1: |
| preview_array = np.repeat(preview_array, 3, axis=2) |
| elif preview_array.ndim == 2: |
| preview_array = np.stack([preview_array] * 3, axis=2) |
|
|
| preview_image = Image.fromarray(preview_array) |
| return ("PNG", preview_image, max_size) |
| except Exception as exc: |
| self.logger.error(f"实时预览生成失败: 批次 {batch_index + 1}: {str(exc)[:80]}") |
| return None |
|
|
|
|
| class ErrorCanvas: |
| """统一的错误画布与字体管理。""" |
|
|
| def __init__(self, logger_instance=logger): |
| self.logger = logger_instance |
| self._font_cache: Dict[int, ImageFont.ImageFont] = {} |
|
|
| def _get_error_font_paths(self) -> List[str]: |
| candidates = [] |
| windir = os.environ.get("WINDIR") |
| if windir: |
| for name in ("msyh.ttc", "msyh.ttf", "msjh.ttc", "simhei.ttf", "msmincho.ttc"): |
| candidates.append(os.path.join(windir, "Fonts", name)) |
| candidates.append(os.path.join(os.path.dirname(__file__), "msyh.ttc")) |
| return candidates |
|
|
| def _load_error_font(self, size: int) -> ImageFont.ImageFont: |
| cached = self._font_cache.get(size) |
| if cached is not None: |
| return cached |
| for font_path in self._get_error_font_paths(): |
| if font_path and os.path.exists(font_path): |
| try: |
| font = ImageFont.truetype(font_path, size) |
| self._font_cache[size] = font |
| return font |
| except Exception: |
| continue |
| fallback = ImageFont.load_default() |
| self._font_cache[size] = fallback |
| return fallback |
|
|
| @staticmethod |
| def _wrap_text_segments(draw: ImageDraw.ImageDraw, text: str, |
| font: ImageFont.ImageFont, max_width: int) -> List[str]: |
| if not text: |
| return [""] |
| segments: List[str] = [] |
| current = "" |
| for ch in text: |
| tentative = current + ch |
| if draw.textlength(tentative, font=font) <= max_width or not current: |
| current = tentative |
| else: |
| segments.append(current) |
| current = ch |
| if current: |
| segments.append(current) |
| return segments |
|
|
| def build_error_image_tensor( |
| self, |
| title: str, |
| lines: List[str], |
| size: Tuple[int, int] = (640, 640), |
| ) -> torch.Tensor: |
| lines = [line.strip() for line in lines if line and line.strip()] |
| if not lines: |
| lines = ["发生未知错误"] |
|
|
| width, height = size |
| background = (248, 248, 248) |
| accent = (255, 235, 235) |
| title_color = (180, 30, 30) |
| text_color = (45, 45, 45) |
|
|
| img = Image.new("RGB", (width, height), background) |
| draw = ImageDraw.Draw(img) |
| font_title = self._load_error_font(26) |
| font_body = self._load_error_font(18) |
|
|
| margin = 32 |
| y = margin |
| max_text_width = max(10, width - 2 * margin) |
| max_y = height - margin |
|
|
| draw.rectangle([margin - 6, margin - 6, width - margin + 6, y + 40], fill=accent) |
| draw.text((margin, y), title, fill=title_color, font=font_title) |
| y += font_title.getbbox(title)[3] - font_title.getbbox(title)[1] + 16 |
|
|
| for line in lines: |
| wrapped = self._wrap_text_segments(draw, line, font_body, max_text_width) |
| for seg in wrapped: |
| bbox = font_body.getbbox(seg) |
| line_height = bbox[3] - bbox[1] + 6 |
| if y + line_height > max_y: |
| break |
| draw.text((margin, y), seg, fill=text_color, font=font_body) |
| y += line_height |
|
|
| arr = np.array(img).astype(np.float32) / 255.0 |
| return torch.from_numpy(arr).unsqueeze(0) |
|
|
| def build_error_tensor_from_text(self, title: str, text: str) -> torch.Tensor: |
| normalized = text.replace("\r\n", "\n").replace("\r", "\n") |
| lines = [line.strip() for line in normalized.split("\n") if line.strip()] |
| if not lines: |
| lines = ["发生未知错误"] |
| return self.build_error_image_tensor(title, lines) |
|
|
