vislupus commited on
Commit
d364b55
·
verified ·
1 Parent(s): 861b276

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -2
app.py CHANGED
@@ -13,8 +13,12 @@ if torch.cuda.is_available():
13
  else:
14
  torch_dtype = torch.float32
15
 
16
- pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
17
- pipe = pipe.to(device)
 
 
 
 
18
 
19
  MAX_SEED = np.iinfo(np.int32).max
20
  MAX_IMAGE_SIZE = 512
 
13
  else:
14
  torch_dtype = torch.float32
15
 
16
+ pipe = DiffusionPipeline.from_pretrained(model_repo_id, revision="fp16", torch_dtype=torch_dtype)
17
+ pipe.enable_model_cpu_offload()
18
+ pipe.enable_vae_slicing()
19
+
20
+ def safety_checker(images, **kwargs): return images, False
21
+ pipe.safety_checker = safety_checker
22
 
23
  MAX_SEED = np.iinfo(np.int32).max
24
  MAX_IMAGE_SIZE = 512