sd_cpp_node.py

"""CLI 黑盒节点: 通过 sd-cli.exe (stable-diffusion.cpp) 生成图像。

设计:
- 所有组件直接在节点上选: diffusion_model + 可选 llm/clip_l/clip_g/t5xxl/vae
   (不再有 profiles.json 配置, 换模型/文本编码器/VAE 不用改文件)
- sd.cpp 从 GGUF 头自动识别架构, 无需 --type
- 单文件模型 (sd1.5/sdxl 含内嵌 clip+vae) 不选任何配套组件 -> 自动用 -m 加载
- 解码/分块 VAE 全部在 sd.cpp 内部完成 (自带 vae-tiling)
"""
import glob
import os
import random
import shlex
import subprocess
import uuid

import numpy as np
import torch
from PIL import Image

import folder_paths

NODE_DIR = os.path.dirname(os.path.abspath(__file__))
SD_CLI = os.path.join(NODE_DIR, "bin", "sd-cli.exe")

# sd.cpp master 新命名 (str_to_sample_method 精确匹配, 无旧名别名):
# dpmpp_2m -> dpm++2m, dpmpp_2m_v2 -> dpm++2mv2, dpmpp_2s_a -> dpm++2s_a
SAMPLERS = ["euler", "euler_a", "heun", "dpm2", "dpm++2s_a", "dpm++2m",
            "dpm++2mv2", "ipndm", "ipndm_v", "lcm", "ddim_trailing", "tcd",
            "res_multistep", "res_2s", "er_sde", "euler_cfg_pp",
            "euler_a_cfg_pp", "euler_ge", "dpm++2m_sde", "dpm++2m_sde_bt", "lms"]
# --scheduler 的 sigma 调度器 (str_to_scheduler 精确匹配):
# discrete=ComfyUI normal; 不指定时 sd.cpp 按模型自动选 (SD1.5/SDXL->discrete,
# LCM/TCD->lcm, Flux->flux...)。"(auto)" = 不传参数用默认。
SCHEDULERS = ["(auto)", "discrete", "karras", "exponential", "ays", "gits",
              "sgm_uniform", "simple", "smoothstep", "kl_optimal", "lcm",
              "bong_tangent", "ltx2", "logit_normal", "flux2", "flux", "beta"]
MODEL_EXTS = (".gguf", ".safetensors", ".ckpt", ".pth", ".pt")
TE_EXTS = (".gguf", ".safetensors", ".bin", ".pt")
LORA_EXTS = (".safetensors", ".ckpt", ".pt")

NONE = "(none)"

CREATE_NO_WINDOW = 0x08000000 if os.name == "nt" else 0


def _list_folder(folder_names, exts):
    """跨多个 ComfyUI 目录类型收集文件, 返回 [filename, ...]

    注意: 不用 folder_paths.get_filename_list —— 它按原生支持扩展名过滤,
    会滤掉 .gguf (ComfyUI 原生只认 pt/ckpt/safetensors 等)。
    """
    out = set()
    for fn in folder_names:
        for d in folder_paths.get_folder_paths(fn):
            if not os.path.isdir(d):
                continue
            for name in os.listdir(d):
                full = os.path.join(d, name)
                if name.lower().endswith(exts) and os.path.isfile(full):
                    out.add(name)
    return sorted(out)


def _resolve(folder_names, filename):
    """按目录类型顺序解析文件名 -> 绝对路径; 找不到返回 None"""
    for fn in folder_names:
        try:
            p = folder_paths.get_full_path(fn, filename)
        except Exception:
            p = None
        if p and os.path.exists(p):
            return p
    return None


def _combo(items, default):
    """combo 下拉: 指定默认值; 若默认值不在列表(如用户没有该模型文件)则退回列表第一个"""
    return (items, {"default": default}) if default in items else (items,)


