nodes.py

import hashlib
import io
import urllib.request
import urllib.error

import numpy as np
import torch
from PIL import Image, ImageOps


class LoadImageFromHttpURLNorm:
    """Loads an image from an HTTP/HTTPS URL with automatic EXIF orientation correction."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image_url": ("STRING", {"default": "", "multiline": False}),
            },
        }

    RETURN_TYPES = ("IMAGE", "MASK", "INT", "INT")
    RETURN_NAMES = ("image", "mask", "width", "height")
    FUNCTION = "load_image"
    CATEGORY = "image"

    def load_image(self, image_url: str):
        if not image_url or not image_url.strip():
            raise ValueError("image_url is empty")

        image_url = image_url.strip()

        req = urllib.request.Request(image_url, headers={"User-Agent": "ComfyUI"})
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = resp.read()

        img = Image.open(io.BytesIO(data))

        # Apply EXIF orientation — handles all 8 orientation tag values
        # (rotation, mirroring) and strips the tag so downstream code
        # doesn't double-rotate.
        img = ImageOps.exif_transpose(img)

        has_alpha = img.mode == "RGBA" or "A" in img.getbands()

        img = img.convert("RGBA") if has_alpha else img.convert("RGB")

        # IMAGE tensor: (1, H, W, C) float32 in [0, 1]
        image_np = np.array(img).astype(np.float32) / 255.0

        if has_alpha:
            rgb = image_np[:, :, :3]
            alpha = image_np[:, :, 3]
        else:
            rgb = image_np
            alpha = np.ones((image_np.shape[0], image_np.shape[1]), dtype=np.float32)

        image_tensor = torch.from_numpy(rgb).unsqueeze(0)  # (1, H, W, 3)
        mask_tensor = 1.0 - torch.from_numpy(alpha).unsqueeze(0)  # (1, H, W)

        return (image_tensor, mask_tensor, image_tensor.shape[2], image_tensor.shape[1])

    @classmethod
    def IS_CHANGED(cls, image_url: str):
        return hashlib.sha256(image_url.strip().encode()).hexdigest()