Update app.py

app.py

@@ -9,6 +9,7 @@ from typing import Iterable
 from gradio.themes import Soft
 from gradio.themes.utils import colors, fonts, sizes
 
+# --- Theme Configuration ---
 colors.orange_red = colors.Color(
     name="orange_red",
     c50="#FFF0E5",
@@ -77,6 +78,7 @@ class OrangeRedTheme(Soft):
 
 orange_red_theme = OrangeRedTheme()
 
+# --- Device Setup ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 dtype = torch.bfloat16
 
@@ -85,6 +87,7 @@ from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
 from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
 from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
+# --- Model Loading ---
 print("Loading Qwen Image Edit Pipeline...")
 pipe = QwenImageEditPlusPipeline.from_pretrained(
     "Qwen/Qwen-Image-Edit-2509",
@@ -214,16 +217,30 @@ def infer(
     
     width, height = update_dimensions_on_upload(img1_pil)
 
-    result = pipe(
-        image=[img1_pil, img2_pil],
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        height=height,
-        width=width,
-        num_inference_steps=steps,
-        generator=generator,
-        true_cfg_scale=guidance_scale,
-    ).images[0]
+    # --- Fix: Explicit Memory Management ---
+    # Clear cache before starting the heavy inference process
+    torch.cuda.empty_cache()
+
+    try:
+        # Use no_grad to prevent gradient calculation and save memory
+        with torch.no_grad():
+            result = pipe(
+                image=[img1_pil, img2_pil],
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                height=height,
+                width=width,
+                num_inference_steps=steps,
+                generator=generator,
+                true_cfg_scale=guidance_scale,
+            ).images[0]
+    except Exception as e:
+        # If an error occurs, ensure we still clear cache before raising
+        torch.cuda.empty_cache()
+        raise e
+    
+    # Clear cache after inference is done
+    torch.cuda.empty_cache()
 
     return result, seed
 
@@ -231,6 +248,10 @@ def infer(
 def infer_example(image_1, image_2, prompt, lora_adapter):
     if image_1 is None or image_2 is None:
         return None, 0
+    
+    # Optional: Clear cache before example inference as well
+    torch.cuda.empty_cache()
+    
     result, seed = infer(
         image_1.convert("RGB"), 
         image_2.convert("RGB"),