fffiloni commited on
Commit
cc16ea6
1 Parent(s): 5963b10

Update to SD 1.5

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -16,7 +16,7 @@ YOUR_TOKEN=MY_SECRET_TOKEN
16
 
17
  device="cpu"
18
 
19
- pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=YOUR_TOKEN)
20
  pipe.to(device)
21
 
22
  gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
@@ -27,7 +27,7 @@ def infer(prompt):
27
  images_list = pipe([prompt] * 4)
28
  images = []
29
  safe_image = Image.open(r"unsafe.png")
30
- for i, image in enumerate(images_list["sample"]):
31
  if(images_list["nsfw_content_detected"][i]):
32
  images.append(safe_image)
33
  else:
 
16
 
17
  device="cpu"
18
 
19
+ pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", use_auth_token=YOUR_TOKEN)
20
  pipe.to(device)
21
 
22
  gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[2], height="auto")
 
27
  images_list = pipe([prompt] * 4)
28
  images = []
29
  safe_image = Image.open(r"unsafe.png")
30
+ for i, image in enumerate(images_list["images"]):
31
  if(images_list["nsfw_content_detected"][i]):
32
  images.append(safe_image)
33
  else: