Update app.py

app.py

@@ -5,14 +5,14 @@ import base64
 import torch
 import os
 
-
+auth_token = os.environ.get("auth_token")
 device = "cuda"
 generator = torch.Generator(device=device)
 
 seed = 496012807434005 #generator.seed()
 generator = generator.manual_seed(seed)
-HF_TOKEN = os.getenv('HF_TOKEN')
-hf_writer =gr.HuggingFaceDatasetSaver(HF_TOKEN, "dst-movie-poster-demo")
+#HF_TOKEN = os.getenv('HF_TOKEN')
+hf_writer =gr.HuggingFaceDatasetSaver(auth_token, "dst-movie-poster-demo")
 
 
 def improve_image(img):
@@ -23,7 +23,7 @@ def improve_image(img):
   resp_img = gr.processing_utils.decode_base64_to_image((resp_obj.json())['data'][0])
   return resp_img
 
-pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",use_auth_token= HF_TOKEN)
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4",use_auth_token= auth_token)
 pipe = pipe.to("cuda")
 
 def generate(celebrity, setting):
@@ -41,6 +41,6 @@ gr.Interface(
   outputs = gr.Image(type='pill'),
   allow_flagging="manual",
   flagging_options = ['Incorrect movie poster','Incorrect Actor','Other Problem'],
-  #flagging_callback=hf_writer,
+  flagging_callback=hf_writer,
   flagging_dir='flagged_data'
 ).launch()