ms1449 commited on
Commit
64cdf13
·
verified ·
1 Parent(s): a70b44b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -1
app.py CHANGED
@@ -8,10 +8,17 @@ from huggingface_hub import login
8
  hf_token = os.getenv("HF_TOKEN") # This should be set in your Hugging Face Space Secrets
9
  login(token=hf_token)
10
 
 
 
 
11
  # Load the model using the API key and move it to GPU
12
  def load_model():
13
  model_id = "NVIDIA/sdxl-turbo"
14
- pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_auth_token=True)
 
 
 
 
15
  pipe = pipe.to("cuda")
16
  return pipe
17
 
 
8
  hf_token = os.getenv("HF_TOKEN") # This should be set in your Hugging Face Space Secrets
9
  login(token=hf_token)
10
 
11
+ # Access the NVIDIA API key from the environment
12
+ nvidia_api_key = os.getenv("NVIDIA_API_KEY")
13
+
14
  # Load the model using the API key and move it to GPU
15
  def load_model():
16
  model_id = "NVIDIA/sdxl-turbo"
17
+ pipe = StableDiffusionPipeline.from_pretrained(
18
+ model_id,
19
+ torch_dtype=torch.float16,
20
+ use_auth_token=nvidia_api_key # Use the API key here
21
+ )
22
  pipe = pipe.to("cuda")
23
  return pipe
24