MushroomFleet-ComfyUI-ZeroCLIP-nodes/nodes_shared.py

"""Shared utility nodes for ZeroClip-nodes."""

import struct
import torch

from .utils import (
    ZEROCLIP_SEED_TYPE,
    zero_conditioning,
    zero_sdxl_conditioning,
)


class ZeroClip_SeedPack:
    """Packs a 4-tuple seed into a passable object for ZeroClip conditioner nodes."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "concept_id": ("INT", {
                    "default": 1000, "min": 0, "max": 65535, "step": 1,
                    "tooltip": "Primary semantic axis (u16, 0-65535)"
                }),
                "style_id": ("INT", {
                    "default": 500, "min": 0, "max": 65535, "step": 1,
                    "tooltip": "Secondary style axis (u16, 0-65535)"
                }),
                "mood_salt": ("INT", {
                    "default": 0, "min": 0, "max": 65535, "step": 1,
                    "tooltip": "Mood channel salt (u16, 0-65535)"
                }),
                "world_seed": ("INT", {
                    "default": 42, "min": 0, "max": 0xFFFFFFFFFFFFFFFF,
                    "step": 1,
                    "tooltip": "Universe identifier (u64)"
                }),
            }
        }

    RETURN_TYPES = (ZEROCLIP_SEED_TYPE,)
    RETURN_NAMES = ("seed",)
    FUNCTION = "pack"
    CATEGORY = "ZeroClip"

    def pack(self, concept_id, style_id, mood_salt, world_seed):
        return ((concept_id, style_id, mood_salt, world_seed),)


class ZeroClip_SeedFromRandom:
    """Derive a ZeroCLIP seed tuple from a single random integer.

    Uses FNV-1a child hashes to split one seed into four independent values,
    enabling ComfyUI's standard seed randomization to drive ZeroCLIP exploration.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "seed": ("INT", {
                    "default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF,
                    "tooltip": "Master seed (use ComfyUI's randomize widget)"
                }),
            }
        }

    RETURN_TYPES = (ZEROCLIP_SEED_TYPE,)
    RETURN_NAMES = ("seed",)
    FUNCTION = "derive"
    CATEGORY = "ZeroClip"

    def derive(self, seed):
        from .core.hash import fnv64, child_hash
        base = fnv64(struct.pack('<Q', seed & 0xFFFFFFFFFFFFFFFF))
        concept_id = child_hash(base, 0) & 0xFFFF
        style_id = child_hash(base, 1000) & 0xFFFF
        mood_salt = child_hash(base, 2000) & 0xFFFF
        world_seed = child_hash(base, 3000)
        return ((concept_id, style_id, mood_salt, world_seed),)


class ZeroClip_EmptyConditioning:
    """Produces zero-vector conditioning for uncond/negative input."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dimension": ("INT", {
                    "default": 768, "min": 64, "max": 4096, "step": 1,
                    "tooltip": "Embedding dimension: 768 (SD1.x), 1024 (SD2.x), 2048 (SDXL seq)"
                }),
                "sdxl": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Enable SDXL format (adds pooled_output to conditioning dict)"
                }),
            },
            "optional": {
                "pooled_dimension": ("INT", {
                    "default": 1280, "min": 64, "max": 4096, "step": 1,
                    "tooltip": "SDXL pooled embedding dimension (typically 1280)"
                }),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION = "generate"
    CATEGORY = "ZeroClip"

    def generate(self, dimension, sdxl, pooled_dimension=1280):
        if sdxl:
            return (zero_sdxl_conditioning(dimension, pooled_dimension),)
        return (zero_conditioning(dimension),)


class ZeroClip_ConditioningBlend:
    """Blend two CONDITIONING inputs via weighted interpolation.

    Performs linear interpolation on the conditioning tensors and
    L2-normalizes the result to stay on the unit sphere. Also blends
    pooled_output if present (for SDXL compatibility).
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning_a": ("CONDITIONING",),
                "conditioning_b": ("CONDITIONING",),
                "weight_b": ("FLOAT", {
                    "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
                    "tooltip": "0.0=100% A, 1.0=100% B, 0.5=equal blend"
                }),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    RETURN_NAMES = ("conditioning",)
    FUNCTION = "blend"
    CATEGORY = "ZeroClip"

    def blend(self, conditioning_a, conditioning_b, weight_b):
        t_a = conditioning_a[0][0]
        t_b = conditioning_b[0][0]
        dict_a = conditioning_a[0][1]
        dict_b = conditioning_b[0][1]

        # Lerp and renormalize
        t_out = t_a * (1.0 - weight_b) + t_b * weight_b
        norm = torch.norm(t_out, dim=-1, keepdim=True).clamp(min=1e-8)
        t_out = t_out / norm

        # Handle pooled_output for SDXL
        out_dict = dict(dict_a)
        if "pooled_output" in dict_a and "pooled_output" in dict_b:
            p_a = dict_a["pooled_output"]
            p_b = dict_b["pooled_output"]
            p_out = p_a * (1.0 - weight_b) + p_b * weight_b
            p_norm = torch.norm(p_out, dim=-1, keepdim=True).clamp(min=1e-8)
            out_dict["pooled_output"] = p_out / p_norm

        return ([(t_out, out_dict)],)


class ZeroClip_ConditioningInfo:
    """Display information about a CONDITIONING tensor.

    Shows dimension, norm, and basic statistics. Useful for debugging
    and understanding what a ZeroCLIP conditioning vector contains.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "conditioning": ("CONDITIONING",),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("info",)
    FUNCTION = "inspect"
    CATEGORY = "ZeroClip"
    OUTPUT_NODE = True

    def inspect(self, conditioning):
        t = conditioning[0][0]
        d = conditioning[0][1]
        B, S, D = t.shape

        vec = t[0, 0]
        norm = float(torch.norm(vec))
        mean_val = float(vec.mean())
        std_val = float(vec.std())
        min_val = float(vec.min())
        max_val = float(vec.max())

        lines = [
            f"Shape: [{B}, {S}, {D}]",
            f"Dtype: {t.dtype}",
            f"Norm: {norm:.6f}",
            f"Mean: {mean_val:.6f}",
            f"Std:  {std_val:.6f}",
            f"Range: [{min_val:.4f}, {max_val:.4f}]",
        ]

        if "pooled_output" in d:
            p = d["pooled_output"]
            lines.append(f"Pooled: shape={list(p.shape)}, "
                         f"norm={float(torch.norm(p)):.6f}")

        info = "\n".join(lines)
        return {"ui": {"text": [info]}, "result": (info,)}