Update app.py

app.py

@@ -111,7 +111,7 @@ def validate_image(image_content: bytes):
         raise HTTPException(status_code=400, detail=f"Image validation error: {str(e)}")
 
 def generate_gemini_image(images: List[Image.Image], prompt: str):
-    """Generate or edit image using Gemini API and resize to 1024x1024 pixels"""
+    """Generate or edit image using Gemini API and resize to 1024x1024 pixels with aspect ratio preservation"""
     try:
         contents = images + [prompt]
         response = genai.GenerativeModel(GEMINI_MODEL).generate_content(contents)
@@ -126,15 +126,31 @@ def generate_gemini_image(images: List[Image.Image], prompt: str):
                 with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
                     temp_path = tmp.name
                     generated_image = Image.open(BytesIO(part.inline_data.data))
-                    # Resize image to 1024x1024 pixels
-                    generated_image = generated_image.resize((1024, 1024), Image.LANCZOS)
+
+                    # Preserve aspect ratio by padding to 1024x1024
+                    original_width, original_height = generated_image.size
+                    target_size = 1024
+                    ratio = min(target_size / original_width, target_size / original_height)
+                    new_width = int(original_width * ratio)
+                    new_height = int(original_height * ratio)
+
+                    # Resize while preserving aspect ratio
+                    generated_image = generated_image.resize((new_width, new_height), Image.LANCZOS)
+
+                    # Create a new 1024x1024 canvas with black background
+                    new_image = Image.new('RGB', (target_size, target_size), (0, 0, 0))
+                    # Paste the resized image in the center
+                    offset = ((target_size - new_width) // 2, (target_size - new_height) // 2)
+                    new_image.paste(generated_image, offset)
+
                     # Verify dimensions
-                    if generated_image.size != (1024, 1024):
-                        logger.error(f"Resized image dimensions {generated_image.size} do not match expected 1024x1024")
-                        raise HTTPException(status_code=500, detail="Failed to resize image to 1024x1024")
-                    generated_image.save(temp_path, format="PNG")
+                    if new_image.size != (1024, 1024):
+                        logger.error(f"Output image dimensions {new_image.size} do not match expected 1024x1024")
+                        raise HTTPException(status_code=500, detail="Failed to create 1024x1024 image")
+
+                    new_image.save(temp_path, format="PNG")
                     image_path = temp_path
-                    logger.info(f"Generated and resized image to 1024x1024 saved to: {temp_path} with prompt: {prompt}")
+                    logger.info(f"Generated and padded image to 1024x1024 saved to: {temp_path} with prompt: {prompt}")
 
         if image_path:
             return image_path, ""