Initial app.py

app.py

@@ -0,0 +1,44 @@
+import gradio as gr
+from diffusers import StableDiffusionPipeline
+import requests
+import base64
+import torch
+import os
+
+
+device = "cuda"
+generator = torch.Generator(device=device)
+
+HF_TOKEN = os.getenv('HF_TOKEN')
+hf_writer =gr.HuggingFaceDatasetSaver(HF_TOKEN, "dst-movie-poster-demo")
+
+
+def improve_image(img):
+  # ANSWER HERE
+  img_in_base64 = gr.processing_utils.encode_pil_to_base64(img)
+  scale=3
+  resp_obj = requests.post('https://hf.space/embed/abidlabs/GFPGAN/+/api/predict',json={'data':[img_in_base64,scale]})
+  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")
+pipe = pipe.to("cuda")
+
+def generate(celebrity, setting):
+  # ANSWER HERE
+  prompt = 'A movie poster of {} in the movie{}'.format(celebrity,setting)
+  latent_sample = torch.randn((1,4,64,64),generator = generator,device=device)
+  gen_img = pipe(prompt,latents=latent_sample,num_inference_steps=100,guidance_scale=g_scale[i]).images[0]
+  image =  improve_image(gen_img)
+  return image
+
+gr.Interface(
+  # ANSWER HERE    
+  fn=generate,
+  inputs=[gr.Textbox(label='Celebrity'), gr.Dropdown(['The Godfather', 'Titanic', 'Fast and Furious'], label='Movie')],
+  outputs = gr.Image(type='pill'),
+  allow_flagging="manual",
+  flagging_options = ['Incorrect movie poster','Incorrect Actor','Other Problem'],
+  flagging_callback=hf_writer,
+  flagging_dir='/flagged_data'
+).launch()