From 18435772d12ac59e98dfb2555ad380359c6efd9f Mon Sep 17 00:00:00 2001
From: chflame163 <chflame@163.com>
Date: Thu, 11 Dec 2025 21:28:45 +0800
Subject: [PATCH] fix bug of Lama node

---
 py/iopaint/helper.py     |  6 ++++-
 py/iopaint/imghdr.py     | 55 ++++++++++++++++++++++++++++++++++++++++
 py/iopaint/model/base.py |  7 +++++
 py/iopaint/model/fcf.py  | 10 ++++++++
 pyproject.toml           |  4 +--
 requirements.txt         |  4 ++-
 6 files changed, 82 insertions(+), 4 deletions(-)
 create mode 100644 py/iopaint/imghdr.py

diff --git a/py/iopaint/helper.py b/py/iopaint/helper.py
index 35d348e..4e805d5 100644
--- a/py/iopaint/helper.py
+++ b/py/iopaint/helper.py
@@ -1,5 +1,5 @@
 import base64
-import imghdr
+
 import io
 import os
 import sys
@@ -16,6 +16,10 @@ from torch.hub import download_url_to_file, get_dir
 import hashlib
 from .const import DEFAULT_MODEL_DIR
 
+try :
+    import imghdr
+except ImportError:
+    from .imghdr import *
 
 def md5sum(filename):
     md5 = hashlib.md5()
diff --git a/py/iopaint/imghdr.py b/py/iopaint/imghdr.py
new file mode 100644
index 0000000..380f197
--- /dev/null
+++ b/py/iopaint/imghdr.py
@@ -0,0 +1,55 @@
+# imghdr.py - Determine the type of an image
+# Copied from Python standard library
+
+import io
+import os
+
+def what(file, h=None):
+    if h is None:
+        if isinstance(file, str):
+            f = None
+            try:
+                f = open(file, 'rb')
+                h = f.read(32)
+            finally:
+                if f:
+                    f.close()
+        elif hasattr(file, 'read'):
+            pos = file.tell()
+            h = file.read(32)
+            file.seek(pos)
+        else:
+            return None
+    for name, test in tests:
+        res = test(h)
+        if res:
+            return name
+    return None
+
+def test_jpeg(h):
+    if h[0:3] == b'\xff\xd8\xff':
+        return 'jpeg'
+def test_png(h):
+    if h[:8] == b'\211PNG\r\n\032\n':
+        return 'png'
+def test_gif(h):
+    if h[:6] in (b'GIF87a', b'GIF89a'):
+        return 'gif'
+def test_tiff(h):
+    if h[:2] in (b'MM', b'II'):
+        return 'tiff'
+def test_bmp(h):
+    if h[:2] == b'BM':
+        return 'bmp'
+def test_webp(h):
+    if h[0:4] == b'RIFF' and h[8:12] == b'WEBP':
+        return 'webp'
+
+tests = [
+    ('jpeg', test_jpeg),
+    ('png', test_png),
+    ('gif', test_gif),
+    ('tiff', test_tiff),
+    ('bmp', test_bmp),
+    ('webp', test_webp),
+]
diff --git a/py/iopaint/model/base.py b/py/iopaint/model/base.py
index e992e00..058b32f 100644
--- a/py/iopaint/model/base.py
+++ b/py/iopaint/model/base.py
@@ -76,6 +76,13 @@ class InpaintModel:
 
         if config.sd_keep_unmasked_area:
             mask = mask[:, :, np.newaxis]
+            # Ensure mask shape is (H, W, 1)
+            if mask.ndim == 4 and mask.shape[-1] == 1:
+                mask = mask.squeeze(-1)  # remove last dim
+            if mask.ndim == 4:
+                mask = mask.squeeze()
+            if mask.ndim == 3 and mask.shape[-1] != 1:
+                mask = mask[:, :, :1]
             result = result * (mask / 255) + image[:, :, ::-1] * (1 - (mask / 255))
         return result
 
diff --git a/py/iopaint/model/fcf.py b/py/iopaint/model/fcf.py
index 0eb2c1f..3fb5b75 100644
--- a/py/iopaint/model/fcf.py
+++ b/py/iopaint/model/fcf.py
@@ -1692,6 +1692,16 @@ class FcF(InpaintModel):
             )
 
             original_pixel_indices = crop_mask < 127
+
+            # --- FIX: Ensure boolean mask matches image shape ---
+            if original_pixel_indices.ndim == 3 and original_pixel_indices.shape[-1] == 1:
+                # Convert (H, W, 1) -> (H, W)
+                original_pixel_indices = original_pixel_indices.squeeze(-1)
+
+            if original_pixel_indices.ndim == 2:
+                # Convert (H, W) -> (H, W, 1) for channel broadcast
+                original_pixel_indices = np.repeat(original_pixel_indices[:, :, None], 3, axis=2)
+
             inpaint_result[original_pixel_indices] = crop_image[:, :, ::-1][
                 original_pixel_indices
             ]
diff --git a/pyproject.toml b/pyproject.toml
index f5b690c..c8fd189 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,9 +1,9 @@
 [project]
 name = "comfyui_layerstyle_advance"
 description = "The nodes detached from ComfyUI Layer Style are mainly those with complex requirements for dependency packages."
-version = "2.0.27"
+version = "2.0.28"
 license = { text = "MIT License" }
-dependencies = ["numpy", "matplotlib", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "blend_modes", "transformers", "diffusers", "loguru", "colour-science", "huggingface_hub", "segment_anything", "addict", "omegaconf", "yapf", "wget", "iopath", "mediapipe", "typer_config", "fastapi", "rich", "google-generativeai", "ultralytics", "transparent-background", "accelerate", "onnxruntime", "bitsandbytes", "peft", "protobuf", "hydra-core", "blind-watermark", "qrcode", "pyzbar", "psd-tools", "wandb", "zhipuai", "openai","google-genai"]
+dependencies = ["numpy", "matplotlib", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "blend_modes", "transformers", "diffusers", "loguru", "colour-science", "huggingface_hub", "segment_anything", "addict", "omegaconf", "yapf", "wget", "iopath", "mediapipe", "typer_config", "fastapi", "rich", "google-generativeai", "ultralytics", "transparent-background", "accelerate", "onnxruntime", "bitsandbytes", "peft", "protobuf", "hydra-core", "blind-watermark", "qrcode", "pyzbar", "psd-tools", "wandb", "zhipuai", "openai","google-genai", "fastapi","typer-config"]
 
 [project.urls]
 Repository = "https://github.com/chflame163/ComfyUI_LayerStyle_Advance"
diff --git a/requirements.txt b/requirements.txt
index 4f826e7..8d188e6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -37,4 +37,6 @@ qrcode
 pyzbar
 psd-tools
 zhipuai
-openai
\ No newline at end of file
+openai
+typer-config
+fastapi
\ No newline at end of file