Fix torch.compile optimization issues in image editing

- Add

@torch
.compiler.allow_in_graph decorators for PIL operations
- Create writable numpy arrays to avoid compilation errors
- Isolate torch.compile to core inference function only
- Preserve timing functionality outside compiled scope

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>

gradio_tabs/img_edit.py

@@ -37,16 +37,20 @@ labels_v = [
 ]
 
 
+@torch.compiler.allow_in_graph
 def load_image(img, size):
     img = Image.open(img).convert('RGB')
     w, h = img.size
     img = img.resize((size, size))
     img = np.asarray(img)
+    # Make a writable copy to avoid torch.compile issues
+    img = np.copy(img)
     img = np.transpose(img, (2, 0, 1))    # 3 x 256 x 256
 
     return img / 255.0, w, h
 
 
+@torch.compiler.allow_in_graph
 def img_preprocessing(img_path, size):
     img, w, h = load_image(img_path, size)  # [0, 1]
     img = torch.from_numpy(img).unsqueeze(0).float()  # [0, 1]
@@ -105,11 +109,15 @@ def img_postprocessing(img, w, h):
 
 def img_edit(gen, device):
 
+    @torch.compile
+    def compiled_inference(image_tensor, selected_s):
+        """Compiled version of just the model inference"""
+        return gen.edit_img(image_tensor, labels_v, selected_s)
+
     @spaces.GPU
     @torch.inference_mode()
-    @torch.compile
     def edit_img(image, *selected_s):
-        # Start timing
+        # Start timing (outside compiled function)
         start_time = time.time()
         print(f"[edit_img] Starting image editing...")
 
@@ -120,9 +128,9 @@ def img_edit(gen, device):
         preprocess_end = time.time()
         print(f"[edit_img] Preprocessing took: {(preprocess_end - preprocess_start) * 1000:.2f} ms")
 
-        # Model inference timing
+        # Model inference timing (compile only the core computation)
         inference_start = time.time()
-        edited_image_tensor = gen.edit_img(image_tensor, labels_v, selected_s)
+        edited_image_tensor = compiled_inference(image_tensor, selected_s)
         inference_end = time.time()
         print(f"[edit_img] Model inference took: {(inference_end - inference_start) * 1000:.2f} ms")