gokayfem-ComfyUI-Texture-Simple/__init__.py

"""Interactive material and texture-map preview node for ComfyUI."""

from __future__ import annotations

import os
from typing import Any

import folder_paths
import numpy as np
from PIL import Image

try:
    from .texture_nodes import (
        NODE_CLASS_MAPPINGS as TOOL_NODE_CLASS_MAPPINGS,
        NODE_DISPLAY_NAME_MAPPINGS as TOOL_NODE_DISPLAY_NAME_MAPPINGS,
    )
except (ImportError, ModuleNotFoundError):  # Standalone import used by pytest.
    from texture_nodes import (  # type: ignore[no-redef]
        NODE_CLASS_MAPPINGS as TOOL_NODE_CLASS_MAPPINGS,
        NODE_DISPLAY_NAME_MAPPINGS as TOOL_NODE_DISPLAY_NAME_MAPPINGS,
    )


MAP_MODES = {
    "color": "RGB",
    "displacement": "L",
    "normal": "RGB",
    "ao": "L",
    "metalness": "L",
    "roughness": "L",
    "alpha": "L",
}


def _as_pil(image: Any, mode: str, *, bit_depth: int = 8) -> Image.Image:
    array = image.detach().cpu().float().numpy()
    array = np.nan_to_num(array, nan=0.0, posinf=1.0, neginf=0.0)
    array = np.clip(array, 0.0, 1.0)

    if array.ndim == 2:
        source_mode = "L"
    elif array.ndim == 3 and array.shape[-1] == 1:
        array = array[..., 0]
        source_mode = "L"
    elif array.ndim == 3 and array.shape[-1] >= 3:
        array = array[..., :3]
        source_mode = "RGB"
    else:
        raise ValueError(f"Expected an HxW, HxWx1, or HxWx3+ image, got {array.shape}.")

    if bit_depth == 16 and mode == "L":
        if source_mode == "RGB":
            array = array[..., 0] * 0.2126 + array[..., 1] * 0.7152 + array[..., 2] * 0.0722
        converted = Image.fromarray((array * 65535.0).round().astype(np.uint16))
        return converted
    converted = Image.fromarray((array * 255.0).round().astype(np.uint8), mode=source_mode)
    return converted.convert(mode)


def _save_map(
    image: Image.Image,
    *,
    map_type: str,
    batch_number: int,
) -> dict[str, str]:
    output_dir = folder_paths.get_temp_directory()
    full_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
        "texture_viewer",
        output_dir,
        image.width,
        image.height,
    )
    filename = filename.replace("%batch_num%", str(batch_number))
    image_name = f"{filename}_{counter:05}_{map_type}.png"
    image.save(os.path.join(full_folder, image_name), compress_level=1)
    return {"filename": image_name, "subfolder": subfolder, "type": "temp"}


class TextureViewer:
    """Preview PBR texture maps on built-in or browser-loaded 3D meshes."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {},
            "optional": {
                "color_map": ("IMAGE",),
                "displacement_map": ("IMAGE",),
                "normal_map": ("IMAGE",),
                "ao_map": ("IMAGE",),
                "metalness_map": ("IMAGE",),
                "roughness_map": ("IMAGE",),
                "alpha_map": ("IMAGE",),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("passthrough",)
    OUTPUT_NODE = True
    FUNCTION = "process_images"
    CATEGORY = "visualization/3D"
    DESCRIPTION = (
        "Interactively previews PBR texture maps on selectable primitives or a "
        "browser-loaded GLB/OBJ mesh. Supports batches, screenshots, and mesh export."
    )

    def process_images(
        self,
        color_map=None,
        displacement_map=None,
        normal_map=None,
        ao_map=None,
        metalness_map=None,
        roughness_map=None,
        alpha_map=None,
    ):
        provided = {
            "color": color_map,
            "displacement": displacement_map,
            "normal": normal_map,
            "ao": ao_map,
            "metalness": metalness_map,
            "roughness": roughness_map,
            "alpha": alpha_map,
        }
        counts = [len(batch) for batch in provided.values() if batch is not None]
        target_count = max(counts, default=0)

        for map_type, batch in provided.items():
            if batch is not None and len(batch) not in (1, target_count):
                raise ValueError(
                    f"{map_type}_map has {len(batch)} images, but the largest input "
                    f"batch has {target_count}. Inputs must match or contain one image."
                )

        saved: dict[str, list[dict[str, str]]] = {
            map_type: [] for map_type in provided
        }
        for map_type, batch in provided.items():
            if batch is None:
                continue
            mode = MAP_MODES[map_type]
            saved[map_type] = [
                _save_map(
                    _as_pil(
                        image,
                        mode,
                        bit_depth=16 if map_type == "displacement" else 8,
                    ),
                    map_type=map_type,
                    batch_number=index,
                )
                for index, image in enumerate(batch)
            ]

        passthrough = next((batch for batch in provided.values() if batch is not None), None)
        if passthrough is None:
            import torch

            passthrough = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
        return {"ui": saved, "result": (passthrough,)}


NODE_CLASS_MAPPINGS = {"TextureViewer": TextureViewer, **TOOL_NODE_CLASS_MAPPINGS}
NODE_DISPLAY_NAME_MAPPINGS = {
    "TextureViewer": "Texture Viewer Pro",
    **TOOL_NODE_DISPLAY_NAME_MAPPINGS,
}
WEB_DIRECTORY = "./web"

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]