183 lines
6.8 KiB
Python
183 lines
6.8 KiB
Python
# ComfyUI-PIH/pih_node.py
|
|
import os, sys, tempfile, uuid, subprocess, glob, shutil
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
from .model import Model_Composite_PL, Model_Composite
|
|
import folder_paths as comfy_paths
|
|
import torch.nn.functional as F
|
|
|
|
|
|
# Comfy types
|
|
from typing import Any, Dict, Tuple
|
|
|
|
class PIH_Harmonize:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"Background Image": ("IMAGE",), # [B,H,W,C] float [0,1], RGB
|
|
"Foreground Image": ("IMAGE",), # [B,H,W,C], RGB or RGBA
|
|
"Strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), # Strength of harmonization
|
|
"PIH Processor": ("PIH_PROCESSOR",), # Model_Composite_PL instance
|
|
},
|
|
"optional": {
|
|
"mask": ("MASK",), # [B,H,W] float [0,1] (optional, multiplies FG alpha)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",) # Composited, harmonized RGB
|
|
FUNCTION = "run"
|
|
CATEGORY = "Harmonization/PIH"
|
|
|
|
def run(self, **kw):
|
|
bg = kw["Background Image"] # [B,H,W,C]
|
|
fg = kw["Foreground Image"] # [B,H,W,C] (RGB or RGBA)
|
|
strength = kw["Strength"] # float [0,1]
|
|
model = kw["PIH Processor"]
|
|
mask_in = kw.get("mask", None) # [B,H,W] or None
|
|
|
|
# device
|
|
device = next(model.parameters()).device if any(True for _ in model.parameters()) else ("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# to float32 on device
|
|
bg = bg.to(device=device, dtype=torch.float32)
|
|
fg = fg.to(device=device, dtype=torch.float32)
|
|
if mask_in is not None:
|
|
mask_in = mask_in.to(device=device, dtype=torch.float32)
|
|
|
|
B, Hf, Wf, Cf = fg.shape
|
|
Bb, Hb, Wb, Cb = bg.shape
|
|
if B != Bb:
|
|
raise ValueError("Batch size mismatch between Background and Foreground")
|
|
|
|
# Resize BG to FG spatial size if needed
|
|
if (Hb, Wb) != (Hf, Wf):
|
|
bg = F.interpolate(bg.permute(0,3,1,2), size=(Hf,Wf), mode="bilinear", align_corners=False).permute(0,2,3,1)
|
|
|
|
# Resize mask to FG spatial size if needed (ComfyUI MASK is [B,H,W])
|
|
if mask_in is not None and mask_in.shape[-2:] != (Hf, Wf):
|
|
mask_in = F.interpolate(mask_in.unsqueeze(1), size=(Hf, Wf), mode="bilinear", align_corners=False).squeeze(1)
|
|
|
|
# Split FG into RGB (+ alpha if present)
|
|
if Cf == 4:
|
|
alpha_fg = fg[..., 3:4] # [B,H,W,1], assumed straight alpha in [0,1]
|
|
rgb_fg = fg[..., :3] # [B,H,W,3]
|
|
else:
|
|
alpha_fg = None
|
|
rgb_fg = fg
|
|
|
|
# Effective mask: (provided mask) * (alpha if present) else (alpha) else ones
|
|
if mask_in is not None:
|
|
m = mask_in.unsqueeze(-1) # [B,H,W,1]
|
|
if alpha_fg is not None:
|
|
m = m * alpha_fg
|
|
else:
|
|
m = alpha_fg if alpha_fg is not None else torch.ones((B,Hf,Wf,1), device=device)
|
|
|
|
m = m.clamp(0,1)
|
|
|
|
# Convert to BCHW
|
|
bg_c = bg.permute(0,3,1,2).contiguous() # [B,3,H,W]
|
|
fg_c = rgb_fg.permute(0,3,1,2).contiguous() # [B,3,H,W]
|
|
mask_c = m.permute(0,3,1,2).contiguous() # [B,1,H,W]
|
|
|
|
# --- initial composite at full resolution (FG over BG using alpha) ---
|
|
print(f"Shapes - Mask: {mask_c.shape}, Foreground: {fg_c.shape}, Background: {bg_c.shape}")
|
|
comp_full = fg_c * mask_c + bg_c * (1 - mask_c) # [B,3,H,W]
|
|
|
|
# --- low-res pass (512x512) ---
|
|
target = (512, 512)
|
|
bg_low = F.interpolate(bg_c, size=target, mode="bilinear", align_corners=False)
|
|
comp_low = F.interpolate(comp_full, size=target, mode="bilinear", align_corners=False)
|
|
mask_low = F.interpolate(mask_c, size=target, mode="bilinear", align_corners=False)
|
|
|
|
with torch.no_grad():
|
|
_inter, _out, _scalars, _pl = model(bg_low, comp_low, mask_low)
|
|
|
|
# --- high-res reapply: LUT then gain, still respecting alpha mask ---
|
|
hr_intermediate = model.PL3D(model.pl_table, comp_full) * mask_c + (1 - mask_c) * bg_c
|
|
gain_full = F.interpolate(model.gainmap, size=(Hf, Wf), mode="bilinear", align_corners=False)
|
|
composite_c = (mask_c) * fg_c + (1 - mask_c) * bg_c
|
|
out_full = hr_intermediate * gain_full * strength * mask_c + (1 - strength * mask_c) * composite_c # [B,3,H,W]
|
|
|
|
# Back to Comfy format [B,H,W,C]
|
|
out_img = out_full.permute(0,2,3,1).clamp(0,1).contiguous()
|
|
return (out_img,)
|
|
|
|
|
|
|
|
class PIHModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("PIH_PROCESSOR",)
|
|
RETURN_NAMES = ("PIH Processor",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "Adobe PIH"
|
|
|
|
def load_model(self, checkpoint_folder="pretrained"):
|
|
print("[PIH] PIHModelLoader.load_model() called")
|
|
|
|
# device
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
# paths
|
|
pih_folder = os.path.join(comfy_paths.models_dir, "PIH")
|
|
ckpt_path = os.path.join(pih_folder, "ckpt_g39.pth")
|
|
|
|
if not os.path.exists(ckpt_path):
|
|
os.makedirs(pih_folder, exist_ok=True)
|
|
print(f"[PIH] Checkpoint not found — downloading to {ckpt_path} ...")
|
|
try:
|
|
import gdown
|
|
gdown.download(id="1seW8qSnaBOQ4_S9bQ4ThVOdeJGYJ-f74", output=ckpt_path, quiet=False)
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
f"[PIH] Auto-download failed: {e}\n"
|
|
f"Download manually from https://drive.google.com/file/d/1seW8qSnaBOQ4_S9bQ4ThVOdeJGYJ-f74/view "
|
|
f"and place it at: {ckpt_path}"
|
|
)
|
|
|
|
# build model
|
|
model = Model_Composite_PL(
|
|
dim=32,
|
|
masking=True,
|
|
brush=True,
|
|
maskoffset=0.6,
|
|
swap=True,
|
|
Vit_bool=False,
|
|
onlyupsample=True,
|
|
aggupsample=True,
|
|
light=False,
|
|
Eff_bool=False,
|
|
).to(device)
|
|
|
|
# load weights
|
|
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
|
|
sd = ckpt.get("state_dict", ckpt) # tolerate raw SDs
|
|
missing, unexpected = model.load_state_dict(sd, strict=False)
|
|
if missing or unexpected:
|
|
print(f"[PIH] load_state_dict: missing={missing}, unexpected={unexpected}")
|
|
|
|
model.eval()
|
|
for p in model.parameters():
|
|
p.requires_grad = False
|
|
|
|
print("[PIH] Model loaded and ready on", device)
|
|
return (model,)
|
|
|
|
|
|
|
|
# Node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"PIH Harmonize": PIH_Harmonize,
|
|
"PIHModelLoader": PIHModelLoader,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"PIH Harmonize": "PIH Harmonize",
|
|
"PIHModelLoader": "PIH Model Loader",
|
|
}
|
|
|