Update app.py

app.py

@@ -55,6 +55,32 @@ def rewrite_prompt(input_prompt):
 
 # --- 2. Preprocessor Functions ---
 
+def resize_image(input_image, max_size=1024):
+    """
+    Resizes an image so that its longest side is `max_size` pixels,
+    maintaining aspect ratio. The final dimensions are made divisible by 8.
+    """
+    w, h = input_image.size
+    aspect_ratio = w / h
+
+    if w > h:
+        new_w = max_size
+        new_h = int(new_w / aspect_ratio)
+    else:
+        new_h = max_size
+        new_w = int(new_h * aspect_ratio)
+
+    # Make dimensions divisible by 8
+    new_w = new_w - (new_w % 8)
+    new_h = new_h - (new_h % 8)
+
+    # Handle potential zero dimensions after rounding
+    if new_w == 0: new_w = 8
+    if new_h == 0: new_h = 8
+
+    return input_image.resize((new_w, new_h), Image.Resampling.LANCZOS)
+
+
 def extract_canny(input_image):
     image = np.array(input_image)
     image = cv2.Canny(image, 100, 200)
@@ -104,7 +130,7 @@ anyline = AnylineDetector.from_pretrained("TheMistoAI/MistoLine", filename="MTEE
 print("All models loaded.")
 
 def get_control_image(input_image, control_mode):
-    """A master function to select and run the correct preprocessor."""
+    """A master function to select and run the correct preprocessor on a pre-resized image."""
     if control_mode == "Canny":
         return extract_canny(input_image)
     elif control_mode == "Soft Edge":
@@ -143,6 +169,8 @@ def generate(
     if not prompt:
         raise gr.Error("Please enter a prompt.")
 
+    resized_image = resize_image(image, max_size=1024)
+
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
@@ -151,7 +179,7 @@ def generate(
         print(f"Original prompt: {prompt}\nEnhanced prompt: {enhanced_prompt}")
         prompt = enhanced_prompt
 
-    control_image = get_control_image(image, conditioning)
+    control_image = get_control_image(resized_image, conditioning)
     generator = torch.Generator(device=device).manual_seed(int(seed))
 
     generated_image = pipe(
@@ -159,8 +187,8 @@ def generate(
         negative_prompt=negative_prompt,
         control_image=control_image,
         controlnet_conditioning_scale=controlnet_conditioning_scale,
-        width=image.width,
-        height=image.height,
+        width=resized_image.width,
+        height=resized_image.height,
         num_inference_steps=int(num_inference_steps),
         guidance_scale=guidance_scale,
         generator=generator,
@@ -229,8 +257,4 @@ with gr.Blocks(css=css, theme=gr.themes.Citrus()) as demo:
     )
 
 if __name__ == "__main__":
-    if not os.path.exists("assets"):
-        os.makedirs("assets")
-        print("Created 'assets' directory. Please add example images for the Gradio examples to work.")
-    
     demo.launch()