Use Gradio's postMessage zerogpu-headers for regen quota attribution

Instead of the unreliable x-ip-token relay, now use the same mechanism
Gradio's own JS client uses: postMessage("zerogpu-headers") to the HF
parent frame, which responds with x-zerogpu-token and x-zerogpu-uuid.
These are the actual headers that ZeroGPU uses for Pro quota attribution.
The HF token input in Settings is kept as a fallback.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -1976,44 +1976,51 @@ _GLOBAL_JS = """
   if (window._wf_global_listener) return;  // already registered
   window._wf_global_listener = true;
 
-  // Fetch a fresh x-ip-token via a Gradio endpoint where HF's proxy injects it.
-  // Returns a Promise<string> — resolves with the token (or '' on failure).
-  // Called immediately before each regen fetch so the token is always fresh
-  // (the JWT expires in ~170s, so caching across regen calls is unreliable).
-  function _fetchIpToken() {
+  // ── ZeroGPU quota attribution ──
+  // HF Spaces run inside an iframe on huggingface.co. Gradio's own JS client
+  // gets ZeroGPU auth headers (x-zerogpu-token, x-zerogpu-uuid) by sending a
+  // postMessage("zerogpu-headers") to the parent frame. The parent responds
+  // with a Map of headers that must be included on queue/join calls.
+  // We replicate this exact mechanism so our raw regen fetch() calls are
+  // attributed to the logged-in user's Pro quota.
+  function _fetchZerogpuHeaders() {
     return new Promise(function(resolve) {
-      var ipFnIndex = null;
-      var deps = window.gradio_config && window.gradio_config.dependencies;
-      if (deps) deps.forEach(function(d, i) { if (d.api_name === 'get_ip_token') ipFnIndex = i; });
-      if (ipFnIndex === null) { resolve(''); return; }
-      fetch('/gradio_api/queue/join', {
-        method: 'POST',
-        credentials: 'include',
-        headers: {'Content-Type': 'application/json'},
-        body: JSON.stringify({
-          data: [], fn_index: ipFnIndex,
-          session_hash: window.__gradio_session_hash__,
-          event_data: null, trigger_id: null
-        })
-      }).then(function(r) { return r.json(); }).then(function(j) {
-        if (!j.event_id) { resolve(''); return; }
-        var es = new EventSource('/gradio_api/queue/data?session_hash=' + window.__gradio_session_hash__);
-        var done = false;
-        es.onmessage = function(e) {
-          var msg; try { msg = JSON.parse(e.data); } catch(_) { return; }
-          if (msg.event_id !== j.event_id) return;
-          if (msg.msg === 'process_completed') {
-            es.close(); done = true;
-            var token = msg.output && msg.output.data && msg.output.data[0];
-            console.log('[zerogpu] fresh x-ip-token acquired, length:', token ? token.length : 0);
-            resolve(token || '');
+      // Check if we're in an HF iframe with zerogpu support
+      if (typeof window === 'undefined' || window.parent === window || !window.supports_zerogpu_headers) {
+        console.log('[zerogpu] not in HF iframe or no zerogpu support');
+        resolve({});
+        return;
+      }
+      // Determine origin — same logic as Gradio's client
+      var hostname = window.location.hostname;
+      var hfhubdev = 'dev.spaces.huggingface.tech';
+      var origin = hostname.includes('.dev.')
+        ? 'https://moon-' + hostname.split('.')[1] + '.' + hfhubdev
+        : 'https://huggingface.co';
+      // Use MessageChannel just like Gradio's post_message helper
+      var channel = new MessageChannel();
+      var done = false;
+      channel.port1.onmessage = function(ev) {
+        channel.port1.close();
+        done = true;
+        var headers = ev.data;
+        if (headers && typeof headers === 'object') {
+          // Convert Map to plain object if needed
+          var obj = {};
+          if (typeof headers.forEach === 'function') {
+            headers.forEach(function(v, k) { obj[k] = v; });
+          } else {
+            obj = headers;
           }
-          if (msg.msg === 'close_stream') { es.close(); if (!done) { done = true; resolve(''); } }
-        };
-        es.onerror = function() { es.close(); if (!done) { done = true; resolve(''); } };
-        // Timeout safety: if SSE hangs, don't block regen forever
-        setTimeout(function() { if (!done) { done = true; es.close(); resolve(''); } }, 5000);
-      }).catch(function() { resolve(''); });
+          console.log('[zerogpu] got headers from parent:', Object.keys(obj).join(', '));
+          resolve(obj);
+        } else {
+          resolve({});
+        }
+      };
+      window.parent.postMessage('zerogpu-headers', origin, [channel.port2]);
+      // Timeout: don't block regen if parent doesn't respond
+      setTimeout(function() { if (!done) { done = true; channel.port1.close(); resolve({}); } }, 3000);
     });
   }
 
@@ -2139,24 +2146,27 @@ _GLOBAL_JS = """
       console.warn('[fireRegen] fn_index not found for api_name:', apiName);
       return;
     }
