create app.py

app.py

@@ -0,0 +1,42 @@
+#import requests
+import gradio as gr
+from gradio_client import Client
+from PIL import Image
+from io import BytesIO
+from diffusers import StableDiffusionUpscalePipeline
+import torch
+import os
+
+HF_TOKEN = os.environ.get('HF_TOKEN')
+client_if = Client("huggingface-projects/IF", hf_token=HF_TOKEN)
+client_pick = Client("yuvalkirstain/PickScore") 
+
+# load upscaling model and scheduler
+model_id = "stabilityai/stable-diffusion-x4-upscaler"
+pipeline_upscale = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+pipeline_upscale = pipeline_upscale.to("cuda")
+
+def get_IF_op(prompt, neg_prompt):
+  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  
+
+def get_pickscores(prompt, file_list):
+  #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)
+  probabilities_all = list(probabilities1) + list(probabilities2)
+  max_score = max(probabilities_all)
+  max_score_index = probabilities_all.index(max_score)
+  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
+  low_res_img = Image.open(best_match_image).convert("RGB")
+  low_res_img = low_res_img.resize((128, 128))
+  upscaled_image = pipeline_upscale(prompt=prompt, image=low_res_img).images[0]
+  return upscaled_image
+