cap image size, ensure result dimension is divisible by 8

app.py

@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import torch
 import spaces
 import gradio as gr
@@ -17,6 +19,29 @@ pipe = FluxInpaintPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
 
 
+def resize_image_dimensions(
+    original_resolution_wh: Tuple[int, int],
+    maximum_dimension: int = 2048
+) -> Tuple[int, int]:
+    width, height = original_resolution_wh
+
+    if width > height:
+        scaling_factor = maximum_dimension / width
+    else:
+        scaling_factor = maximum_dimension / height
+
+    new_width = int(width * scaling_factor)
+    new_height = int(height * scaling_factor)
+
+    new_width = new_width - (new_width % 8)
+    new_height = new_height - (new_height % 8)
+
+    new_width = min(maximum_dimension, new_width)
+    new_height = min(maximum_dimension, new_height)
+
+    return new_width, new_height
+
+
 @spaces.GPU()
 def process(input_image_editor, input_text, progress=gr.Progress(track_tqdm=True)):
     if not input_text:
@@ -34,7 +59,7 @@ def process(input_image_editor, input_text, progress=gr.Progress(track_tqdm=True
         gr.Info("Please draw a mask on the image.")
         return None
 
-    width, height = image.size
+    width, height = resize_image_dimensions(original_resolution_wh=image.size)
 
     return pipe(
         prompt=input_text,
@@ -42,7 +67,8 @@ def process(input_image_editor, input_text, progress=gr.Progress(track_tqdm=True
         mask_image=mask_image,
         width=width,
         height=height,
-        strength=0.7
+        strength=0.7,
+        num_inference_steps=2
     ).images[0]