enable text_encoder cpu

2022-09-30 21:39:23 +08:00
parent db1d7d5c48
commit dba7b01da7
4 changed files with 21 additions and 10 deletions
--- a/lama_cleaner/model/sd.py
+++ b/lama_cleaner/model/sd.py
@@ -91,7 +91,8 @@ class SD(InpaintModel):

        if kwargs['sd_cpu_textencoder']:
            logger.info("Run Stable Diffusion TextEncoder on CPU")
-            self.model.text_encoder = self.model.text_encoder.to(torch.device('cpu'))
+            self.model.text_encoder = self.model.text_encoder.to(torch.device('cpu'), non_blocking=True)
+            self.model.text_encoder = self.model.text_encoder.to(torch.float32, non_blocking=True )

        self.callbacks = kwargs.pop("callbacks", None)