Cthomasdesign-ComfyUI-Krea2-MLX/nodes.py

"""ComfyUI nodes for Krea-2-Turbo on Apple MLX.

Runs the vendored `krea2_engine` MLX pipeline in-process and returns ComfyUI IMAGE tensors — a
self-contained "island": load model → (stack LoRAs) → generate/img2img/edit → IMAGE, then any
downstream ComfyUI node can save/upscale/post-process. It does not expose a torch MODEL, so it
does not plug into ComfyUI's native KSampler/ControlNet.

Nodes (category "Krea2 MLX"):
  Krea2ModelLoader   pick a transformer build in ComfyUI/models/krea2 -> KREA2_PIPE (cached)
  Krea2LoRA          [stack] name strength   -> KREA2_LORASTACK        (chain to stack multiple)
  Krea2Generate      pipe prompt … [stack]   -> IMAGE                  (text-to-image)
  Krea2Img2Img       pipe image prompt …     -> IMAGE                  (denoise-from-source)
  Krea2Edit          pipe image prompt …     -> IMAGE                  (in-context edit, optional
                                                                        Qwen3-VL grounding)
  Krea2Unload        ()                      -> frees the cached model + Metal cache
"""

import gc
import os

import numpy as np
import torch

import folder_paths  # ComfyUI
from comfy.utils import ProgressBar
import comfy.model_management as mm

from .krea2_engine.pipeline import Krea2Pipeline

# custom socket types for passing objects between our nodes
KREA2_PIPE = "KREA2_PIPE"
KREA2_LORASTACK = "KREA2_LORASTACK"

# Transformer builds live in a "krea2" model folder. Register the install default
# (ComfyUI/models/krea2); ComfyUI's extra_model_paths (e.g. a Comfy Desktop shared store) adds
# any others. All lookups go through folder_paths so every registered location is searched — not
# just the install folder.
KREA2_MODELS = os.path.join(folder_paths.models_dir, "krea2")
folder_paths.add_model_folder_path("krea2", KREA2_MODELS)

# one model resident at a time (a 12.9B transformer won't share unified memory with another)
_PIPE_CACHE: dict[str, Krea2Pipeline] = {}


def _precision_for(filename):
    """Infer the quantization recipe from a transformer filename (None if unrecognized)."""
    n = filename.lower()
    if "mixed" in n or "4_8" in n or "4-8" in n:
        return "mixed-4-8"
    if "8bit" in n or "8-bit" in n:
        return "8bit"
    if "bf16" in n or "turbo" in n:
        return "bf16"
    return None


def _krea2_dirs():
    """Every directory ComfyUI searches for krea2 builds (honors extra_model_paths)."""
    try:
        dirs = list(folder_paths.get_folder_paths("krea2"))
    except KeyError:
        dirs = []
    if KREA2_MODELS not in dirs:
        dirs.append(KREA2_MODELS)
    return dirs


def _list_models():
    """Recipe-inferable transformer builds across all krea2 search paths."""
    seen = []
    for d in _krea2_dirs():
        if os.path.isdir(d):
            for f in sorted(os.listdir(d)):
                if f.endswith(".safetensors") and _precision_for(f) and f not in seen:
                    seen.append(f)
    return seen


def _model_path(model_file):
    """Resolve a build filename to a full path across the krea2 search paths."""
    p = folder_paths.get_full_path("krea2", model_file)
    if p:
        return p
    for d in _krea2_dirs():
        cand = os.path.join(d, model_file)
        if os.path.isfile(cand):
            return cand
    return None


def _to_image_tensor(pil_images):
    """List of PIL images -> ComfyUI IMAGE tensor, float32 [B, H, W, C] in 0..1."""
    arrs = [np.asarray(im.convert("RGB"), dtype=np.float32) / 255.0 for im in pil_images]
    return torch.from_numpy(np.stack(arrs, axis=0))


def _from_image_tensor(image):
    """ComfyUI IMAGE tensor [B, H, W, C] in 0..1 -> the first frame as a PIL image."""
    from PIL import Image

    arr = (image[0].cpu().numpy() * 255.0).round().clip(0, 255).astype(np.uint8)
    return Image.fromarray(arr)


def _get_pipe(model_file):
    if model_file not in _PIPE_CACHE:
        _PIPE_CACHE.clear()          # free the previous build first
        gc.collect()
        import mlx.core as mx
        mx.clear_cache()
        _PIPE_CACHE[model_file] = Krea2Pipeline(_model_path(model_file),
                                                precision=_precision_for(model_file),
                                                base_dir=os.environ.get("KREA2_BASE_DIR"))
    return _PIPE_CACHE[model_file]


# --- shared node plumbing -----------------------------------------------------------------
# The three generation nodes differ only in which pipeline method they call; everything around
# that call (LoRA application, progress bar + Cancel, the license-required safety filter, and
# the PIL->IMAGE conversion) is identical and lives here.

