Update app.py

app.py

@@ -105,26 +105,35 @@ pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight",
 pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 MAX_SEED = np.iinfo(np.int32).max
 
-# --- Helper Function for Aspect Ratio ---
+# --- Helper Function for Aspect Ratio (Corrected) ---
+@torch.no_grad()
 def update_dimensions_on_upload(image):
+    # *** FIX: This function now correctly preserves aspect ratio for all image sizes. ***
     if image is None:
-        return 1024, 1024
-    
+        return 1024, 1024  # Default for no image
+
     original_width, original_height = image.size
-    
-    # Cap max dimension to 1024 while preserving aspect ratio
-    if original_width > original_height:
-        new_width = 1024
-        new_height = int(1024 * original_height / original_width)
+    max_dim = 1024
+
+    if original_width > max_dim or original_height > max_dim:
+        # If the image is larger than the max dimension, scale it down
+        if original_width > original_height:
+            new_width = max_dim
+            new_height = int(max_dim * original_height / original_width)
+        else:
+            new_height = max_dim
+            new_width = int(max_dim * original_width / original_height)
     else:
-        new_height = 1024
-        new_width = int(1024 * original_width / original_height)
-        
-    # Ensure dimensions are multiples of 8 for model compatibility
-    new_width = (new_width // 8) * 8
-    new_height = (new_height // 8) * 8
-    
-    return new_width, new_height
+        # If the image is smaller, use its original dimensions
+        new_width = original_width
+        new_height = original_height
+
+    # Ensure final dimensions are multiples of 8 for model compatibility
+    final_width = (new_width // 8) * 8
+    final_height = (new_height // 8) * 8
+
+    return final_width, final_height
+
 
 # --- Main Inference Function ---
 @spaces.GPU
@@ -177,17 +186,12 @@ def infer(
 # --- Wrapper for Examples ---
 @spaces.GPU
 def infer_example(input_image, prompt, lora_adapter):
-    # *** FIX: The input is already a PIL Image object from Gradio, not a path. ***
-    # We no longer need Image.open() here.
     input_pil = input_image.convert("RGB")
-    
-    # Calculate aspect ratio for the example image
+    # Calculate correct aspect ratio for the example image using the corrected function
     width, height = update_dimensions_on_upload(input_pil)
-    
-    # Set reasonable default values for example inference to get good results
+    # Set reasonable default values for example inference
     guidance_scale = 1.0
     steps = 4
-    
     # Call the main infer function
     result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps, width, height)
     return result, seed