From 069c3b20922a55ea485343eb3eeffef6b11d339d Mon Sep 17 00:00:00 2001
From: leafiy <t@leafiy.com>
Date: Tue, 11 Nov 2025 10:10:05 +0800
Subject: [PATCH] update

---
 README.md          |  7 ++--
 psd_mockup_node.py | 80 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 6c9fbf6..8f1f174 100644
--- a/README.md
+++ b/README.md
@@ -12,9 +12,10 @@ matches what you would get when replacing the layer inside Photoshop.
   bounding box, transform box). Use it to confirm the layer names you want to
   target.
 - **PSD Mockup Embedder** – Uploads PSD/PSB files directly from the node,
-  replaces the selected smart objects with your image, resizes the result to any
-  width (default 800 px), and emits the composited tensor plus an auto-saved
-  PNG.
+  reuses the PSD’s original clipping stack (shadows, highlights, rounded
+  corners, etc.), projects your image into the smart objects, resizes the result
+  to any width (default 800 px), and emits the composited tensor plus an
+  auto-saved PNG.
 - Works with PSD/PSB files, multiple smart objects, and either ComfyUI `IMAGE`
   tensors or file paths.
 
diff --git a/psd_mockup_node.py b/psd_mockup_node.py
index 6fc69f4..27221f5 100644
--- a/psd_mockup_node.py
+++ b/psd_mockup_node.py
@@ -233,6 +233,70 @@ def _polygon_mask(quad: Sequence[float], canvas_size: Tuple[int, int]) -> Image.
     return mask
 
 
+def _collect_clipping_stack(layer: Layer) -> List[Layer]:
+    stack = [layer]
+    parent = layer.parent
+    if parent is None:
+        return stack
+    siblings = list(parent)
+    try:
+        idx = siblings.index(layer)
+    except ValueError:
+        return stack
+    for sib in siblings[idx + 1 :]:
+        if not getattr(sib, "clipping", False):
+            break
+        stack.append(sib)
+    return stack
+
+
+def _composite_layers(psd: PSDImage, layer_ids: Sequence[int]) -> Image.Image:
+    include = set(layer_ids)
+
+    def filter_fn(layer: Layer) -> bool:
+        if hasattr(layer, "layer_id") and layer.layer_id in include:
+            return True
+        return False
+
+    return psd.composite(layer_filter=filter_fn)
+
+
+def _compute_shading_components(stack_image: Image.Image, base_image: Image.Image):
+    stack_rgba = stack_image.convert("RGBA")
+    base_rgba = base_image.convert("RGBA")
+    stack_arr = np.array(stack_rgba).astype(np.float32) / 255.0
+    base_arr = np.array(base_rgba).astype(np.float32) / 255.0
+    base_rgb = base_arr[..., :3]
+    stack_rgb = stack_arr[..., :3]
+    base_alpha = base_arr[..., 3:4]
+
+    multiplier = np.ones_like(base_rgb, dtype=np.float32)
+    mask = base_alpha > 1e-3
+    if np.any(mask):
+        denom = np.clip(base_rgb, 1e-3, None)
+        ratio = stack_rgb / denom
+        ratio = np.clip(ratio, 0.0, 4.0)
+        multiplier[mask.repeat(3, axis=2)] = ratio[mask.repeat(3, axis=2)]
+
+    base_alpha_rgb = np.repeat(base_alpha, 3, axis=2)
+    multiplier = multiplier * base_alpha_rgb + (1.0 - base_alpha_rgb)
+    alpha = stack_arr[..., 3:4]
+    alpha_img = Image.fromarray((np.clip(alpha, 0.0, 1.0) * 255).astype(np.uint8).squeeze(-1), mode="L")
+    return multiplier, alpha_img
+
+
+def _apply_shading_map(image: Image.Image, multiplier: np.ndarray) -> Image.Image:
+    if multiplier is None:
+        return image
+    rgba = image.convert("RGBA")
+    arr = np.array(rgba).astype(np.float32) / 255.0
+    h, w, _ = arr.shape
+    if multiplier.shape[:2] != (h, w):
+        return image
+    arr[..., :3] = np.clip(arr[..., :3] * multiplier, 0.0, 1.0)
+    return Image.fromarray((arr * 255).astype(np.uint8), mode="RGBA")
+
+
 def _output_directory() -> str:
     if folder_paths and hasattr(folder_paths, "get_output_directory"):
         return folder_paths.get_output_directory()
@@ -347,9 +411,21 @@ class PSDMockupEmbedder:
             if quad is None:
                 left, top, right, bottom = layer.bbox
                 quad = (left, top, right, top, right, bottom, left, bottom)
+
+            stack_layers = _collect_clipping_stack(layer)
+            stack_ids = [l.layer_id for l in stack_layers if hasattr(l, "layer_id")]
+            stack_image = _composite_layers(psd, stack_ids) if stack_ids else None
+            base_image = _composite_layers(psd, [layer.layer_id])
+            shading_multiplier = None
+            alpha_mask = None
+            if stack_image is not None:
+                shading_multiplier, alpha_mask = _compute_shading_components(stack_image, base_image)
+
             warped = _warp_image_to_canvas(source_image, quad, canvas_size)
-            mask = _polygon_mask(quad, canvas_size)
-            base_rgba = Image.composite(warped, base_rgba, mask)
+            warped = _apply_shading_map(warped, shading_multiplier)
+            fallback_mask = _polygon_mask(quad, canvas_size)
+            composite_mask = alpha_mask if alpha_mask and alpha_mask.getbbox() else fallback_mask
+            base_rgba = Image.composite(warped, base_rgba, composite_mask)
             metadata.append(
                 {
                     "layer_id": layer.layer_id,