Update app.py

app.py

@@ -1,11 +1,41 @@
 import gradio as gr
 import torch
-
+import numpy as np
 import diffusers
 import os
 hf_token = os.environ.get("HF_TOKEN")
 from diffusers import StableDiffusionXLInpaintPipeline, DDIMScheduler, UNet2DConditionModel
 
+ratios_map =  {
+    0.5:{"width":704,"height":1408},
+    0.57:{"width":768,"height":1344},
+    0.68:{"width":832,"height":1216},
+    0.72:{"width":832,"height":1152},
+    0.78:{"width":896,"height":1152},
+    0.82:{"width":896,"height":1088},
+    0.88:{"width":960,"height":1088},
+    0.94:{"width":960,"height":1024},
+    1.00:{"width":1024,"height":1024},
+    1.13:{"width":1088,"height":960},
+    1.21:{"width":1088,"height":896},
+    1.29:{"width":1152,"height":896},
+    1.38:{"width":1152,"height":832},
+    1.46:{"width":1216,"height":832},
+    1.67:{"width":1280,"height":768},
+    1.75:{"width":1344,"height":768},
+    2.00:{"width":1408,"height":704}
+}
+ratios = np.array(ratios_map.keys())
+
+def get_size(init_image):
+    w,h=init_image.size
+    curr_ratio = w/h
+    ind = np.argmin(np.abs(curr_ratio-ratios))
+    ratio = ratios[ind]
+    chosen_ratio  = ratios_map[ratio]
+    w,h = chosen_ratio['width'], chosen_ratio['height']
+
+    return w,h
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 unet = UNet2DConditionModel.from_pretrained(
@@ -42,8 +72,20 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=5, steps=30, str
     if negative_prompt == "":
         negative_prompt = None
 
-    init_image = dict["image"].convert("RGB").resize((1024, 1024))
-    mask = dict["mask"].convert("RGB").resize((1024, 1024))
+    
+    init_image = dict["image"].convert("RGB")#.resize((1024, 1024))
+    mask = dict["mask"].convert("RGB")#.resize((1024, 1024))
+
+    w,h = get_size(init_image)
+
+    init_image = init_image.resize((w, h))
+    mask = mask.resize((w, h))
+    
+    # Resize to nearest ratio ?
+    
+    mask = np.array(mask)
+    mask[mask>0]=255
+    mask = Image.fromarray(mask)
     
     output = pipe(prompt = prompt, negative_prompt=negative_prompt, image=init_image, mask_image=mask, guidance_scale=guidance_scale, num_inference_steps=int(steps), strength=strength)