From e848839b2f360eb9e3079bf27745a16f044ded9f Mon Sep 17 00:00:00 2001 From: PozzettiAndrea Date: Sun, 1 Mar 2026 14:40:30 +0000 Subject: [PATCH] utils update --- nodes/image_utils.py | 24 ++++++++++++------------ nodes/load_image_exif.py | 2 +- nodes/predict.py | 28 ++++++++++++++++++++-------- requirements.txt | 2 +- 4 files changed, 34 insertions(+), 22 deletions(-) diff --git a/nodes/image_utils.py b/nodes/image_utils.py index 55b7752..205251e 100644 --- a/nodes/image_utils.py +++ b/nodes/image_utils.py @@ -15,55 +15,55 @@ import torch from PIL import Image -# ── ComfyUI ↔ numpy ────────────────────────────────────────────── +# -- ComfyUI <-> numpy --------------------------------------------------- def comfy_to_numpy(image: torch.Tensor, index: int = 0) -> np.ndarray: - """ComfyUI IMAGE [B,H,W,C] float32 → numpy [H,W,C] uint8.""" + """ComfyUI IMAGE [B,H,W,C] float32 -> numpy [H,W,C] uint8.""" return (image[index].cpu().numpy() * 255).clip(0, 255).astype(np.uint8) def numpy_to_comfy(image: np.ndarray) -> torch.Tensor: - """numpy [H,W,C] uint8 → ComfyUI IMAGE [1,H,W,C] float32.""" + """numpy [H,W,C] uint8 -> ComfyUI IMAGE [1,H,W,C] float32.""" return torch.from_numpy(image.astype(np.float32) / 255.0).unsqueeze(0) -# ── ComfyUI ↔ PIL ──────────────────────────────────────────────── +# -- ComfyUI <-> PIL ----------------------------------------------------- def comfy_to_pil(image: torch.Tensor, index: int = 0) -> Image.Image: - """ComfyUI IMAGE [B,H,W,C] float32 → PIL Image (RGB).""" + """ComfyUI IMAGE [B,H,W,C] float32 -> PIL Image (RGB).""" return Image.fromarray(comfy_to_numpy(image, index)) def pil_to_comfy(image: Image.Image) -> torch.Tensor: - """PIL Image → ComfyUI IMAGE [1,H,W,C] float32.""" + """PIL Image -> ComfyUI IMAGE [1,H,W,C] float32.""" return numpy_to_comfy(np.array(image.convert("RGB"))) -# ── ComfyUI ↔ CHW (model input) ───────────────────────────────── +# -- ComfyUI <-> CHW (model input) --------------------------------------- def comfy_to_chw(image: torch.Tensor) -> torch.Tensor: - """ComfyUI IMAGE [B,H,W,C] → model input [B,C,H,W].""" + """ComfyUI IMAGE [B,H,W,C] -> model input [B,C,H,W].""" return image.movedim(-1, 1) def chw_to_comfy(image: torch.Tensor) -> torch.Tensor: - """Model output [B,C,H,W] → ComfyUI IMAGE [B,H,W,C].""" + """Model output [B,C,H,W] -> ComfyUI IMAGE [B,H,W,C].""" return image.movedim(1, -1) -# ── Mask conversions ───────────────────────────────────────────── +# -- Mask conversions ---------------------------------------------------- def mask_to_image(mask: torch.Tensor) -> torch.Tensor: - """ComfyUI MASK [B,H,W] → ComfyUI IMAGE [B,H,W,3] (grayscale RGB).""" + """ComfyUI MASK [B,H,W] -> ComfyUI IMAGE [B,H,W,3] (grayscale RGB).""" return mask.unsqueeze(-1).expand(-1, -1, -1, 3) def image_to_mask(image: torch.Tensor) -> torch.Tensor: - """ComfyUI IMAGE [B,H,W,C] → ComfyUI MASK [B,H,W] (luminance).""" + """ComfyUI IMAGE [B,H,W,C] -> ComfyUI MASK [B,H,W] (luminance).""" if image.shape[-1] == 1: return image.squeeze(-1) # ITU-R BT.601 luminance diff --git a/nodes/load_image_exif.py b/nodes/load_image_exif.py index 847c99f..0bf9c0a 100644 --- a/nodes/load_image_exif.py +++ b/nodes/load_image_exif.py @@ -116,7 +116,7 @@ class LoadImageWithExif(io.ComfyNode): description="Load an image and extract focal length from EXIF metadata (35mm equivalent).", inputs=[ io.Combo.Input("image", options=_get_image_files(), - image_upload=True), + upload=io.UploadType.image), io.Float.Input("default_focal_mm", default=30.0, min=1.0, max=500.0, step=0.1, optional=True, tooltip="Default focal length (35mm equiv) if not found in EXIF"), diff --git a/nodes/predict.py b/nodes/predict.py index 46501e9..4f6d336 100644 --- a/nodes/predict.py +++ b/nodes/predict.py @@ -33,6 +33,18 @@ _encode_cache = { } +def _monodepth_to(output, device): + """Move all tensors in a MonodepthOutput to the given device.""" + from .sharp.model import MonodepthOutput + return MonodepthOutput( + disparity=output.disparity.to(device), + encoder_features=[t.to(device) for t in output.encoder_features], + decoder_features=output.decoder_features.to(device), + output_features=[t.to(device) for t in output.output_features], + intermediate_features=[t.to(device) for t in output.intermediate_features], + ) + + def _compute_image_hash(image_np: np.ndarray) -> str: """Compute hash of image for cache key.""" return hashlib.sha256(image_np.tobytes()).hexdigest()[:16] @@ -242,8 +254,8 @@ def _predict_image_cached( if _encode_cache["image_hash"] == image_hash: # Cache hit - reuse encoded features log.info("Cache hit - reusing encoded features (focal_length change is instant)") - monodepth_output = _encode_cache["monodepth_output"] - image_resized_pt = _encode_cache["image_resized"] + monodepth_output = _monodepth_to(_encode_cache["monodepth_output"], device) + image_resized_pt = _encode_cache["image_resized"].to(device) else: # Cache miss - need to encode log.info("Encoding...") @@ -270,15 +282,15 @@ def _predict_image_cached( monodepth_output, _ = predictor.encode(image_resized_pt) log.info(f"Encode time: {time.time() - encode_start:.2f}s") + # Update cache (store on CPU to free VRAM) + _encode_cache["image_hash"] = image_hash + _encode_cache["monodepth_output"] = _monodepth_to(monodepth_output, "cpu") + _encode_cache["image_resized"] = image_resized_pt.cpu() + _encode_cache["original_shape"] = (height, width) + # Release fragmented GPU memory after heavy encode (35 ViT passes) comfy.model_management.soft_empty_cache() - # Update cache - _encode_cache["image_hash"] = image_hash - _encode_cache["monodepth_output"] = monodepth_output - _encode_cache["image_resized"] = image_resized_pt - _encode_cache["original_shape"] = (height, width) - # Decode - always run with current focal length disparity_factor = torch.tensor([f_px / width]).float().to(device) diff --git a/requirements.txt b/requirements.txt index e1fae2f..4401e06 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,4 +13,4 @@ pillow-heif>=0.10.0 huggingface-hub>=0.16.0 # Hugging Face transfer extension (faster downloads on Windows) -hf_xet +hf_xet \ No newline at end of file