import hashlib import os import numpy as np import torch from PIL import Image, ImageOps, ImageSequence HEIC_EXTS = {".heic", ".heif"} def _register_preview_route_if_possible() -> None: """Register /heic_preview once the PromptServer is available.""" try: from aiohttp import web from io import BytesIO from server import PromptServer instance = getattr(PromptServer, "instance", None) if instance is None or not hasattr(instance, "routes"): return if getattr(instance, "_heic_preview_route_registered", False): return instance._heic_preview_route_registered = True @instance.routes.get("/heic_preview") async def heic_preview(request: web.Request): filename = request.rel_url.query.get("filename") if not filename: return web.Response(status=400, text="filename is required") folder_paths, _ = _get_comfy_modules() if not folder_paths.exists_annotated_filepath(filename): return web.Response(status=404, text="file not found") image_path = folder_paths.get_annotated_filepath(filename) _try_register_heif_opener() try: img = Image.open(image_path) img = ImageOps.exif_transpose(img) img = img.convert("RGBA") except Exception as e: return web.Response(status=500, text=f"failed to decode image: {e}") bio = BytesIO() img.save(bio, format="PNG") return web.Response(body=bio.getvalue(), content_type="image/png") except Exception: return def _get_comfy_modules(): try: import folder_paths # provided by ComfyUI import node_helpers # provided by ComfyUI except Exception as e: raise RuntimeError( "This node must run inside a ComfyUI environment (missing folder_paths/node_helpers)." ) from e return folder_paths, node_helpers def _try_register_heif_opener() -> bool: """Register HEIF/HEIC opener for Pillow. Returns True if available.""" try: from pillow_heif import register_heif_opener # type: ignore register_heif_opener() return True except Exception: return False def _is_heic_path(name: str) -> bool: _, ext = os.path.splitext(str(name)) return ext.lower() in HEIC_EXTS def _list_heic_files_in_input_dir() -> list[str]: folder_paths, _ = _get_comfy_modules() input_dir = folder_paths.get_input_directory() try: files = [ f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) ] except Exception: return [] allowed_ext = {".heic", ".heif"} out = [] for f in files: _, ext = os.path.splitext(f) if ext.lower() in allowed_ext: out.append(f) return sorted(out) def _list_image_files_in_input_dir() -> list[str]: """Return all supported image files (PNG, JPG, WEBP, HEIC/HEIF).""" folder_paths, _ = _get_comfy_modules() input_dir = folder_paths.get_input_directory() try: files = [ f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) ] except Exception: return [] # Support all common image formats + HEIC allowed_ext = {".png", ".jpg", ".jpeg", ".webp", ".heic", ".heif", ".bmp", ".gif"} out = [] for f in files: _, ext = os.path.splitext(f) if ext.lower() in allowed_ext: out.append(f) return sorted(out) class LoadImagePlusHEIC: @classmethod def INPUT_TYPES(cls): _register_preview_route_if_possible() files = _list_image_files_in_input_dir() return { "required": { "image": (files, {"image_upload": True}), }, } CATEGORY = "image" RETURN_TYPES = ("IMAGE", "MASK") FUNCTION = "load_image" def load_image(self, image): if _is_heic_path(image) and not _try_register_heif_opener(): raise RuntimeError( "HEIC/HEIF support is not available. Install dependency: pip install pillow-heif" ) else: _try_register_heif_opener() folder_paths, node_helpers = _get_comfy_modules() if not folder_paths.exists_annotated_filepath(image): raise FileNotFoundError(f"Image not found in input path: {image}") image_path = folder_paths.get_annotated_filepath(image) try: img = node_helpers.pillow(Image.open, image_path) except Exception as e: if _is_heic_path(image): raise RuntimeError( "Failed to open HEIC/HEIF image. Ensure pillow-heif is installed: pip install pillow-heif\n" f"File: {image}\nError: {e}" ) raise RuntimeError(f"Failed to open image: {image}\nError: {e}") output_images = [] output_masks = [] w, h = None, None excluded_formats = ["MPO"] for i in ImageSequence.Iterator(img): i = node_helpers.pillow(ImageOps.exif_transpose, i) if i.mode == "I": i = i.point(lambda i: i * (1 / 255)) image_rgb = i.convert("RGB") if len(output_images) == 0: w = image_rgb.size[0] h = image_rgb.size[1] if image_rgb.size[0] != w or image_rgb.size[1] != h: continue image_np = np.array(image_rgb).astype(np.float32) / 255.0 image_t = torch.from_numpy(image_np)[None,] if "A" in i.getbands(): mask_np = np.array(i.getchannel("A")).astype(np.float32) / 255.0 mask_t = 1.0 - torch.from_numpy(mask_np) elif i.mode == "P" and "transparency" in i.info: mask_np = ( np.array(i.convert("RGBA").getchannel("A")).astype(np.float32) / 255.0 ) mask_t = 1.0 - torch.from_numpy(mask_np) else: mask_t = torch.zeros((64, 64), dtype=torch.float32, device="cpu") output_images.append(image_t) output_masks.append(mask_t.unsqueeze(0)) if not output_images: raise RuntimeError(f"No frames could be decoded from image: {image}") if len(output_images) > 1 and getattr(img, "format", None) not in excluded_formats: output_image = torch.cat(output_images, dim=0) output_mask = torch.cat(output_masks, dim=0) else: output_image = output_images[0] output_mask = output_masks[0] return (output_image, output_mask) @classmethod def IS_CHANGED(cls, image): folder_paths, _ = _get_comfy_modules() image_path = folder_paths.get_annotated_filepath(image) m = hashlib.sha256() with open(image_path, "rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): m.update(chunk) return m.digest().hex() @classmethod def VALIDATE_INPUTS(cls, image): folder_paths, _ = _get_comfy_modules() if not folder_paths.exists_annotated_filepath(image): return "Invalid image file: {}".format(image) return True NODE_CLASS_MAPPINGS = { "LoadImagePlusHEIC": LoadImagePlusHEIC, } NODE_DISPLAY_NAME_MAPPINGS = { "LoadImagePlusHEIC": "Load Image (HEIC)", }