Update app.py

app.py

@@ -49,84 +49,94 @@ except Exception as e:
 
 # --- Helper Functions ---
 
-# JS script to fix ZeroGPU headers propagation
-zerogpu_fix_js = """
-<script>
-    (function() {
-        const is_zerogpu = typeof window !== "undefined" && window.parent !== window;
-        if (!is_zerogpu) return;
-
-        console.log("ZeroGPU Fix: Detected iframe, attempting to fetch headers.");
-
-        function post_message(message) {
-            return new Promise((resolve, reject) => {
-                const channel = new MessageChannel();
-                channel.port1.onmessage = (event) => {
-                    channel.port1.close();
-                    resolve(event.data);
-                };
-                window.parent.postMessage(message, "*", [channel.port2]);
-            });
-        }
-
-        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;
-                        }
+        // JS script to fix ZeroGPU headers propagation
+        zerogpu_fix_js = """
+        <script>
+            (function() {
+                const is_zerogpu = typeof window !== "undefined" && window.parent !== window;
+                if (!is_zerogpu) {
+                    console.log("ZeroGPU Fix: Not in iframe, skipping.");
+                    return;
+                }
+        
+                console.log("ZeroGPU Fix: Detected iframe, initializing...");
+        
+                function post_message(message) {
+                    return new Promise((resolve, reject) => {
+                        const channel = new MessageChannel();
+                        channel.port1.onmessage = (event) => {
+                            channel.port1.close();
+                            resolve(event.data);
+                        };
+                        window.parent.postMessage(message, "*", [channel.port2]);
                     });
-                    return options;
-                };
-
-                if (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
-                        if (typeof url === 'string' && (url.includes('/queue/join') || url.includes('/call/') || url.includes('/api/'))) {
+                }
+        
+                async function initZeroGPU() {
+                    try {
+                        console.log("ZeroGPU Fix: Requesting initial 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 = {};
                             
-                            // 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);
+                            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 {
-                                     // Fallback to initial headers if fetch fails
-                                     options = updateHeaders(options, headers);
+                                    options.headers[key] = value;
                                 }
-                            } catch (e) {
-                                // Fallback to initial headers
-                                options = updateHeaders(options, headers);
-                            }
+                            });
+                            return options;
+                        };
+        
+                        if (headers) {
+                            console.log("ZeroGPU Fix: Received initial headers", headers);
+                            
+                            const originalFetch = window.fetch;
+                            window.fetch = async function(url, options) {
+                                // Inject headers for ALL requests to be safe, especially queue/join
+                                // We filter out only obviously external or static asset requests if needed, 
+                                // but for debugging, let's be aggressive.
+                                if (typeof url === 'string' && (url.includes('/queue/') || url.includes('/api/') || url.includes('run/predict'))) {
+                                    console.log("ZeroGPU Fix: Intercepting request to", url);
+                                    
+                                    try {
+                                        // Attempt to refresh headers
+                                        const currentHeaders = await post_message("zerogpu-headers");
+                                        if (currentHeaders) {
+                                            options = updateHeaders(options, currentHeaders);
+                                            console.log("ZeroGPU Fix: Injected FRESH headers");
+                                        } else {
+                                            options = updateHeaders(options, headers);
+                                            console.log("ZeroGPU Fix: Injected CACHED headers (refresh failed)");
+                                        }
+                                    } catch (e) {
+                                        options = updateHeaders(options, headers);
+                                        console.log("ZeroGPU Fix: Injected CACHED headers (error)");
+                                    }
+                                }
+                                
+                                return originalFetch(url, options);
+                            };
+                            console.log("ZeroGPU Fix: Fetch interceptor installed.");
+                        } else {
+                            console.warn("ZeroGPU Fix: No headers received from parent.");
                         }
-                        
-                        return originalFetch(url, options);
-                    };
+                    } catch (e) {
+                        console.error("ZeroGPU Fix: Error initializing", e);
+                    }
                 }
-            } catch (e) {
-                console.error("ZeroGPU Fix: Error fetching headers", e);
-            }
-        }
-        
-        initZeroGPU();
-    })();
-</script>
-"""
+                
+                // Run immediately
+                initZeroGPU();
+            })();
+        </script>
+        """
 
 def get_zerogpu_headers(request_headers):
     """