diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
index 02fdfb2..b43ee37 100644
--- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
+++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
@@ -10,7 +10,7 @@ import SettingBlock from './SettingBlock'
 export enum AIModel {
   LAMA = 'lama',
   LDM = 'ldm',
-  ZITS = 'ZITS',
+  ZITS = 'zits',
 }
 
 function ModelSettingBlock() {
diff --git a/lama_cleaner/model/zits.py b/lama_cleaner/model/zits.py
index c08e173..9f2f175 100644
--- a/lama_cleaner/model/zits.py
+++ b/lama_cleaner/model/zits.py
@@ -229,8 +229,8 @@ class ZITS(InpaintModel):
     def wireframe_edge_and_line(self, items, enable: bool):
         # 最终向 items 中添加 edge 和 line key
         if not enable:
-            items["edge"] = torch.zeros_like(items["mask_256"])
-            items["line"] = torch.zeros_like(items["mask_256"])
+            items["edge"] = torch.zeros_like(items["masks"])
+            items["line"] = torch.zeros_like(items["masks"])
             return
 
         start = time.time()
@@ -297,6 +297,7 @@ class ZITS(InpaintModel):
         masks: [H, W]
         return: BGR IMAGE
         """
+        mask = mask[:, :, 0]
         items = load_image(image, mask, device=self.device)
 
         self.wireframe_edge_and_line(items, config.zits_wireframe)