def _run(krea2_pipe, steps, safety_filter, lora_stack, generate_fn):
    """Apply LoRAs, run `generate_fn(step_callback)` with progress/Cancel wired to ComfyUI,
    apply the NSFW filter, and convert to an IMAGE tensor tuple."""
    krea2_pipe.set_loras(lora_stack or [])
    pbar = ProgressBar(steps)

    def cb(step, total):
        mm.throw_exception_if_processing_interrupted()  # honor ComfyUI's Cancel button
        pbar.update_absolute(step, total)

    imgs = generate_fn(cb)
    # NSFW content filter (on by default). The Krea 2 Community License (§4.2) requires
    # reasonable content-filtering in deployments; flagged images are redacted.
    if safety_filter:
        from .krea2_engine import safety
        imgs, _ = safety.apply(imgs, enabled=True)
    return (_to_image_tensor(imgs),)


# Widget factories return fresh spec dicts (never shared/aliased between nodes). Assembly order
# in each node's INPUT_TYPES is load-bearing: ComfyUI stores workflow widget values positionally.

def _prompt_widget(default):
    return ("STRING", {"multiline": True, "default": default,
                       "tooltip": "What to generate. Turbo uses no negative prompt (guidance 0)."})


def _size_widgets():
    tip = "Must be a multiple of 16. For image inputs, match the source aspect ratio."
    return {
        "width": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 16, "tooltip": tip}),
        "height": ("INT", {"default": 1024, "min": 256, "max": 2048, "step": 16, "tooltip": tip}),
    }


def _sampling_widgets():
    return {
        "steps": ("INT", {"default": 8, "min": 1, "max": 50,
                          "tooltip": "Turbo is distilled for 8 steps; more rarely helps."}),
        "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF,
                         "tooltip": "Same seed + settings reproduces the same image."}),
        "num_images": ("INT", {"default": 1, "min": 1, "max": 8,
                               "tooltip": "Batch size. One 12.9B model stays resident; large "
                                          "batches at 1024² get memory-tight on 32 GB."}),
    }


def _safety_widget():
    return ("BOOLEAN", {"default": True, "label_on": "on", "label_off": "off",
                        "tooltip": "NSFW filter (redacts flagged images). The Krea 2 Community "
                                   "License requires content filtering in deployments."})


class Krea2ModelLoader:
    DESCRIPTION = ("Loads a Krea-2-Turbo MLX transformer build from ComfyUI/models/krea2 "
                   "(plus the Qwen3-VL text encoder and VAE) and keeps it resident. "
                   "Precision is inferred from the filename (mixed_4_8 / 8bit / bf16).")

    @classmethod
    def INPUT_TYPES(cls):
        models = _list_models() or ["(put transformer_*.safetensors in ComfyUI/models/krea2)"]
        return {"required": {"model": (models, {"tooltip": "Transformer build to load."})}}

    RETURN_TYPES = (KREA2_PIPE,)
    RETURN_NAMES = ("krea2_pipe",)
    FUNCTION = "load"
    CATEGORY = "Krea2 MLX"

    def load(self, model):
        if not _list_models():
            raise RuntimeError(
                "No Krea-2 transformer found in any 'krea2' model folder "
                f"({', '.join(_krea2_dirs())}). Download a build, e.g. "
                "transformer_mixed_4_8.safetensors from "
                "huggingface.co/avlp12/Krea-2-Turbo-Alis-MLX-mixed-4-8")
        return (_get_pipe(model),)


class Krea2LoRA:
    DESCRIPTION = ("Adds a LoRA from ComfyUI/models/loras to the stack. Chain several Krea2LoRA "
                   "nodes to stack adapters; both Krea/diffusers and ai-toolkit naming work "
                   "(incl. the krea2_edit identity LoRA).")

    @classmethod
    def INPUT_TYPES(cls):
        loras = folder_paths.get_filename_list("loras")
        return {
            "required": {
                "lora_name": (loras if loras else ["None"],
                              {"tooltip": "LoRA .safetensors from ComfyUI/models/loras."}),
                "strength": ("FLOAT", {"default": 0.9, "min": -2.0, "max": 2.0, "step": 0.05,
                                       "tooltip": "Adapter strength (0 skips it; the identity "
                                                  "edit LoRA is trained for 1.0)."}),
            },
            "optional": {"lora_stack": (KREA2_LORASTACK,
                                        {"tooltip": "Chain from a previous Krea2LoRA to stack."})},
        }

    RETURN_TYPES = (KREA2_LORASTACK,)
    RETURN_NAMES = ("lora_stack",)
    FUNCTION = "append"
    CATEGORY = "Krea2 MLX"

    def append(self, lora_name, strength, lora_stack=None):
        stack = list(lora_stack) if lora_stack else []
        if lora_name and lora_name != "None" and strength != 0:
            path = folder_paths.get_full_path("loras", lora_name)
            if path:
                stack.append((path, float(strength)))
        return (stack,)


