Update app.py

app.py

@@ -1,14 +1,7 @@
 import subprocess
 import shlex
-# Install the custom component if needed
-subprocess.run(
-    shlex.split(
-        "pip install ./gradio_magicquillv2-0.0.1-py3-none-any.whl"
-    )
-)
 import sys
 import os
-import gradio as gr
 import tempfile
 import numpy as np
 import io
@@ -16,13 +9,21 @@ import base64
 import json
 import uvicorn
 import torch
+from PIL import Image
+
+# Install the custom component if needed
+subprocess.run(
+    shlex.split(
+        "pip install ./gradio_magicquillv2-0.0.1-py3-none-any.whl"
+    )
+)
+
+import gradio as gr
 from fastapi import FastAPI, Request
 from fastapi.concurrency import run_in_threadpool
 from fastapi.middleware.cors import CORSMiddleware
 from gradio_client import Client, handle_file
 from gradio_magicquillv2 import MagicQuillV2
-from PIL import Image
-
 
 from util import (
     read_base64_image as read_base64_image_utils,
@@ -32,18 +33,13 @@ from util import (
 
 # --- Configuration ---
 # Set this to the URL of your backend Space (running app_backend.py)
-# Example: "https://huggingface.co/spaces/username/backend-space"
 BACKEND_URL = "LiuZichen/MagicQuillV2"
 SAM_URL = "LiuZichen/MagicQuillHelper"
 
-print(f"Connecting to backend at: {BACKEND_URL}")
-
-try:
-    backend_client = Client(BACKEND_URL)
-except Exception as e:
-    print(f"Failed to connect to backend: {e}")
-    backend_client = None
+print(f"Target Backend URL: {BACKEND_URL}")
 
+# We still initialize SAM client globally as it might not require ZeroGPU quotas
+# or is a helper CPU space.
 print(f"Connecting to SAM client at: {SAM_URL}")
 try:
     sam_client = Client(SAM_URL)
@@ -53,7 +49,32 @@ except Exception as e:
 
 # --- Helper Functions ---
 
-def generate_image_handler(x, negative_prompt, fine_edge, fix_perspective, grow_size, edge_strength, color_strength, local_strength, seed, steps, cfg):
+def get_zerogpu_headers(request_headers):
+    """
+    Extracts ZeroGPU specific headers from the incoming request headers.
+    These are required to forward the user's quota token to the backend.
+    """
+    headers = {}
+    if request_headers:
+        # These are the headers HF injects for ZeroGPU authentication and tracking
+        target_headers = [
+            "x-ip-token", 
+            "x-zerogpu-token", 
+            "x-zerogpu-uuid", 
+            "authorization", 
+            "cookie"
+        ]
+        for h in target_headers:
+            val = request_headers.get(h)
+            if val:
+                headers[h] = val
+    return headers
+
+def generate_image_handler(x, negative_prompt, fine_edge, fix_perspective, grow_size, edge_strength, color_strength, local_strength, seed, steps, cfg, request: gr.Request):
+    """
+    Handler for the Gradio UI. 
+    Note the 'request: gr.Request' argument - Gradio automatically injects this.
+    """
     merged_image = x['from_frontend']['img']
     total_mask = x['from_frontend']['total_mask']
     original_image = x['from_frontend']['original_image']
@@ -64,15 +85,17 @@ def generate_image_handler(x, negative_prompt, fine_edge, fix_perspective, grow_
     add_prop_image = x['from_frontend']['add_prop_image']
     positive_prompt = x['from_backend']['prompt']
 
-    if backend_client is None:
-        print("Backend client not initialized")
-        x["from_backend"]["generated_image"] = None
-        return x
+    # 1. Extract headers from the current user's request
+    forward_headers = get_zerogpu_headers(request.headers)
 
     try:
+        # 2. Instantiate a client specifically for this request with the forwarded headers.
+        # This ensures the backend sees the 'x-zerogpu-token' of the user, not the server.
+        # gradio_client caches schemas, so re-init is relatively cheap but necessary for headers.
+        client = Client(BACKEND_URL, headers=forward_headers)
+
         # Call the backend API
-        # The order of arguments must match app_backend.py input list
-        res_base64 = backend_client.predict(
+        res_base64 = client.predict(
             merged_image,       # merged_image
             total_mask,         # total_mask
             original_image,     # original_image
@@ -107,8 +130,8 @@ with gr.Blocks(title="MagicQuill V2") as demo:
     with gr.Row(elem_classes="row"):
         text = gr.Markdown(
         """
-        # Welcome to MagicQuill V2! Give us a [GitHub star](https://github.com/zliucz/magicquillv2) if you are interested. 
-        Click the [link](https://magicquill.art/v2) to view our demo and tutorial. The paper is on [ArXiv](https://arxiv.org/abs/2512.03046) now.
+        # Welcome to MagicQuill V2! Give us a [GitHub star] if you are interested. 
+        Click the [link] to view our demo and tutorial. The paper is on [ArXiv] now.
         """)
         
     with gr.Row():
@@ -132,6 +155,8 @@ with gr.Blocks(title="MagicQuill V2") as demo:
 
     btn.click(
         generate_image_handler, 
+        # Note: We do NOT need to explicitly add 'request' to inputs here.
+        # Gradio handles type hinting for gr.Request automatically.
         inputs=[ms, negative_prompt, fine_edge, fix_perspective, grow_size, edge_strength, color_strength, local_strength, seed, steps, cfg], 
         outputs=ms
     )
@@ -150,18 +175,20 @@ app.add_middleware(
 def get_root_url(request: Request, route_path: str, root_path: str | None):
     return root_path
 gr.route_utils.get_root_url = get_root_url
-# gr.mount_gradio_app(app, demo, path="/demo", root_path="/demo")
 
 @app.post("/magic_quill/generate_image")
 async def generate_image(request: Request):
     data = await request.json()
     
-    if backend_client is None:
-        return {'error': 'Backend client not connected'}
+    # 1. Extract headers from the FastAPI request object
+    forward_headers = get_zerogpu_headers(request.headers)
 
     try:
+        # 2. Instantiate client with headers
+        client = Client(BACKEND_URL, headers=forward_headers)
+
         res = await run_in_threadpool(
-            backend_client.predict,
+            client.predict,
             data["merged_image"],
             data["total_mask"],
             data["original_image"],
@@ -214,7 +241,7 @@ async def segmentation(request: Request):
     if sam_client is None:
         return {"error": "sam client not initialized"}
 
-    # Process coordinates and bboxes (copied from original app.py)
+    # Process coordinates and bboxes
     pos_coordinates = None
     if coordinates_positive and len(coordinates_positive) > 0:
         pos_coordinates = []
@@ -345,8 +372,6 @@ async def segmentation(request: Request):
         else:
             seg_bbox = {'startX': 0, 'startY': 0, 'endX': 0, 'endY': 0}
 
-        print(seg_bbox)
-        
         buffered = io.BytesIO()
         res_pil.save(buffered, format="PNG")
         image_base64_res = base64.b64encode(buffered.getvalue()).decode("utf-8")