Update app.py

app.py

@@ -71,27 +71,48 @@ zerogpu_fix_js = """
 
         async function initZeroGPU() {
             try {
+                // Initial check for headers
                 const headers = await post_message("zerogpu-headers");
+                
+                // Helper to update headers object
+                const updateHeaders = (options, headers) => {
+                    if (!options) options = {};
+                    if (!options.headers) options.headers = {};
+                    
+                    const headerMap = headers instanceof Map ? headers : new Map(Object.entries(headers));
+                    
+                    headerMap.forEach((value, key) => {
+                        if (options.headers instanceof Headers) {
+                            options.headers.append(key, value);
+                        } else {
+                            options.headers[key] = value;
+                        }
+                    });
+                    return options;
+                };
+
                 if (headers) {
-                    console.log("ZeroGPU Fix: Received headers", headers);
+                    console.log("ZeroGPU Fix: Received initial headers", headers);
                     
                     const originalFetch = window.fetch;
                     window.fetch = async function(url, options) {
-                        // Only inject headers for prediction requests to avoid rate limiting on heartbeats/static assets
-                        // Gradio requests typically go to /queue/join, /api/predict, etc.
+                        // Only inject headers for prediction requests
                         if (typeof url === 'string' && (url.includes('/queue/join') || url.includes('/call/') || url.includes('/api/'))) {
-                             if (!options) options = {};
-                             if (!options.headers) options.headers = {};
-                             
-                             const headerMap = headers instanceof Map ? headers : new Map(Object.entries(headers));
-                             
-                             headerMap.forEach((value, key) => {
-                                  if (options.headers instanceof Headers) {
-                                      options.headers.append(key, value);
-                                  } else {
-                                      options.headers[key] = value;
-                                  }
-                             });
+                            
+                            // Re-fetch headers before every important request to ensure token freshness
+                            // This is critical because tokens might expire or change
+                            try {
+                                const currentHeaders = await post_message("zerogpu-headers");
+                                if (currentHeaders) {
+                                    options = updateHeaders(options, currentHeaders);
+                                } else {
+                                     // Fallback to initial headers if fetch fails
+                                     options = updateHeaders(options, headers);
+                                }
+                            } catch (e) {
+                                // Fallback to initial headers
+                                options = updateHeaders(options, headers);
+                            }
                         }
                         
                         return originalFetch(url, options);
@@ -145,6 +166,13 @@ def generate_image_handler(x, negative_prompt, fine_edge, fix_perspective, grow_
 
     # 1. Extract headers from the current user's request
     forward_headers = get_zerogpu_headers(request.headers)
+    
+    # 打印调试信息，确认是否收到了 x-zerogpu-token
+    print(f"Debug: Received headers keys: {list(request.headers.keys())}")
+    if "x-zerogpu-token" in forward_headers:
+        print("Debug: x-zerogpu-token found in request headers.")
+    else:
+        print("Debug: x-zerogpu-token NOT found in request headers.")
 
     try:
         # 2. Instantiate a client specifically for this request with the forwarded headers.