class SDCppGenerator:
    """sd.cpp 黑盒生成节点: prompt -> 图 (Vulkan 后端)"""

    @classmethod
    def INPUT_TYPES(cls):
        models = _list_folder(("checkpoints", "diffusion_models", "unet"), MODEL_EXTS)
        loras = [NONE] + _list_folder(("loras",), LORA_EXTS)
        te_opts = [NONE] + _list_folder(("text_encoders", "clip"), TE_EXTS)
        vae_opts = [NONE] + _list_folder(("vae",), (".safetensors", ".bin", ".pt"))
        return {
            "required": {
                "diffusion_model": _combo([NONE] + models, NONE),
                "prompt": ("STRING", {"multiline": True}),
                "negative_prompt": ("STRING", {"multiline": True}),
                "seed": ("INT", {"default": -1, "min": -1, "max": 0x7FFFFFFF}),
                "steps": ("INT", {"default": 8, "min": 1, "max": 200}),
                "cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 30.0, "step": 0.1}),
                "sampling_method": (SAMPLERS,),
                "scheduler": (SCHEDULERS,),
                "width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
                "height": ("INT", {"default": 1024, "min": 64, "max": 4096, "step": 8}),
                "batch_count": ("INT", {"default": 1, "min": 1, "max": 8}),
                "vae_tiling": ("BOOLEAN", {"default": True}),
                "offload_to_cpu": ("BOOLEAN", {"default": True}),
                "diffusion_fa": ("BOOLEAN", {"default": True}),
                "clip_on_cpu": ("BOOLEAN", {"default": True}),
                "vae_on_cpu": ("BOOLEAN", {"default": False}),
                "threads": ("INT", {"default": 0, "min": 0, "max": 128}),
            },
            "optional": {
                # 配套组件: 全部留空 = 单文件模型 (用 -m 加载, 内嵌 clip/vae)
                "llm": (te_opts,),
                "clip_l": (te_opts,),
                "clip_g": (te_opts,),
                "t5xxl": (te_opts,),
                "vae": (vae_opts,),
                "lora": (loras,),
                "lora_strength": ("FLOAT", {"default": 1.0, "min": -2.0, "max": 2.0, "step": 0.05}),
                # --- 图生图 (v1.1): 只在末尾追加, 不动上面已有 widget 顺序 ---
                # init_image/ref_image 是 IMAGE socket, 不占 widgets_values 位,
                # 旧工作流(文生图)完全不受影响
                # v1.2: 顺序改为 lora_strength → ref_preset → denoise_strength → extra_args,
                #   图生图参数聚在一起, 读工作流时先看完两条路线再落到 extra_args
                # v1.2: ref_preset 去掉 "(auto)" — auto 实测对 krea2 默认套
                #   krea2_ostris_edit, 配 identity-edit LoRA 会出废图, 对新手是坑;
                #   默认 default(=不强制指定, 走 sd-cli 通用逻辑), 用编辑 LoRA 时手动选
                "init_image": ("IMAGE",),
                "ref_image": ("IMAGE",),
                "ref_preset": (["default", "krea2_edit", "krea2_ostris_edit", "flux_kontext",
                                "qwen", "z_image_omni", "cosmos_reference"],),
                "denoise_strength": ("FLOAT", {"default": 0.4, "min": 0.0, "max": 1.0, "step": 0.01}),
                "extra_args": ("STRING", {"multiline": True}),
            },
        }

    RETURN_TYPES = ("IMAGE", "INT")
    RETURN_NAMES = ("image", "seed")
    FUNCTION = "generate"
    CATEGORY = "sd.cpp"
    DESCRIPTION = ("调用 sd-cli.exe (stable-diffusion.cpp, Vulkan 后端) 文生图/图生图。"
                   "diffusion_model 必选, llm/clip_l/clip_g/t5xxl/vae 按需选 "
                   "(全留空=单文件 sd1.5/sdxl 自动用内嵌 clip+vae)。\n"
                   "图生图: init_image=重绘 (配 denoise_strength); "
                   "ref_image=参考图编辑 (配 ref_preset, Krea-2 需编辑 LoRA)。"
                   "详细说明见工作流 Note 节点。")

    def generate(self, diffusion_model, prompt, negative_prompt, seed, steps,
                 cfg_scale, sampling_method, scheduler, width, height, batch_count,
                 vae_tiling, offload_to_cpu, diffusion_fa, clip_on_cpu, vae_on_cpu,
                 threads, llm=NONE, clip_l=NONE, clip_g=NONE, t5xxl=NONE, vae=NONE,
                 lora=NONE, lora_strength=1.0,
                 ref_preset="default", denoise_strength=0.4, extra_args="",
                 init_image=None, ref_image=None):
        if not os.path.exists(SD_CLI):
            raise RuntimeError(
                f"找不到 sd-cli.exe: {SD_CLI}\n"
                "请从 https://github.com/leejet/stable-diffusion.cpp/releases 下载 "
                "sd-master-*-bin-win-vulkan-x64.zip 解压到本节点目录的 bin/ 下")

        if not diffusion_model or diffusion_model == NONE:
            raise RuntimeError("请先在 diffusion_model 下拉里选择扩散模型\n"
                               "(sd1.5/sdxl 单文件放 models/checkpoints/, GGUF 放 models/diffusion_models/)")
        model_path = _resolve(("checkpoints", "diffusion_models", "unet"), diffusion_model)
        if not model_path:
            raise RuntimeError(f"找不到扩散模型: {diffusion_model}")

        # --- 配套组件: 选了什么传什么, 全空 = 单文件 ---
        te_map = {"llm": llm, "clip_l": clip_l, "clip_g": clip_g, "t5xxl": t5xxl}
        te_paths = {}
        for key, name in te_map.items():
            if name and name != NONE:
                p = _resolve(("text_encoders", "clip"), name)
                if not p:
                    raise RuntimeError(f"找不到 {key} 文件: {name}\n"
                                       "放到 models/text_encoders/ 或 models/clip/")
                te_paths[key] = p
        vae_path = None
        if vae and vae != NONE:
            vae_path = _resolve(("vae",), vae)
            if not vae_path:
                raise RuntimeError(f"找不到 vae 文件: {vae}\n放到 models/vae/")

        if seed is None or seed < 0:
            seed = random.randint(0, 0x7FFFFFFF)

        # --- 输出目录 ---
        out_dir = os.path.join(folder_paths.get_temp_directory(), "sdcpp", uuid.uuid4().hex[:8])
        os.makedirs(out_dir, exist_ok=True)
        out_path = os.path.join(out_dir, "out_%03d.png" if batch_count > 1 else "out.png")

        # 单文件模式(完整模型, clip+vae 内嵌): sd1.5/sdxl 用 -m 加载,
        # 键名原样进 name_conversion(conditioner.embedders.* -> cond_stage_model.*)。
        # --diffusion-model 会给所有键强加 model.diffusion_model. 前缀,
        # 破坏单文件里的 CLIP/VAE 键名 -> metadata validation failed。
        single_file = not te_paths
        cmd = [SD_CLI, "-m" if single_file else "--diffusion-model", model_path,
               "--output", out_path,
               "--prompt", prompt, "--steps", str(steps), "--cfg-scale", str(cfg_scale),
               "--sampling-method", sampling_method, "--width", str(width),
               "--height", str(height), "--batch-count", str(batch_count),
               "--seed", str(seed)]
        if scheduler and scheduler != "(auto)":
            cmd += ["--scheduler", scheduler]
        if negative_prompt.strip():
            cmd += ["--negative-prompt", negative_prompt]
        for key, p in te_paths.items():
            cmd += [f"--{key}", p]
        if vae_path:
            cmd += ["--vae", vae_path]
        if vae_tiling:
            cmd += ["--vae-tiling"]
        if offload_to_cpu:
            cmd += ["--offload-to-cpu"]
        if diffusion_fa:
            cmd += ["--diffusion-fa"]
        if clip_on_cpu:
            cmd += ["--clip-on-cpu"]
        if vae_on_cpu:
            cmd += ["--vae-on-cpu"]
        if threads > 0:
            cmd += ["--threads", str(threads)]

        if lora != NONE:
            lora_dir = _resolve(("loras",), lora)
            lora_dir = os.path.dirname(lora_dir) if lora_dir else folder_paths.get_folder_paths("loras")[0]
            stem = os.path.splitext(os.path.basename(lora))[0]
            if f"<lora:{stem}:" not in prompt:
                prompt_eff = prompt + f"<lora:{stem}:{lora_strength}>"
            else:
                prompt_eff = prompt
            cmd += ["--lora-model-dir", lora_dir, "--prompt", prompt_eff]

        # --- 图生图输入: IMAGE(HWC float 0-1) -> 临时 PNG 交给 sd-cli ---
        # v1.2 注: 走 ref_image 时 denoise_strength 无效 (那是 init_image 路线的参数);
        #          走 init_image 时编辑 LoRA/ref_preset 不参与 (lora 只是公共参数被加载)
        def _save_img(t, tag):
            arr = t
            if arr.dim() == 4:  # [B,H,W,C] 取第一张
                arr = arr[0]
            img = Image.fromarray(
                (arr.cpu().numpy().clip(0, 1) * 255.0).astype(np.uint8))
            p = os.path.join(out_dir, f"{tag}.png")
            img.save(p)
            return p

        # init_image: 传统 img2img(加噪重绘)。sd1.5/sdxl 全强度可用;
        #   Krea-2-Turbo 建议 strength<=0.5; z-image-turbo 劣化禁用
        if init_image is not None:
            cmd += ["--init-img", _save_img(init_image, "init"),
                    "--strength", str(denoise_strength)]
        # ref_image: 参考图编辑(Flux Kontext/Krea2 edit 等 pass_to_dit 架构)。
        #   Krea-2 编辑务必配套编辑 LoRA, 否则预期出现颜色失真+网纹伪影
        if ref_image is not None:
            cmd += ["--ref-image", _save_img(ref_image, "ref")]
            if ref_preset and ref_preset != "default":
                cmd += ["--ref-image-args", f"preset={ref_preset}"]

        if extra_args.strip():
            cmd += shlex.split(extra_args)

        # 默认允许 ggml-vulkan 在设备显存不足时回退到系统内存:
        # Intel Iris Xe 等核显/UMA 设备的 device-local 只有 ~128-256MB,
        # 采样计算缓冲(345MB)直接分配会 ErrorOutOfDeviceMemory。
        # 注意: ggml-vulkan 只检查环境变量"存在性"(非空即启用, 不看值),
        # 所以显式禁用必须删掉变量(设 =0 无效)。
        env = os.environ.copy()
        if env.get("GGML_VK_ALLOW_SYSMEM_FALLBACK") is None:
            env["GGML_VK_ALLOW_SYSMEM_FALLBACK"] = "1"
        elif env.get("GGML_VK_ALLOW_SYSMEM_FALLBACK") == "0":
            env.pop("GGML_VK_ALLOW_SYSMEM_FALLBACK")
        try:
            proc = subprocess.run(cmd, capture_output=True, text=True,
                                  encoding="utf-8", errors="replace",
                                  creationflags=CREATE_NO_WINDOW, env=env)
        except FileNotFoundError:
            raise RuntimeError(f"无法启动 sd-cli.exe: {SD_CLI}")

        if proc.returncode != 0:
            tail = (proc.stderr or "")[-2000:]
            hint = ""
            if "ErrorOutOfDeviceMemory" in tail:
                hint = (
                    "\n\n[提示] Vulkan 显存分配失败(常见于核显/低显存卡)。"
                    "可在 extra_args 填:\n"
                    "--max-vram vulkan0=0.2 --params-backend diffusion=cpu\n"
                    "(Iris Xe 实测 512x512≈156s / 1024x1024≈903s; "
                    "注意 --max-vram 图分段在 z-image 上会改变出图,\n"
                    "需要标准结果请改用 --backend te=cpu,diffusion=cpu 全 CPU; "
                    "sd1.5/sdxl 不受分段影响)")
            raise RuntimeError(f"sd-cli 失败 (exit {proc.returncode}):\n{tail}{hint}")

        files = sorted(glob.glob(os.path.join(out_dir, "out*.png")))
        if not files:
            raise RuntimeError("sd-cli 成功退出但没有产出图像文件")

        images = []
        for f in files:
            with Image.open(f) as im:
                im = im.convert("RGB")
                arr = np.array(im).astype(np.float32) / 255.0
            images.append(torch.from_numpy(arr)[None, ...])
        out = torch.cat(images, dim=0)
        return (out, seed)


NODE_CLASS_MAPPINGS = {"SDCppGenerator": SDCppGenerator}
NODE_DISPLAY_NAME_MAPPINGS = {"SDCppGenerator": "SD.cpp Generator (Vulkan)"}