class Krea2Generate:
    DESCRIPTION = ("Text-to-image with Krea-2-Turbo on MLX. 8-step flow matching, no negative "
                   "prompt/CFG. Progress bar and Cancel supported.")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "krea2_pipe": (KREA2_PIPE,),
                "prompt": _prompt_widget("a fox in the snow"),
                **_size_widgets(),
                **_sampling_widgets(),
                "safety_filter": _safety_widget(),
            },
            "optional": {"lora_stack": (KREA2_LORASTACK,)},
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "Krea2 MLX"

    def generate(self, krea2_pipe, prompt, width, height, steps, seed, num_images,
                 safety_filter=True, lora_stack=None):
        return _run(krea2_pipe, steps, safety_filter, lora_stack,
                    lambda cb: krea2_pipe.generate(
                        prompt, width=width, height=height, steps=steps, seed=seed,
                        num_images=num_images, step_callback=cb))


class Krea2Img2Img:
    DESCRIPTION = ("Image-to-image: VAE-encodes the source and denoises from it. denoise 0→1 "
                   "goes from near-copy to full reinterpretation (1.0 = plain text-to-image). "
                   "Reinterprets the whole frame — for targeted edits use Krea2 Edit.")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "krea2_pipe": (KREA2_PIPE,),
                "image": ("IMAGE", {"tooltip": "Source image (the starting point)."}),
                "prompt": _prompt_widget("a fox in the snow"),
                "denoise": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.0, "step": 0.05,
                                      "tooltip": "How far to travel from the source: ~0.3 polish, "
                                                 "~0.5-0.7 restyle, 1.0 ignores the source."}),
                **_size_widgets(),
                **_sampling_widgets(),
                "safety_filter": _safety_widget(),
            },
            "optional": {"lora_stack": (KREA2_LORASTACK,)},
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "Krea2 MLX"

    def generate(self, krea2_pipe, image, prompt, denoise, width, height, steps, seed, num_images,
                 safety_filter=True, lora_stack=None):
        return _run(krea2_pipe, steps, safety_filter, lora_stack,
                    lambda cb: krea2_pipe.generate_img2img(
                        prompt, _from_image_tensor(image), denoise=denoise,
                        width=width, height=height, steps=steps, seed=seed,
                        num_images=num_images, step_callback=cb))


class Krea2Edit:
    DESCRIPTION = ("In-context edit: the source rides the sequence as a clean reference frame "
                   "while a fresh target is generated — identity-preserving edits, not img2img. "
                   "Pair with the krea2_edit identity LoRA (via Krea2 LoRA) and keep the output "
                   "aspect ratio matched to the source. grounding_px > 0 additionally reads the "
                   "instruction while looking at the source (Qwen3-VL).")

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "krea2_pipe": (KREA2_PIPE,),
                "image": ("IMAGE", {"tooltip": "Source/reference image (scene, for two-ref edits)."}),
                "prompt": _prompt_widget("make the shirt blue"),
                **_size_widgets(),
                **_sampling_widgets(),
                "grounding_px": ("INT", {"default": 768, "min": 0, "max": 1536, "step": 64,
                                          "tooltip": "0 = plain text; >0 reads the instruction while "
                                                     "looking at the source via Qwen3-VL (longest side "
                                                     "cap). Lower = stronger edit adherence, higher = "
                                                     "stronger identity; 768 is balanced."}),
                "safety_filter": _safety_widget(),
            },
            "optional": {
                "image_b": ("IMAGE", {"tooltip": "2nd reference (subject) for multi-ref edit LoRAs; "
                                                 "training order is scene first, subject second."}),
                "lora_stack": (KREA2_LORASTACK,),
            },
        }

    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "generate"
    CATEGORY = "Krea2 MLX"

    def generate(self, krea2_pipe, image, prompt, width, height, steps, seed, num_images,
                 grounding_px=768, safety_filter=True, image_b=None, lora_stack=None):
        return _run(krea2_pipe, steps, safety_filter, lora_stack,
                    lambda cb: krea2_pipe.generate_edit(
                        prompt, _from_image_tensor(image),
                        image_b=_from_image_tensor(image_b) if image_b is not None else None,
                        grounding_px=grounding_px, width=width, height=height, steps=steps,
                        seed=seed, num_images=num_images, step_callback=cb))


class Krea2Unload:
    DESCRIPTION = ("Frees the cached Krea-2 model and clears the MLX Metal cache — reclaim "
                   "unified memory before loading other large models.")

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

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("status",)
    FUNCTION = "unload"
    CATEGORY = "Krea2 MLX"
    OUTPUT_NODE = True

    def unload(self):
        _PIPE_CACHE.clear()
        gc.collect()
        import mlx.core as mx
        mx.clear_cache()
        return ("Krea2 model unloaded; Metal cache cleared.",)


NODE_CLASS_MAPPINGS = {
    "Krea2ModelLoader": Krea2ModelLoader,
    "Krea2LoRA": Krea2LoRA,
    "Krea2Generate": Krea2Generate,
    "Krea2Img2Img": Krea2Img2Img,
    "Krea2Edit": Krea2Edit,
    "Krea2Unload": Krea2Unload,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "Krea2ModelLoader": "Krea2 Model Loader (MLX)",
    "Krea2LoRA": "Krea2 LoRA (MLX)",
    "Krea2Generate": "Krea2 Generate (MLX)",
    "Krea2Img2Img": "Krea2 Img2Img (MLX)",
    "Krea2Edit": "Krea2 Edit (MLX)",
    "Krea2Unload": "Krea2 Unload (MLX)",
}