lewtun HF staff commited on
Commit
2df430a
1 Parent(s): 0bdb4f3

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -1,4 +1,7 @@
1
  import torch
 
 
 
2
 
3
  device = "cuda" if torch.cuda.is_available() else "cpu"
4
  torch_dtype = torch.float16 if device == "cuda" else None
@@ -7,7 +10,7 @@ from diffusers import StableDiffusionPipeline
7
 
8
  model_id = "CompVis/stable-diffusion-v1-4"
9
  pipe = StableDiffusionPipeline.from_pretrained(
10
- model_id, revision="fp16", torch_dtype=torch_dtype
11
  ).to(device)
12
 
13
  def predict(prompt):
 
1
  import torch
2
+ import os
3
+
4
+ auth_token = os.getenv("HF_TOKEN")
5
 
6
  device = "cuda" if torch.cuda.is_available() else "cpu"
7
  torch_dtype = torch.float16 if device == "cuda" else None
 
10
 
11
  model_id = "CompVis/stable-diffusion-v1-4"
12
  pipe = StableDiffusionPipeline.from_pretrained(
13
+ model_id, use_auth_token=auth_token, revision="fp16", torch_dtype=torch_dtype
14
  ).to(device)
15
 
16
  def predict(prompt):