Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix mps crash with safety checker

Browse Source
This commit is contained in:
Lincoln Stein 2022-10-30 16:54:06 -04:00
parent 330b417a7b
commit 23d54ee69e
2 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
ldm/generate.py
View File

@ -217,6 +217,7 @@ class Generate:
safety_model_id = "CompVis/stable-diffusion-safety-checker" safety_model_id = "CompVis/stable-diffusion-safety-checker"
self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, local_files_only=True) self.safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, local_files_only=True)
self.safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id, local_files_only=True) self.safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id, local_files_only=True)
self.safety_checker.to(self.device)
except Exception: except Exception:
print('** An error was encountered while installing the safety checker:') print('** An error was encountered while installing the safety checker:')
print(traceback.format_exc()) print(traceback.format_exc())

1
ldm/invoke/generator/base.py
View File

@ -197,6 +197,7 @@ class Generator():
checker = self.safety_checker['checker'] checker = self.safety_checker['checker']
extractor = self.safety_checker['extractor'] extractor = self.safety_checker['extractor']
features = extractor([image], return_tensors="pt") features = extractor([image], return_tensors="pt")
features.to(self.model.device)
# unfortunately checker requires the numpy version, so we have to convert back # unfortunately checker requires the numpy version, so we have to convert back
x_image = np.array(image).astype(np.float32) / 255.0 x_image = np.array(image).astype(np.float32) / 255.0