Update app.py

app.py

@@ -24,25 +24,36 @@ img_pipe = StableDiffusionImg2ImgPipeline.from_pretrained("AkiKagura/mkgen-diffu
 img_pipe.safety_checker = empty_checker
 img_pipe.to(device)
 
-source_img = gr.Image(source="upload", type="filepath", label="init_img | 512*512 px")
+source_img = gr.Image(source="upload", type="filepath", label="init_img")
 gallery = gr.Gallery(label="Generated images", show_label=False, elem_id="gallery").style(grid=[1], height="auto")
 
-def resize(value,img):
+def resize(img):
   #baseheight = value
   img = Image.open(img)
   #hpercent = (baseheight/float(img.size[1]))
   #wsize = int((float(img.size[0])*float(hpercent)))
   #img = img.resize((wsize,baseheight), Image.Resampling.LANCZOS)
-  img = img.resize((value,value), Image.Resampling.LANCZOS)
-  return img
+  hsize = img.size[1]
+  wsize = img.size[0]
+  if 6*wsize <= 5*hsize:
+      wsize = 512
+      hsize = 768
+  elif 4*wsize >= 5*hsize:
+      wsize = 768
+      hsize = 512
+  else:
+      wsize = 512
+      hsize = 512
+  img = img.resize((wsize,hsize), Image.Resampling.LANCZOS)
+  return img, wsize, hsize
 
 
 def infer(source_img, prompt, guide, steps, seed, strength): 
     generator = torch.Generator('cpu').manual_seed(seed)
     
-    source_image = resize(512, source_img)
+    source_image, img_w, img_h = resize(source_img)
     source_image.save('source.png')
-    images_list = img_pipe([prompt] * 1, init_image=source_image, strength=strength, guidance_scale=guide, num_inference_steps=steps)
+    images_list = img_pipe([prompt] * 1, init_image=source_image, strength=strength, guidance_scale=guide, num_inference_steps=steps, width=img_w, height=img_h)
     images = []
     
     for i, image in enumerate(images_list["images"]):