Upload salia_square.py

salia_square.py

@@ -0,0 +1,454 @@
+import os
+import random
+import threading
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+from PIL import Image
+
+import folder_paths
+
+
+# -----------------------------------------------------------------------------
+# Global caches (shared across node instances)
+# -----------------------------------------------------------------------------
+_CKPT_CACHE: Dict[str, Dict[str, Any]] = {}
+_CONTROLNET_CACHE: Dict[str, Dict[str, Any]] = {}
+_CACHE_LOCK = threading.RLock()
+
+_ALLOWED_UPSCALE_FACTORS = (1, 2, 4, 6, 8, 10, 12, 14, 16)
+
+
+# -----------------------------------------------------------------------------
+# Lazy imports / caching
+# -----------------------------------------------------------------------------
+def _lazy_import_nodes():
+    import nodes  # comfy-core nodes module
+    return nodes
+
+
+def _get_mtime(path: Optional[str]) -> Optional[float]:
+    if not path:
+        return None
+    try:
+        return float(os.path.getmtime(path))
+    except Exception:
+        return None
+
+
+def _load_checkpoint_cached(ckpt_name: str):
+    """
+    Returns (model, clip, vae) for ckpt_name.
+    Cached by (ckpt_name + file mtime).
+    """
+    nodes = _lazy_import_nodes()
+
+    ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
+    if not ckpt_path:
+        raise FileNotFoundError(f"Checkpoint not found: {ckpt_name}")
+    mtime = _get_mtime(ckpt_path)
+
+    with _CACHE_LOCK:
+        entry = _CKPT_CACHE.get(ckpt_name)
+        if entry and entry.get("mtime") == mtime:
+            return entry["model"], entry["clip"], entry["vae"]
+
+    loader = nodes.CheckpointLoaderSimple()
+    model, clip, vae = loader.load_checkpoint(ckpt_name)
+
+    with _CACHE_LOCK:
+        _CKPT_CACHE[ckpt_name] = {"mtime": mtime, "model": model, "clip": clip, "vae": vae}
+
+    return model, clip, vae
+
+
+def _load_controlnet_cached(controlnet_name: str):
+    """
+    Returns controlnet for controlnet_name.
+    Cached by (controlnet_name + file mtime).
+    """
+    nodes = _lazy_import_nodes()
+
+    cn_path = folder_paths.get_full_path("controlnet", controlnet_name)
+    if not cn_path:
+        raise FileNotFoundError(f"ControlNet not found: {controlnet_name}")
+    mtime = _get_mtime(cn_path)
+
+    with _CACHE_LOCK:
+        entry = _CONTROLNET_CACHE.get(controlnet_name)
+        if entry and entry.get("mtime") == mtime:
+            return entry["controlnet"]
+
+    loader = nodes.ControlNetLoader()
+    (controlnet,) = loader.load_controlnet(control_net_name=controlnet_name)
+
+    with _CACHE_LOCK:
+        _CONTROLNET_CACHE[controlnet_name] = {"mtime": mtime, "controlnet": controlnet}
+
+    return controlnet
+
+
+# -----------------------------------------------------------------------------
+# Comfy tensor helpers (IMAGE/MASK)
+# -----------------------------------------------------------------------------
+def _ensure_batched_image(image: torch.Tensor) -> torch.Tensor:
+    # Accept [H,W,C] or [B,H,W,C]
+    if image.ndim == 3:
+        return image.unsqueeze(0)
+    return image
+
+
+def _ensure_batched_mask(mask: torch.Tensor) -> torch.Tensor:
+    # Accept [H,W] or [B,H,W] or [B,H,W,1]
+    if mask.ndim == 2:
+        return mask.unsqueeze(0)
+    if mask.ndim == 4 and mask.shape[-1] == 1:
+        return mask[..., 0]
+    return mask
+
+
+def _comfy_image_batch_to_pil_list(img_bhwc: torch.Tensor) -> List[Image.Image]:
+    img_bhwc = _ensure_batched_image(img_bhwc).detach().cpu().float().clamp(0.0, 1.0)
+    b, h, w, c = img_bhwc.shape
+    if c not in (3, 4):
+        raise ValueError(f"Expected IMAGE with 3 or 4 channels, got {c} channels.")
+
+    out: List[Image.Image] = []
+    for i in range(b):
+        arr = (img_bhwc[i].numpy() * 255.0).round().astype(np.uint8)
+        mode = "RGB" if c == 3 else "RGBA"
+        out.append(Image.fromarray(arr, mode=mode))
+    return out
+
+
+def _pil_list_to_comfy_image_batch(pils: List[Image.Image], want_channels: int) -> torch.Tensor:
+    if want_channels not in (3, 4):
+        raise ValueError("want_channels must be 3 or 4")
+
+    tensors: List[torch.Tensor] = []
+    for p in pils:
+        p = p.convert("RGB") if want_channels == 3 else p.convert("RGBA")
+        arr = np.array(p).astype(np.float32) / 255.0
+        tensors.append(torch.from_numpy(arr))
+    return torch.stack(tensors, dim=0)
+
+
+def _resize_comfy_image_lanczos(img_bhwc: torch.Tensor, width: int, height: int) -> torch.Tensor:
+    img_bhwc = _ensure_batched_image(img_bhwc)
+    if width <= 0 or height <= 0:
+        raise ValueError("width/height must be > 0")
+
+    b, h, w, c = img_bhwc.shape
+    if (w == width) and (h == height):
+        return img_bhwc
+
+    pils = _comfy_image_batch_to_pil_list(img_bhwc)
+    resized = [p.resize((width, height), resample=Image.LANCZOS) for p in pils]
+    return _pil_list_to_comfy_image_batch(resized, want_channels=c)
+
+
+def _resize_comfy_mask_lanczos(mask_bhw: torch.Tensor, width: int, height: int) -> torch.Tensor:
+    mask_bhw = _ensure_batched_mask(mask_bhw).detach().cpu().float().clamp(0.0, 1.0)
+    b, h, w = mask_bhw.shape
+    if width <= 0 or height <= 0:
+        raise ValueError("width/height must be > 0")
+    if (w == width) and (h == height):
+        return mask_bhw
+
+    out: List[torch.Tensor] = []
+    for i in range(b):
+        arr = (mask_bhw[i].numpy() * 255.0).round().astype(np.uint8)
+        pil = Image.fromarray(arr, mode="L")
+        pil = pil.resize((width, height), resample=Image.LANCZOS)
+        arr2 = np.array(pil).astype(np.float32) / 255.0
+        out.append(torch.from_numpy(arr2))
+    return torch.stack(out, dim=0)
+
+
+def _repeat_batch_if_needed(t: torch.Tensor, target_b: int) -> torch.Tensor:
+    if t.ndim == 4:
+        b = t.shape[0]
+        if b == target_b:
+            return t
+        if b == 1:
+            return t.repeat(target_b, 1, 1, 1)
+        raise ValueError(f"Batch mismatch: tensor batch {b} vs target {target_b}")
+    if t.ndim == 3:
+        b = t.shape[0]
+        if b == target_b:
+            return t
+        if b == 1:
+            return t.repeat(target_b, 1, 1)
+        raise ValueError(f"Batch mismatch: tensor batch {b} vs target {target_b}")
+    raise ValueError("Unsupported tensor rank for batching")
+
+
+def _alpha_over_composite_at_xy(base_bhwc: torch.Tensor, overlay_bhwc: torch.Tensor, x: int, y: int) -> torch.Tensor:
+    """
+    Alpha composite overlay (must be RGBA) over base at (x,y).
+    Output channels match base channels (RGB stays RGB; RGBA stays RGBA).
+    """
+    base_bhwc = _ensure_batched_image(base_bhwc).detach().cpu().float().clamp(0.0, 1.0)
+    overlay_bhwc = _ensure_batched_image(overlay_bhwc).detach().cpu().float().clamp(0.0, 1.0)
+
+    b0, H, W, Cb = base_bhwc.shape
+    b1, h, w, Co = overlay_bhwc.shape
+
+    if Co != 4:
+        raise ValueError("overlay must be RGBA (4 channels).")
+    if Cb not in (3, 4):
+        raise ValueError("base must have 3 or 4 channels.")
+    if b1 != b0:
+        overlay_bhwc = _repeat_batch_if_needed(overlay_bhwc, b0)
+
+    if x < 0 or y < 0 or (x + w) > W or (y + h) > H:
+        raise ValueError(f"Overlay out of bounds: base {W}x{H}, overlay {w}x{h}, x={x}, y={y}")
+
+    out = base_bhwc.clone()
+
+    ov_rgb = overlay_bhwc[..., 0:3]
+    ov_a = overlay_bhwc[..., 3:4]
+
+    region = out[:, y : y + h, x : x + w, :]
+    bd_rgb = region[..., 0:3]
+
+    if Cb == 3:
+        out_rgb = ov_rgb * ov_a + bd_rgb * (1.0 - ov_a)
+        out[:, y : y + h, x : x + w, 0:3] = out_rgb
+        return out.clamp(0.0, 1.0)
+
+    bd_a = region[..., 3:4]
+    out_a = ov_a + bd_a * (1.0 - ov_a)
+    out_rgb_premul = ov_rgb * ov_a + bd_rgb * bd_a * (1.0 - ov_a)
+    out_rgb = torch.where(out_a > 1e-8, out_rgb_premul / out_a, torch.zeros_like(out_rgb_premul))
+
+    out[:, y : y + h, x : x + w, 0:3] = out_rgb
+    out[:, y : y + h, x : x + w, 3:4] = out_a
+    return out.clamp(0.0, 1.0)
+
+
+def _list_asset_pngs_fallback() -> List[str]:
+    """
+    Best-effort asset PNG listing:
+      1) Try comfyui-salia_online/utils/io.py:list_pngs()
+      2) Else scan ../assets/images relative to this file
+    """
+    try:
+        from ..utils.io import list_pngs  # your plugin helper
+        choices = list_pngs()
+        if choices:
+            return choices
+    except Exception:
+        pass
+
+    try:
+        plugin_root = Path(__file__).resolve().parent.parent
+        images_dir = plugin_root / "assets" / "images"
+        if images_dir.exists():
+            return sorted([p.name for p in images_dir.glob("*.png")])
+    except Exception:
+        pass
+
+    return []
+
+
+# -----------------------------------------------------------------------------
+# The one-node workflow
+# -----------------------------------------------------------------------------
+class Salia_OneNode_SquareWorkflow:
+    """
+    One-node replacement for the described workflow.
+    """
+
+    CATEGORY = "image/salia"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "run"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        # Keep INPUT_TYPES light: no model loads here.
+        try:
+            import comfy.samplers as samplers
+            sampler_names = list(getattr(samplers.KSampler, "SAMPLERS", [])) or ["euler"]
+            scheduler_names = list(getattr(samplers.KSampler, "SCHEDULERS", [])) or ["karras"]
+        except Exception:
+            sampler_names = ["euler"]
+            scheduler_names = ["karras"]
+
+        ckpts = folder_paths.get_filename_list("checkpoints") or ["<no checkpoints found>"]
+        cns = folder_paths.get_filename_list("controlnet") or ["<no controlnets found>"]
+        assets = _list_asset_pngs_fallback() or ["<no pngs found>"]
+
+        upscale_choices = [str(v) for v in _ALLOWED_UPSCALE_FACTORS]
+
+        return {
+            "required": {
+                "image": ("IMAGE",),
+
+                "X_coord": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}),
+                "Y_coord": ("INT", {"default": 0, "min": 0, "max": 16384, "step": 1}),
+                "square_size": ("INT", {"default": 384, "min": 1, "max": 8192, "step": 1}),
+
+                "positive_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+
+                "upscale_factor": (upscale_choices, {"default": "4"}),
+
+                "checkpoint_name": (ckpts, {}),
+                "controlnet_name": (cns, {}),
+                "assets_png": (assets, {}),
+
+                "controlnet_strength": ("FLOAT", {"default": 0.33, "min": 0.0, "max": 10.0, "step": 0.01}),
+                "controlnet_start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+                "controlnet_end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
+
+                "steps": ("INT", {"default": 30, "min": 1, "max": 200, "step": 1}),
+                "cfg": ("FLOAT", {"default": 2.6, "min": 0.0, "max": 10.0, "step": 0.05}),
+                "sampler_name": (sampler_names, {"default": "euler"}),
+                "scheduler": (scheduler_names, {"default": "karras"}),
+                "denoise": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
+
+    def run(
+        self,
+        image: torch.Tensor,
+        X_coord: int,
+        Y_coord: int,
+        square_size: int,
+        positive_prompt: str,
+        negative_prompt: str,
+        upscale_factor: str,
+        checkpoint_name: str,
+        controlnet_name: str,
+        assets_png: str,
+        controlnet_strength: float = 0.33,
+        controlnet_start_percent: float = 0.0,
+        controlnet_end_percent: float = 1.0,
+        steps: int = 30,
+        cfg: float = 2.6,
+        sampler_name: str = "euler",
+        scheduler: str = "karras",
+        denoise: float = 0.35,
+    ):
+        # ---- validate ----
+        try:
+            uf = int(upscale_factor)
+        except Exception:
+            raise ValueError(f"Invalid upscale_factor: {upscale_factor}")
+
+        if uf not in _ALLOWED_UPSCALE_FACTORS:
+            raise ValueError(f"upscale_factor must be one of {_ALLOWED_UPSCALE_FACTORS}, got {uf}")
+
+        if square_size <= 0:
+            raise ValueError("square_size must be > 0")
+
+        # ---- crop ----
+        base = _ensure_batched_image(image)
+        b, H, W, C = base.shape
+        if C not in (3, 4):
+            raise ValueError(f"Input image must be RGB or RGBA (3/4 channels), got {C}")
+
+        x = int(X_coord)
+        y = int(Y_coord)
+        s = int(square_size)
+
+        if x < 0 or y < 0 or (x + s) > W or (y + s) > H:
+            raise ValueError(f"Crop out of bounds: image {W}x{H}, x={x}, y={y}, square_size={s}")
+
+        crop = base[:, y : y + s, x : x + s, :]
+        crop_rgb = crop[..., 0:3]
+
+        up_w = int(s * uf)
+        up_h = int(s * uf)
+
+        # ---- upscale crop (Lanczos) for VAEEncode ----
+        crop_up = _resize_comfy_image_lanczos(crop_rgb, width=up_w, height=up_h)
+
+        # ---- depth (Salia_Depth) then upscale depth (Lanczos) ----
+        # lazy import (don’t import transformers at module import time)
+        try:
+            from .salia_depth import Salia_Depth
+        except Exception:
+            from salia_depth import Salia_Depth
+
+        depth_node = Salia_Depth()
+        (depth_img,) = depth_node.execute(image=crop, resolution=-1)  # keep original crop res
+        depth_img = _ensure_batched_image(depth_img)[..., 0:3]
+        depth_up = _resize_comfy_image_lanczos(depth_img, width=up_w, height=up_h)
+
+        # ---- load alpha mask from assets ----
+        try:
+            from .salia_loadimage_assets import LoadImage_SaliaOnline_Assets
+        except Exception:
+            from salia_loadimage_assets import LoadImage_SaliaOnline_Assets
+
+        assets_loader = LoadImage_SaliaOnline_Assets()
+        _asset_img, asset_mask = assets_loader.run(assets_png)
+        asset_mask = _ensure_batched_mask(asset_mask)
+        asset_mask = _resize_comfy_mask_lanczos(asset_mask, width=up_w, height=up_h)
+        asset_mask = _repeat_batch_if_needed(asset_mask, b)
+
+        # ---- load checkpoint + controlnet (cached) ----
+        model, clip, vae = _load_checkpoint_cached(checkpoint_name)
+        controlnet = _load_controlnet_cached(controlnet_name)
+
+        # ---- comfy core pipeline ----
+        nodes = _lazy_import_nodes()
+
+        (pos_cond,) = nodes.CLIPTextEncode().encode(clip=clip, text=positive_prompt)
+        (neg_cond,) = nodes.CLIPTextEncode().encode(clip=clip, text=negative_prompt)
+
+        pos_cn, neg_cn = nodes.ControlNetApplyAdvanced().apply_controlnet(
+            positive=pos_cond,
+            negative=neg_cond,
+            control_net=controlnet,
+            image=depth_up,
+            strength=float(controlnet_strength),
+            start_percent=float(controlnet_start_percent),
+            end_percent=float(controlnet_end_percent),
+            vae=vae,
+        )
+
+        (latent,) = nodes.VAEEncode().encode(pixels=crop_up, vae=vae)
+
+        # No seed input requested: generate a fresh seed per execution
+        seed = random.randint(0, 2**63 - 1)
+
+        (latent_out,) = nodes.KSampler().sample(
+            model=model,
+            seed=seed,
+            steps=int(steps),
+            cfg=float(cfg),
+            sampler_name=sampler_name,
+            scheduler=scheduler,
+            positive=pos_cn,
+            negative=neg_cn,
+            latent_image=latent,
+            denoise=float(denoise),
+        )
+
+        (decoded_rgb,) = nodes.VAEDecode().decode(samples=latent_out, vae=vae)
+
+        # Join alpha -> RGBA
+        (decoded_rgba_up,) = nodes.JoinImageWithAlpha().join(image=decoded_rgb, alpha=asset_mask)
+
+        # Downscale back to original square size (Lanczos)
+        decoded_rgba_down = _resize_comfy_image_lanczos(decoded_rgba_up, width=s, height=s)
+
+        # Composite onto original input at (x,y)
+        out = _alpha_over_composite_at_xy(base, decoded_rgba_down, x=x, y=y)
+
+        return (out,)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Salia_OneNode_SquareWorkflow": Salia_OneNode_SquareWorkflow,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "Salia_OneNode_SquareWorkflow": "Salia One-Node Square Workflow",
+}