-    // Build auth headers for the regen call.
-    // Prefer a user-supplied HF token (Authorization: Bearer) which is the
-    // documented way for ZeroGPU to attribute quota to a logged-in Pro account.
-    // Fall back to the x-ip-token relay approach if no HF token is provided.
-    var userHfToken = '';
-    var hfTokenEl = document.getElementById('hf_token_input');
-    if (hfTokenEl) { var inp = hfTokenEl.querySelector('input,textarea'); if (inp) userHfToken = (inp.value || '').trim(); }
-
-    var _doRegen = function(ipToken) {
+    // Get ZeroGPU auth headers from the HF parent frame (same mechanism
+    // Gradio's own JS client uses), then fire the regen queue/join call.
+    // Falls back to user-supplied HF token if zerogpu headers aren't available.
+    _fetchZerogpuHeaders().then(function(zerogpuHeaders) {
       var regenHeaders = {'Content-Type': 'application/json'};
-      if (userHfToken) {
-        regenHeaders['Authorization'] = 'Bearer ' + userHfToken;
-        console.log('[fireRegen] using HF token for Pro quota attribution');
-      } else if (ipToken) {
-        regenHeaders['x-ip-token'] = ipToken;
-        console.log('[fireRegen] using fresh x-ip-token, len:', ipToken.length);
+      var hasZerogpu = zerogpuHeaders && Object.keys(zerogpuHeaders).length > 0;
+      if (hasZerogpu) {
+        // Merge zerogpu headers (x-zerogpu-token, x-zerogpu-uuid)
+        for (var k in zerogpuHeaders) { regenHeaders[k] = zerogpuHeaders[k]; }
+        console.log('[fireRegen] using zerogpu headers from parent frame');
       } else {
-        console.warn('[fireRegen] no auth available, regen may use anonymous quota');
+        // Fallback: try user-supplied HF token from Settings
+        var userHfToken = '';
+        var hfTokenEl = document.getElementById('hf_token_input');
+        if (hfTokenEl) { var inp = hfTokenEl.querySelector('input,textarea'); if (inp) userHfToken = (inp.value || '').trim(); }
+        if (userHfToken) {
+          regenHeaders['Authorization'] = 'Bearer ' + userHfToken;
+          console.log('[fireRegen] using HF token fallback for quota attribution');
+        } else {
+          console.warn('[fireRegen] no zerogpu headers or HF token — may use anonymous quota');
+        }
       }
       fetch('/gradio_api/queue/join', {
         method: 'POST',
@@ -2179,13 +2189,7 @@ _GLOBAL_JS = """
         var sb = document.getElementById('wf_statusbar_' + slot_id);
         if (sb) { sb.style.color = '#e05252'; sb.textContent = '\u26a0 Request failed: ' + e.message; }
       });
-    };
-    // If user provided HF token, skip x-ip-token relay (token is sufficient)
-    if (userHfToken) {
-      _doRegen('');
-    } else {
-      _fetchIpToken().then(_doRegen);
-    }
+    });
   }
 
   // Subscribe to Gradio SSE stream for an event and apply outputs to DOM.
@@ -2638,20 +2642,10 @@ with gr.Blocks(title="Generate Audio for Video", css=_SLOT_CSS, js=_GLOBAL_JS) a
         api_name="xregen_hunyuan",
     )
 
-    # ---- x-ip-token relay ----
-    # ZeroGPU uses the x-ip-token header to attribute quota to the logged-in user.
-    # HF's proxy injects it on browser requests but NOT on raw JS fetch() calls.
-    # Solution: expose a lightweight Gradio endpoint that reads the token from the
-    # incoming request (where HF DOES inject it) and returns it. The JS _GLOBAL_JS
-    # calls _fetchIpToken() immediately before each regen fetch, getting a fresh
-    # token every time (the JWT expires ~170s so caching is unreliable).
-    _ip_token_tb = gr.Textbox(render=False)
-    def _get_ip_token(request: gr.Request):
-        return request.headers.get("x-ip-token", "")
-    gr.Button(render=False).click(
-        fn=_get_ip_token, inputs=[], outputs=[_ip_token_tb],
-        api_name="get_ip_token",
-    )
+    # NOTE: ZeroGPU quota attribution is now handled via postMessage("zerogpu-headers")
+    # to the HF parent frame — the same mechanism Gradio's own JS client uses.
+    # This replaced the old x-ip-token relay approach which was unreliable.
+    # The HF token input in Settings is kept as a fallback.
 
 print("[startup] app.py fully loaded — regen handlers registered, SSR disabled")
 demo.queue(max_size=10).launch(ssr_mode=False, height=900, allowed_paths=["/tmp"])