Update model.py

model.py

@@ -1,5 +1,6 @@
 from PIL import Image
 from base64 import b64encode
+import os
 
 import torch
 from torch import autocast
@@ -95,7 +96,7 @@ text_tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
 embeddings= text_model.forward(texts, text_tokenizer, device )
 # 1. Load the autoencoder model which will be used to decode the latents into image space.
 vae = AutoencoderKL.from_pretrained(
-    'CompVis/stable-diffusion-v1-4', subfolder='vae', use_auth_token=True)
+    'CompVis/stable-diffusion-v1-4', subfolder='vae', use_auth_token=True, token=os.environ['HF_API_TOKEN'])
 vae = vae.to(device)
 
 # 2. Load the tokenizer and text encoder to tokenize and encode the text.
@@ -104,7 +105,7 @@ text_encoder = text_model
 
 # 3. The UNet model for generating the latents.
 unet = UNet2DConditionModel.from_pretrained(
-    'CompVis/stable-diffusion-v1-4', subfolder='unet', use_auth_token=True)
+    'CompVis/stable-diffusion-v1-4', subfolder='unet', use_auth_token=True, token=os.environ['HF_API_TOKEN'])
 unet = unet.to(device)
 
 # 4. Create a scheduler for inference