update

lama_cleaner/model/base.py

@@ -14,7 +14,7 @@ from lama_cleaner.helper import (
 )
 from lama_cleaner.model.helper.g_diffuser_bot import expand_image
 from lama_cleaner.model.utils import get_scheduler
-from lama_cleaner.schema import Config, HDStrategy, SDSampler
+from lama_cleaner.schema import Config, HDStrategy, SDSampler, ModelInfo
 
 
 class InpaintModel:
@@ -266,9 +266,8 @@ class InpaintModel:
 
 class DiffusionInpaintModel(InpaintModel):
     def __init__(self, device, **kwargs):
-        if kwargs.get("model_id_or_path"):
-            # 用于自定义 diffusers 模型
-            self.model_id_or_path = kwargs["model_id_or_path"]
+        self.model_info: ModelInfo = kwargs["model_info"]
+        self.model_id_or_path = self.model_info.path
         super().__init__(device, **kwargs)
 
     @torch.no_grad()

lama_cleaner/model/controlnet.py

@@ -75,6 +75,11 @@ class ControlNet(DiffusionInpaintModel):
             sd_controlnet_method, torch_dtype=torch_dtype
         )
         if model_info.is_single_file_diffusers:
+            if self.model_info.model_type == ModelType.DIFFUSERS_SD:
+                model_kwargs["num_in_channels"] = 4
+            else:
+                model_kwargs["num_in_channels"] = 9
+
             self.model = PipeClass.from_single_file(
                 model_info.path, controlnet=controlnet
             ).to(torch_dtype)

lama_cleaner/model/instruct_pix2pix.py

@@ -8,7 +8,7 @@ from lama_cleaner.schema import Config
 
 
 class InstructPix2Pix(DiffusionInpaintModel):
-    name = "instruct_pix2pix"
+    name = "timbrooks/instruct-pix2pix"
     pad_mod = 8
     min_size = 512
 

lama_cleaner/model/kandinsky.py

@@ -73,7 +73,7 @@ class Kandinsky(DiffusionInpaintModel):
 
 
 class Kandinsky22(Kandinsky):
-    name = "kandinsky2.2"
+    name = "kandinsky-community/kandinsky-2-2-decoder-inpaint"
     model_name = "kandinsky-community/kandinsky-2-2-decoder-inpaint"
 
     @staticmethod

lama_cleaner/model/paint_by_example.py

@@ -9,7 +9,7 @@ from lama_cleaner.schema import Config
 
 
 class PaintByExample(DiffusionInpaintModel):
-    name = "paint_by_example"
+    name = "Fantasy-Studio/Paint-by-Example"
     pad_mod = 8
     min_size = 512
 

lama_cleaner/model/sd.py

@@ -9,7 +9,7 @@ from loguru import logger
 from lama_cleaner.const import DIFFUSERS_MODEL_FP16_REVERSION
 from lama_cleaner.model.base import DiffusionInpaintModel
 from lama_cleaner.model.helper.cpu_text_encoder import CPUTextEncoderWrapper
-from lama_cleaner.schema import Config
+from lama_cleaner.schema import Config, ModelType
 
 
 class SD(DiffusionInpaintModel):
@@ -36,7 +36,12 @@ class SD(DiffusionInpaintModel):
         use_gpu = device == torch.device("cuda") and torch.cuda.is_available()
         torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
 
-        if os.path.isfile(self.model_id_or_path):
+        if self.model_info.is_single_file_diffusers:
+            if self.model_info.model_type == ModelType.DIFFUSERS_SD:
+                model_kwargs["num_in_channels"] = 4
+            else:
+                model_kwargs["num_in_channels"] = 9
+
             self.model = StableDiffusionInpaintPipeline.from_single_file(
                 self.model_id_or_path, torch_dtype=torch_dtype, **model_kwargs
             )

lama_cleaner/model/sdxl.py

@@ -8,7 +8,7 @@ from diffusers import AutoencoderKL
 from loguru import logger
 
 from lama_cleaner.model.base import DiffusionInpaintModel
-from lama_cleaner.schema import Config
+from lama_cleaner.schema import Config, ModelType
 
 
 class SDXL(DiffusionInpaintModel):
@@ -26,9 +26,16 @@ class SDXL(DiffusionInpaintModel):
         use_gpu = device == torch.device("cuda") and torch.cuda.is_available()
         torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
 
+        if self.model_info.model_type == ModelType.DIFFUSERS_SDXL:
+            num_in_channels = 4
+        else:
+            num_in_channels = 9
+
         if os.path.isfile(self.model_id_or_path):
             self.model = StableDiffusionXLInpaintPipeline.from_single_file(
-                self.model_id_or_path, torch_dtype=torch_dtype
+                self.model_id_or_path,
+                torch_dtype=torch_dtype,
+                num_in_channels=num_in_channels,
             )
         else:
             vae = AutoencoderKL.from_pretrained(