apolinario commited on
Commit
90ca23c
1 Parent(s): 3c43e3d

Add performance upgrades

Browse files
Files changed (2) hide show
  1. app.py +8 -8
  2. requirements.txt +2 -1
app.py CHANGED
@@ -13,6 +13,8 @@ device = "cuda"
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
  #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=True)
18
  word_list = word_list_dataset["train"]['text']
@@ -25,14 +27,12 @@ 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"]):
 
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']
 
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"]):
requirements.txt CHANGED
@@ -2,4 +2,5 @@ diffusers
2
  transformers
3
  nvidia-ml-py3
4
  ftfy
5
- --extra-index-url https://download.pytorch.org/whl/cu113 torch
 
 
2
  transformers
3
  nvidia-ml-py3
4
  ftfy
5
+ --extra-index-url https://download.pytorch.org/whl/cu113
6
+ torch