amazonaws-la commited on
Commit
07c0e5d
1 Parent(s): 739a3ba

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -14,7 +14,7 @@ from PIL import Image
14
  from io import BytesIO
15
  from diffusers import StableDiffusionImg2ImgPipeline, AutoencoderKL, DiffusionPipeline
16
  from diffusers.utils import load_image
17
- from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
18
 
19
  DESCRIPTION = "# SDXL"
20
  if not torch.cuda.is_available():
@@ -64,8 +64,9 @@ def generate(
64
  ):
65
  if torch.cuda.is_available():
66
 
67
- if not use_vae:
68
- pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model, torch_dtype=torch.float16, safety_checker:StableDiffusionSafetyChecker)
 
69
 
70
  if use_vae:
71
  vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)
 
14
  from io import BytesIO
15
  from diffusers import StableDiffusionImg2ImgPipeline, AutoencoderKL, DiffusionPipeline
16
  from diffusers.utils import load_image
17
+ from safety_checker import StableDiffusionSafetyChecker
18
 
19
  DESCRIPTION = "# SDXL"
20
  if not torch.cuda.is_available():
 
64
  ):
65
  if torch.cuda.is_available():
66
 
67
+ if not use_vae:
68
+ safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker")
69
+ pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model, torch_dtype=torch.float16)
70
 
71
  if use_vae:
72
  vae = AutoencoderKL.from_pretrained(vaecall, torch_dtype=torch.float16)