update app.py

app.py

@@ -17,13 +17,15 @@ pipeline_upscale = StableDiffusionUpscalePipeline.from_pretrained(model_id, torc
 pipeline_upscale = pipeline_upscale.to("cuda")
 
 def get_IF_op(prompt, neg_prompt):
+  print("inside get_IF_op")
   filepaths = client_if.predict(prompt, neg_prompt, 1,4,7.0, 'smart100',50, api_name="/generate64")
   folder_path = filepaths[0]
   file_list = os.listdir(folder_path)
   file_list = [os.path.join(folder_path, f) for f in file_list if f != 'captions.json']
-  return file_list  
+  return file_list    
 
 def get_pickscores(prompt, file_list):
+  print("inside get_pickscores")
   #Get the predictons
   probabilities1 = client_pick.predict(prompt, file_list[0], file_list[1], fn_index=0)
   probabilities2 = client_pick.predict(prompt, file_list[2], file_list[3], fn_index=0)
@@ -33,10 +35,30 @@ def get_pickscores(prompt, file_list):
   best_match_image = file_list[max_score_index]
   return best_match_image
 
-def get_upscale_op(prompt, best_match_image):
-  # let's get the image
+def get_upscale_op(prompt, gallery_if):
+  print("inside get_upscale_op")
+  # get pickscores
+  best_match_image = get_pickscores(prompt, gallery_if)
+  # let's get the best pick! 
   low_res_img = Image.open(best_match_image).convert("RGB")
   low_res_img = low_res_img.resize((128, 128))
+  # Upscaling the best pick
   upscaled_image = pipeline_upscale(prompt=prompt, image=low_res_img).images[0]
+  #upscaled_image.save("upsampled.png")
   return upscaled_image
 
+
+with gr.Blocks() as demo:
+  with gr.Row():
+    with gr.Column:
+      prompt = gr.Textbox(label='Prompt')
+      neg_prompt = gr.Textbox(label='Negative Prompt')
+    b1 = gr.Button('Generate')
+  gallery_if = gr.Gallery(label='IF Space outputs')
+  b2 = gr.Button("Get the best generation using Pick-A-Pic")
+  image_picakapic = gr.Image(label="PickAPic Evaluated Output") 
+
+  b1.click(get_IF_op,[prompt, neg_prompt], gallery_if)
+  b1.click(get_upscale_op,[prompt, gallery_if], image_picakapic)
+
+demo.launch(debug=True)