zhiweili commited on
Commit
7688f40
1 Parent(s): d8653bf
Files changed (1) hide show
  1. app_masked.py +5 -1
app_masked.py CHANGED
@@ -10,8 +10,11 @@ from segment_utils import(
10
  )
11
  from diffusers import (
12
  DiffusionPipeline,
 
13
  )
14
 
 
 
15
  BASE_MODEL = "RunDiffusion/Juggernaut-XL-v9"
16
 
17
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -23,8 +26,9 @@ DEFAULT_CATEGORY = "hair"
23
 
24
  basepipeline = DiffusionPipeline.from_pretrained(
25
  BASE_MODEL,
 
26
  torch_dtype=torch.float16,
27
- # use_safetensors=True,
28
  custom_pipeline="./pipelines/masked_stable_diffusion_xl_img2img.py",
29
  )
30
 
 
10
  )
11
  from diffusers import (
12
  DiffusionPipeline,
13
+ AutoencoderKL,
14
  )
15
 
16
+ vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
17
+
18
  BASE_MODEL = "RunDiffusion/Juggernaut-XL-v9"
19
 
20
  DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
26
 
27
  basepipeline = DiffusionPipeline.from_pretrained(
28
  BASE_MODEL,
29
+ vae=vae,
30
  torch_dtype=torch.float16,
31
+ use_safetensors=True,
32
  custom_pipeline="./pipelines/masked_stable_diffusion_xl_img2img.py",
33
  )
34