Merserk-ComfyUI-PiD/pid_sample.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Tuple
import os
import subprocess
import sys
import tempfile

import torch

try:
    from .pid_decode import SEQUENTIAL_OFFLOAD_CHOICES, PiDNodeError, _free_cuda_memory
    from .pid_prepare import PID_PREP_TYPE, PiDPreparedBatch
except ImportError:  # pragma: no cover
    from pid_decode import SEQUENTIAL_OFFLOAD_CHOICES, PiDNodeError, _free_cuda_memory
    from pid_prepare import PID_PREP_TYPE, PiDPreparedBatch


PID_SAMPLES_TYPE = "PID_SAMPLES"


@dataclass
class PiDSampledBatch:
    tensor_cpu: torch.Tensor
    backbone: str
    pid_ckpt_type: str
    infer_image_size: Tuple[int, int]


class PiDSample:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "prepared": (PID_PREP_TYPE,),
                "pid_steps": ("INT", {"default": 4, "min": 1, "max": 64, "step": 1}),
                "cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 2**31 - 1}),
                "aggressive_cleanup": ("BOOLEAN", {"default": True}),
                "sequential_offload": (SEQUENTIAL_OFFLOAD_CHOICES, {"default": "disabled"}),
            }
        }

    RETURN_TYPES = (PID_SAMPLES_TYPE,)
    RETURN_NAMES = ("sampled",)
    FUNCTION = "sample"
    CATEGORY = "PiD/Staged"

    def sample(
        self,
        prepared: PiDPreparedBatch,
        pid_steps: int,
        cfg_scale: float,
        seed: int,
        aggressive_cleanup: bool = True,
        sequential_offload: str = "disabled",
    ):
        if not isinstance(prepared, PiDPreparedBatch):
            raise PiDNodeError("PiD Sample expected a PID_PREP object from PiD Prepare.")
        sequential_offload = str(sequential_offload or "disabled").strip().lower()
        if sequential_offload not in SEQUENTIAL_OFFLOAD_CHOICES:
            raise PiDNodeError(
                f"Unknown sequential_offload={sequential_offload!r}; expected one of {SEQUENTIAL_OFFLOAD_CHOICES}"
            )

        _free_cuda_memory(aggressive=True)
        runner = Path(__file__).resolve().with_name("pid_subprocess_runner.py")
        if not runner.is_file():
            raise PiDNodeError(f"Missing PiD subprocess runner: {runner}")

        with tempfile.TemporaryDirectory(prefix="comfyui_pid_") as tmp:
            tmpdir = Path(tmp)
            input_path = tmpdir / "pid_input.pt"
            output_path = tmpdir / "pid_output.pt"
            payload = {
                "pid_dir": prepared.pid_dir,
                "model_dir": prepared.model_dir,
                "backbone": prepared.backbone,
                "pid_ckpt_type": prepared.pid_ckpt_type,
                "checkpoint_path": prepared.checkpoint_path,
                "caption": prepared.caption,
                "sigma": float(prepared.sigma),
                "scale": int(prepared.scale),
                "infer_image_size": tuple(int(x) for x in prepared.infer_image_size),
                "latent_cpu": prepared.latent_cpu.detach().to("cpu").contiguous(),
                "baseline_cpu": (
                    prepared.baseline_cpu.detach().to("cpu").contiguous()
                    if prepared.baseline_cpu is not None
                    else None
                ),
                "baseline_size": tuple(int(x) for x in prepared.baseline_size),
            }
            torch.save(payload, str(input_path))
            del payload
            _free_cuda_memory(aggressive=True)

            cmd = [
                sys.executable or "python",
                str(runner),
                "--input",
                str(input_path),
                "--output",
                str(output_path),
                "--pid-steps",
                str(int(pid_steps)),
                "--cfg-scale",
                str(float(cfg_scale)),
                "--seed",
                str(int(seed)),
                "--sequential-offload",
                sequential_offload,
            ]
            if aggressive_cleanup:
                cmd.append("--aggressive-cleanup")

            env = os.environ.copy()
            node_dir = str(Path(__file__).resolve().parent)
            env["PYTHONPATH"] = node_dir + os.pathsep + env.get("PYTHONPATH", "")
            proc = subprocess.run(
                cmd,
                cwd=node_dir,
                env=env,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                encoding="utf-8",
                errors="replace",
            )
            if proc.returncode != 0 or not output_path.is_file():
                tail = "\n".join((proc.stdout or "").splitlines()[-120:])
                raise PiDNodeError(
                    "PiD subprocess sampling failed. This usually means the 4K PiD pass still exceeded VRAM, "
                    "or the subprocess could not import/load PiD.\n\n"
                    f"Command: {' '.join(cmd)}\n\n"
                    f"Subprocess log tail:\n{tail}"
                )

            try:
                result = torch.load(str(output_path), map_location="cpu", weights_only=False)
            except TypeError:
                result = torch.load(str(output_path), map_location="cpu")

        _free_cuda_memory(aggressive=True)
        sampled = PiDSampledBatch(
            tensor_cpu=result["tensor_cpu"].detach().to("cpu"),
            backbone=str(result.get("backbone", prepared.backbone)),
            pid_ckpt_type=str(result.get("pid_ckpt_type", prepared.pid_ckpt_type)),
            infer_image_size=tuple(int(x) for x in result.get("infer_image_size", prepared.infer_image_size)),
        )
        return (sampled,)


NODE_CLASS_MAPPINGS = {
    "PiDSample": PiDSample,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PiDSample": "PiD Sample",
}