From 59ee89bd34403b7710b4da3a8e19586ce73c6b8f Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Tue, 3 Jan 2023 21:30:33 +0800
Subject: [PATCH] add --no-half arg

---
 lama_cleaner/app/src/components/shared/Toast.tsx | 1 -
 lama_cleaner/model/paint_by_example.py           | 3 ++-
 lama_cleaner/model/sd.py                         | 3 ++-
 lama_cleaner/parse_args.py                       | 1 +
 lama_cleaner/server.py                           | 1 +
 5 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/lama_cleaner/app/src/components/shared/Toast.tsx b/lama_cleaner/app/src/components/shared/Toast.tsx
index 004d49f..e12d928 100644
--- a/lama_cleaner/app/src/components/shared/Toast.tsx
+++ b/lama_cleaner/app/src/components/shared/Toast.tsx
@@ -74,7 +74,6 @@ const Toast = React.forwardRef<
 })
 
 Toast.defaultProps = {
-  desc: '',
   state: 'loading',
 }
 
diff --git a/lama_cleaner/model/paint_by_example.py b/lama_cleaner/model/paint_by_example.py
index 991e836..d0bd3d5 100644
--- a/lama_cleaner/model/paint_by_example.py
+++ b/lama_cleaner/model/paint_by_example.py
@@ -15,8 +15,9 @@ class PaintByExample(InpaintModel):
     min_size = 512
 
     def init_model(self, device: torch.device, **kwargs):
+        fp16 = not kwargs['no_half']
         use_gpu = device == torch.device('cuda') and torch.cuda.is_available()
-        torch_dtype = torch.float16 if use_gpu else torch.float32
+        torch_dtype = torch.float16 if use_gpu and fp16 else torch.float32
         self.model = DiffusionPipeline.from_pretrained(
             "Fantasy-Studio/Paint-by-Example",
             torch_dtype=torch_dtype,
diff --git a/lama_cleaner/model/sd.py b/lama_cleaner/model/sd.py
index 49b68f5..56817eb 100644
--- a/lama_cleaner/model/sd.py
+++ b/lama_cleaner/model/sd.py
@@ -30,6 +30,7 @@ class SD(InpaintModel):
 
     def init_model(self, device: torch.device, **kwargs):
         from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline
+        fp16 = not kwargs['no_half']
 
         model_kwargs = {"local_files_only": kwargs['sd_run_local']}
         if kwargs['sd_disable_nsfw']:
@@ -42,7 +43,7 @@ class SD(InpaintModel):
         torch_dtype = torch.float16 if use_gpu else torch.float32
         self.model = StableDiffusionInpaintPipeline.from_pretrained(
             self.model_id_or_path,
-            revision="fp16" if use_gpu else "main",
+            revision="fp16" if use_gpu and fp16 else "main",
             torch_dtype=torch_dtype,
             use_auth_token=kwargs["hf_access_token"],
             **model_kwargs
diff --git a/lama_cleaner/parse_args.py b/lama_cleaner/parse_args.py
index e01867d..fb6e035 100644
--- a/lama_cleaner/parse_args.py
+++ b/lama_cleaner/parse_args.py
@@ -12,6 +12,7 @@ def parse_args():
         default="lama",
         choices=["lama", "ldm", "zits", "mat", "fcf", "sd1.5", "cv2", "manga", "sd2", "paint_by_example"],
     )
+    parser.add_argument("--no-half", action="store_true", help="SD/PaintByExample model no half precision")
     parser.add_argument(
         "--hf_access_token",
         default="",
diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
index f447e00..085ea36 100644
--- a/lama_cleaner/server.py
+++ b/lama_cleaner/server.py
@@ -303,6 +303,7 @@ def main(args):
     model = ModelManager(
         name=args.model,
         device=device,
+        no_half=args.no_half,
         hf_access_token=args.hf_access_token,
         sd_disable_nsfw=args.sd_disable_nsfw,
         sd_cpu_textencoder=args.sd_cpu_textencoder,