lots update

lama_cleaner/model/sd.py

@@ -3,9 +3,9 @@ import cv2
 import torch
 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.model.utils import handle_from_pretrained_exceptions
 from lama_cleaner.schema import Config, ModelType
 
 
@@ -40,20 +40,18 @@ class SD(DiffusionInpaintModel):
                 model_kwargs["num_in_channels"] = 9
 
             self.model = StableDiffusionInpaintPipeline.from_single_file(
-                self.model_id_or_path, torch_dtype=torch_dtype, **model_kwargs
+                self.model_id_or_path, dtype=torch_dtype, **model_kwargs
             )
         else:
-            self.model = StableDiffusionInpaintPipeline.from_pretrained(
-                self.model_id_or_path,
-                revision="fp16"
-                if self.model_id_or_path in DIFFUSERS_MODEL_FP16_REVERSION
-                else "main",
-                torch_dtype=torch_dtype,
+            self.model = handle_from_pretrained_exceptions(
+                StableDiffusionInpaintPipeline.from_pretrained,
+                pretrained_model_name_or_path=self.model_id_or_path,
+                variant="fp16",
+                dtype=torch_dtype,
                 **model_kwargs,
             )
 
         if kwargs.get("cpu_offload", False) and use_gpu:
-            # TODO: gpu_id
             logger.info("Enable sequential cpu offload")
             self.model.enable_sequential_cpu_offload(gpu_id=0)
         else:
@@ -98,20 +96,20 @@ class SD(DiffusionInpaintModel):
 
 
 class SD15(SD):
-    name = "sd1.5"
+    name = "runwayml/stable-diffusion-inpainting"
     model_id_or_path = "runwayml/stable-diffusion-inpainting"
 
 
 class Anything4(SD):
-    name = "anything4"
+    name = "Sanster/anything-4.0-inpainting"
     model_id_or_path = "Sanster/anything-4.0-inpainting"
 
 
 class RealisticVision14(SD):
-    name = "realisticVision1.4"
+    name = "Sanster/Realistic_Vision_V1.4-inpainting"
     model_id_or_path = "Sanster/Realistic_Vision_V1.4-inpainting"
 
 
 class SD2(SD):
-    name = "sd2"
+    name = "stabilityai/stable-diffusion-2-inpainting"
     model_id_or_path = "stabilityai/stable-diffusion-2-inpainting"