| #!/usr/bin/env python | |
| import os | |
| import shutil | |
| import sys | |
| from diffusers import StableDiffusionPipeline | |
| from diffusers.pipelines.stable_diffusion.safety_checker import \ | |
| StableDiffusionSafetyChecker | |
| # append project directory to path so predict.py can be imported | |
| sys.path.append('.') | |
| from predict import MODEL_CACHE, MODEL_ID, SAFETY_MODEL_ID | |
| if os.path.exists(MODEL_CACHE): | |
| shutil.rmtree(MODEL_CACHE) | |
| os.makedirs(MODEL_CACHE, exist_ok=True) | |
| saftey_checker = StableDiffusionSafetyChecker.from_pretrained( | |
| SAFETY_MODEL_ID, | |
| cache_dir=MODEL_CACHE, | |
| ) | |
| pipe = StableDiffusionPipeline.from_pretrained( | |
| MODEL_ID, | |
| cache_dir=MODEL_CACHE, | |
| ) | |