Update app.py

app.py

@@ -12,6 +12,7 @@ from torchvision import transforms as tfms
 from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer, logging
 import os
+MY_TOKEN=os.environ.get('Stable_Diffusion')
 import cv2
 import torchvision.transforms as T
 
@@ -21,7 +22,7 @@ logging.set_verbosity_error()
 torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the autoencoder
-vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder='vae')
+vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae",use_auth_token=MY_TOKEN)
 
 # Load tokenizer and text encoder to tokenize and encode the text
 tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
@@ -38,9 +39,9 @@ vae = vae.to(torch_device)
 text_encoder = text_encoder.to(torch_device)
 unet = unet.to(torch_device)
 
-style_files = ['Thumps_up.bin', 'birb_style.bin',
-                     'snoopy.bin', 'pop_art.bin',
-                     'boot-mjstyle.bin']
+style_files = ['bird_style.bin', 'ronaldo.bin',
+                     'pop_art.bin', 'threestooges.bin',
+                     'bflan.bin']
 
 images_without_loss = []
 images_with_loss = []
@@ -251,7 +252,7 @@ def display_images_in_rows(images_with_titles, titles):
     # plt.show()
 
 
-def image_generator(prompt = "snoopy", loss_function=None):
+def image_generator(prompt = "sky", loss_function=None):
   images_without_loss = []
   images_with_loss = []
   if loss_function == "Yes":
@@ -267,7 +268,7 @@ def image_generator(prompt = "snoopy", loss_function=None):
         images_with_loss.append(generated_img)
 
   generated_sd_images = []
-  titles = ["Bird_style", "Boot-mjstyle", "Snoopy Style", "Pop Art Style", "Thumpsup Style"]
+  titles = ["<birb-style>", "'<ronaldo>", "<pop-art>", "<threestooges>", "<Marbled-painting>"]
 
   for i in range(len(titles)):
     generated_sd_images.append((images_without_loss[i], titles[i]))