Handle any size

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+import numpy as np
 import torch
 
 from diffusers import AutoPipelineForInpainting, UNet2DConditionModel
@@ -30,12 +31,46 @@ def predict(dict, prompt="", negative_prompt="", guidance_scale=7.5, steps=20, s
     scheduler = getattr(diffusers, scheduler_class_name)
     pipe.scheduler = scheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler", **add_kwargs)
     
-    init_image = dict["image"].convert("RGB").resize((1024, 1024))
-    mask = dict["mask"].convert("RGB").resize((1024, 1024))
+    init_image = dict["image"].convert("RGB")
+    mask = dict["mask"].convert("RGB")
+    original_height, original_width, original_channel = np.array(init_image).shape
+
+    # Limited to 1 million pixels
+    if 1024 * 1024 < original_width * original_height:
+        factor = ((1024 * 1024) / (original_width * original_height))**0.5
+        process_width = math.floor(original_width * factor)
+        process_height = math.floor(original_height * factor)
+    else:
+        process_width = original_width
+        process_height = original_height
+
+    # Width and height must be multiple of 8
+    if (process_width % 8) != 0 or (process_height % 8) != 0:
+        process_width = process_width - (process_width % 8)
+        process_height = process_height - (process_height % 8)
+
+        if ((process_width + 8) * (process_height + 8)) <= (1024 * 1024):
+            process_width = process_width + 8
+            process_height = process_height + 8
+
+    init_image = init_image.resize((process_width, process_height))
+    mask = mask.resize((process_width, process_height))
     
-    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)
+    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,
+        width = process_width,
+        height = process_height
+    )
+
+    output_image = output.images[0].resize((original_width, original_height))
     
-    return output.images[0], gr.update(visible=True)
+    return output_image, gr.update(visible=True)
 
 
 css = '''
@@ -80,7 +115,7 @@ with image_blocks as demo:
     gr.HTML(read_content("header.html"))
     with gr.Row():
                 with gr.Column():
-                    image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload",height=400)
+                    image = gr.Image(source='upload', tool='sketch', elem_id="image_upload", type="pil", label="Upload")
                     with gr.Row(elem_id="prompt-container", mobile_collapse=False, equal_height=True):
                         with gr.Row():
                             prompt = gr.Textbox(placeholder="Your prompt (what you want in place of what is erased)", show_label=False, elem_id="prompt")
@@ -97,7 +132,7 @@ with image_blocks as demo:
                             scheduler = gr.Dropdown(label="Schedulers", choices=schedulers, value="EulerDiscreteScheduler")
                         
                 with gr.Column():
-                    image_out = gr.Image(label="Output", elem_id="output-img", height=400)
+                    image_out = gr.Image(label="Output", elem_id="output-img")
                     with gr.Group(elem_id="share-btn-container", visible=False) as share_btn_container:
                         community_icon = gr.HTML(community_icon_html)
                         loading_icon = gr.HTML(loading_icon_html)