salia_hf_to_batch.py

import os
import re
import zipfile
import shutil
import tempfile
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError

import numpy as np
import torch
from PIL import Image

try:
    import folder_paths  # ComfyUI helper for temp dirs
except Exception:
    folder_paths = None


def _get_cache_dir() -> str:
    base_dir = None
    if folder_paths is not None:
        try:
            base_dir = folder_paths.get_temp_directory()
        except Exception:
            base_dir = None

    if not base_dir:
        base_dir = tempfile.gettempdir()

    cache_dir = os.path.join(base_dir, "hf_zip_cache")
    os.makedirs(cache_dir, exist_ok=True)
    return cache_dir


def _download_file(url: str, dest_path: str, timeout_sec: int = 60) -> None:
    req = Request(url, headers={"User-Agent": "ComfyUI-HFZipLoader/1.0"})
    try:
        with urlopen(req, timeout=timeout_sec) as resp, open(dest_path, "wb") as out_f:
            shutil.copyfileobj(resp, out_f)
    except HTTPError as e:
        raise ValueError(f"HTTP error while downloading: {url} (status={e.code})") from e
    except URLError as e:
        raise ValueError(f"Network error while downloading: {url} ({e.reason})") from e
    except Exception as e:
        raise ValueError(f"Unexpected error while downloading: {url} ({e})") from e


def _pil_to_tensor_rgb(pil_img: Image.Image) -> torch.Tensor:
    """
    Convert PIL image to ComfyUI IMAGE tensor: [H,W,3] float32 in [0..1].
    """
    if pil_img.mode != "RGB":
        pil_img = pil_img.convert("RGB")

    arr = np.asarray(pil_img, dtype=np.float32) / 255.0  # HWC
    return torch.from_numpy(arr)  # torch float32 HWC


class _ImageSizeMismatchError(ValueError):
    """Raised when images in the zip do not share the same dimensions."""


def _alphanum_key(s: str):
    """
    Natural/alphanumeric sort key for filenames/paths.
    Example: img_2.png comes before img_10.png.

    Sorts by the full zip member name (including folders), case-insensitive.
    """
    s = (s or "").replace("\\", "/")
    parts = re.split(r"(\d+)", s)

    # Build a key composed of tagged tokens so Python never compares int vs str directly.
    key = []
    for p in parts:
        if p.isdigit():
            key.append((0, int(p)))
        else:
            key.append((1, p.lower()))
    return key


def _load_images_from_zip(zip_path: str) -> torch.Tensor:
    """
    Forgiving loader:
      - Accepts all filenames (any depth) in a zip
      - Sorts members in alphanumeric (natural) order
      - Tries to open each file as an image; skips files that PIL cannot read
      - Enforces that all loaded images share the same dimensions

    Returns:
      [B,H,W,3] float32 in [0..1]
    """
    images = []
    shapes = None
    skipped = []

    with zipfile.ZipFile(zip_path, "r") as zf:
        members = [name for name in zf.namelist() if name and not name.endswith("/")]

        if not members:
            raise ValueError("ZIP is empty (no files found).")

        members.sort(key=_alphanum_key)

        for member_name in members:
            try:
                with zf.open(member_name) as fp:
                    with Image.open(fp) as im:
                        # Ensure image data is fully read while the zip file handle is still open
                        im.load()
                        t = _pil_to_tensor_rgb(im)  # HWC, RGB, float32

                if shapes is None:
                    shapes = tuple(t.shape)
                else:
                    if tuple(t.shape) != shapes:
                        raise _ImageSizeMismatchError(
                            f"Image size mismatch in ZIP. Expected {shapes}, got {tuple(t.shape)} "
                            f"for {member_name}. All images must share the same dimensions."
                        )

                images.append(t)

            except _ImageSizeMismatchError:
                # This is a hard error: the batch cannot be formed consistently.
                raise
            except Exception:
                # Forgiving: ignore non-images, unreadable files, etc.
                skipped.append(member_name)
                continue

    if not images:
        raise ValueError(
            "No loadable images found in ZIP. Ensure the archive contains valid image files "
            "(png/jpg/webp/etc.)."
        )

    if skipped:
        print(f"[HFLoadZipImageBatch] Skipped {len(skipped)} non-image/unreadable file(s) in ZIP.")

    return torch.stack(images, dim=0)  # BHWC


class HF_to_Batch:
    """
    Download public ZIP from Hugging Face resolve URL and output IMAGE batch.

    URL format:
      https://huggingface.co/{owner}/{repo}/resolve/{revision}/{index}.zip

    Example:
      owner=saliacoel, repo=pov_fs, revision=main, index=0
      -> https://huggingface.co/saliacoel/pov_fs/resolve/main/0.zip
    """

    CATEGORY = "HuggingFace"
    RETURN_TYPES = ("IMAGE", "STRING", "INT", "STRING")
    RETURN_NAMES = ("images", "source_url", "count", "local_zip_path")
    FUNCTION = "load"

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "repo": ("STRING", {"default": "pov_fs", "multiline": False}),
                "index": ("INT", {"default": 0, "min": 0, "max": 1000000, "step": 1}),
            },
            "optional": {
                "owner": ("STRING", {"default": "saliacoel", "multiline": False}),
                "revision": ("STRING", {"default": "main", "multiline": False}),
                "force_redownload": ("BOOLEAN", {"default": False}),
            },
        }

    def load(
        self,
        repo: str,
        index: int,
        owner: str = "saliacoel",
        revision: str = "main",
        force_redownload: bool = False,
    ):
        repo = (repo or "").strip()
        owner = (owner or "").strip()
        revision = (revision or "").strip()

        if not repo:
            raise ValueError("repo must be a non-empty string (e.g., 'pov_fs' or 'car').")
        if not owner:
            raise ValueError("owner must be a non-empty string (e.g., 'saliacoel').")
        if index is None or int(index) < 0:
            raise ValueError("index must be an integer >= 0.")

        index = int(index)

        source_url = f"https://huggingface.co/{owner}/{repo}/resolve/{revision}/{index}.zip"

        cache_dir = _get_cache_dir()
        local_zip_path = os.path.join(cache_dir, f"{owner}__{repo}__{revision}__{index}.zip")

        if (
            force_redownload
            or (not os.path.exists(local_zip_path))
            or (os.path.getsize(local_zip_path) == 0)
        ):
            _download_file(source_url, local_zip_path)

        images = _load_images_from_zip(local_zip_path)
        count = int(images.shape[0])

        print(f"[HFLoadZipImageBatch] Loaded {count} image(s) from {source_url}")
        return (images, source_url, count, local_zip_path)


NODE_CLASS_MAPPINGS = {
    "HF_to_Batch": HF_to_Batch,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "HF_to_Batch": "HF_to_Batch",
}