Artvik commited on
Commit
6d84b3b
1 Parent(s): bd74936

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +11 -11
app.py CHANGED
@@ -11,10 +11,12 @@ model_id = "CompVis/stable-diffusion-v1-4"
11
  device = "cuda"
12
 
13
  #If you are running this code locally, you need to either do a 'huggingface-cli login` or paste your User Access Token from here https://huggingface.co/settings/tokens into the use_auth_token field below.
14
- pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=False, revision="fp16", torch_dtype=torch.float16)
15
  pipe = pipe.to(device)
 
 
16
  #When running locally, you won`t have access to this, so you can remove this part
17
- word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token="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")
18
  word_list = word_list_dataset["train"]['text']
19
 
20
  def infer(prompt, samples, steps, scale, seed):
@@ -25,16 +27,14 @@ def infer(prompt, samples, steps, scale, seed):
25
 
26
  generator = torch.Generator(device=device).manual_seed(seed)
27
 
28
- #If you are running locally with CPU, you can remove the `with autocast("cuda")`
29
- #with autocast("cuda"):
30
- #images_list = pipe(
31
- #[prompt] * samples,
32
- #num_inference_steps=steps,
33
- #guidance_scale=scale,
34
- #generator=generator,
35
- #)
36
  images = []
37
- #safe_image = Image.open(r"unsafe.png")
38
  for i, image in enumerate(images_list["sample"]):
39
  if(images_list["nsfw_content_detected"][i]):
40
  images.append(safe_image)
 
11
  device = "cuda"
12
 
13
  #If you are running this code locally, you need to either do a 'huggingface-cli login` or paste your User Access Token from here https://huggingface.co/settings/tokens into the use_auth_token field below.
14
+ pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=True, revision="fp16", torch_dtype=torch.float16)
15
  pipe = pipe.to(device)
16
+ torch.backends.cudnn.benchmark = True
17
+
18
  #When running locally, you won`t have access to this, so you can remove this part
19
+ word_list_dataset = load_dataset("stabilityai/word-list", data_files="list.txt", use_auth_token=True)
20
  word_list = word_list_dataset["train"]['text']
21
 
22
  def infer(prompt, samples, steps, scale, seed):
 
27
 
28
  generator = torch.Generator(device=device).manual_seed(seed)
29
 
30
+ images_list = pipe(
31
+ [prompt] * samples,
32
+ num_inference_steps=steps,
33
+ guidance_scale=scale,
34
+ generator=generator,
35
+ )
 
 
36
  images = []
37
+ safe_image = Image.open(r"unsafe.png")
38
  for i, image in enumerate(images_list["sample"]):
39
  if(images_list["nsfw_content_detected"][i]):
40
  images.append(safe_image)