Update app.py

app.py

@@ -2,6 +2,23 @@ import transformers
 import gradio as gr
 import datasets
 import torch
+import os
 from diffusers import StableDiffusionPipeline
-import matplotlib.pyplot as plt
-from mpl_toolkits.axes_grid1 import ImageGrid
+
+# Initialise the token
+HF_TOKEN = os.getenv('HF_TOKEN')
+
+pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", use_auth_token=HF_TOKEN)
+device = torch.device('coda' if torch.cuda.is_available() else 'cpu')
+pipe = pipe.to(device)
+
+hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, 'crowdsourced-movieposter-demo')
+
+gr.Interface(
+  fn=generate,
+  inputs=[gr.Textbox(label='Celebrity Name'), gr.Dropdown(['Mission Impossible', 'Avatar', 'Abyss', 'TOO OLD TO DIE YOUNG', 'The Last Dual','The beginning of everything'], label='Movie'), gr.Slider(label='Image Accuracy', minimum=7.5, maximum=20, value=12)],
+  outputs=gr.Image(type='pil') ,
+  allow_flagging="manual",
+  flagging_options=["No Image", "good", "bad"],
+  flagging_callback=hf_writer
+).launch()