[ZeroGPU] Handshake-based postMessage (#10270)

* handshake-based ZeroGPU postMessage * HF Hub dev mode * window is always defined here * pnpm format:write * add changeset * change --------- Co-authored-by: pngwn <[email protected]> Co-authored-by: gradio-pr-bot <[email protected]> Co-authored-by: Abubakar Abid <[email protected]>
gradio-app · Jan 10, 2025 · bb11a2a · bb11a2a
1 parent be40307
commit bb11a2a
Show file tree

Hide file tree

Showing 8 changed files with 27 additions and 12 deletions.
diff --git a/.changeset/metal-cows-remain.md b/.changeset/metal-cows-remain.md
@@ -0,0 +1,7 @@
+---
+"@gradio/client": minor
+"@self/app": minor
+"gradio": minor
+---
+
+feat:[ZeroGPU] Handshake-based postMessage
diff --git a/client/js/src/globals.d.ts b/client/js/src/globals.d.ts
@@ -7,5 +7,6 @@ declare global {
 		gradio_api_info: ApiInfo<ApiData> | { api: ApiInfo<ApiData> };
 		__is_colab__: boolean;
 		__gradio_space__: string | null;
+		supports_zerogpu_headers?: boolean;
 	}
 }
diff --git a/client/js/src/types.ts b/client/js/src/types.ts
@@ -254,7 +254,6 @@ export interface Dependency {
 	trigger_mode: "once" | "multiple" | "always_last";
 	final_event: Payload | null;
 	show_api: boolean;
-	zerogpu?: boolean;
 	rendered_in: number | null;
 	connection: "stream" | "sse";
 	time_limit: number;

diff --git a/client/js/src/utils/submit.ts b/client/js/src/utils/submit.ts
@@ -566,15 +566,14 @@ export function submit(
 						? `https://moon-${hostname.split(".")[1]}.${hfhubdev}`
 						: `https://huggingface.co`;
 
-					const is_iframe =
+					const is_zerogpu_iframe =
 						typeof window !== "undefined" &&
 						typeof document !== "undefined" &&
-						window.parent != window;
-					const is_zerogpu_space = dependency.zerogpu && config.space_id;
-					const zerogpu_auth_promise =
-						is_iframe && is_zerogpu_space
-							? post_message<Headers>("zerogpu-headers", origin)
-							: Promise.resolve(null);
+						window.parent != window &&
+						window.supports_zerogpu_headers;
+					const zerogpu_auth_promise = is_zerogpu_iframe
+						? post_message<Map<string, string>>("zerogpu-headers", origin)
+						: Promise.resolve(null);
 					const post_data_promise = zerogpu_auth_promise.then((headers) => {
 						return post_data(
 							`${config.root}${api_prefix}/${SSE_DATA_URL}?${url_params}`,

diff --git a/globals.d.ts b/globals.d.ts
@@ -12,6 +12,7 @@ declare global {
 		parentIFrame?: {
 			scrollTo: (x: number, y: number) => void;
 		};
+		supports_zerogpu_headers?: boolean;
 	}
 }
 

diff --git a/gradio/blocks.py b/gradio/blocks.py
@@ -547,7 +547,6 @@ def __init__(
         self.queue = False if fn is None else queue
         self.scroll_to_output = False if utils.get_space() else scroll_to_output
         self.show_api = show_api
-        self.zero_gpu = hasattr(self.fn, "zerogpu")
         self.types_generator = inspect.isgeneratorfunction(
             self.fn
         ) or inspect.isasyncgenfunction(self.fn)
@@ -608,7 +607,6 @@ def get_config(self):
             "trigger_only_on_success": self.trigger_only_on_success,
             "trigger_mode": self.trigger_mode,
             "show_api": self.show_api,
-            "zerogpu": self.zero_gpu,
             "rendered_in": self.rendered_in._id if self.rendered_in else None,
             "connection": self.connection,
             "time_limit": self.time_limit,

diff --git a/gradio/chat_interface.py b/gradio/chat_interface.py
@@ -500,8 +500,6 @@ def _setup_events(self) -> None:
 
         submit_triggers = [self.textbox.submit, self.chatbot.retry]
         submit_fn = self._stream_fn if self.is_generator else self._submit_fn
-        if hasattr(self.fn, "zerogpu"):
-            submit_fn.__func__.zerogpu = self.fn.zerogpu  # type: ignore
 
         synchronize_chat_state_kwargs = {
             "fn": lambda x: x,

diff --git a/js/app/src/routes/[...catchall]/+page.svelte b/js/app/src/routes/[...catchall]/+page.svelte
@@ -215,6 +215,18 @@
 		css_ready = true;
 		window.__is_colab__ = config.is_colab;
 
+		const supports_zerogpu_headers = "supports-zerogpu-headers";
+		window.addEventListener("message", (event) => {
+			if (event.data === supports_zerogpu_headers) {
+				window.supports_zerogpu_headers = true;
+			}
+		});
+		const hostname = window.location.hostname;
+		const origin = hostname.includes(".dev.")
+			? `https://moon-${hostname.split(".")[1]}.dev.spaces.huggingface.tech`
+			: `https://huggingface.co`;
+		window.parent.postMessage(supports_zerogpu_headers, origin);
+
 		dispatch("loaded");
 
 		if (config.dev_mode) {
-Original file line number
+Diff line change
@@ Expand Up / @@ -12,6 +12,7 @@ declare global { @@
     		parentIFrame?: {
     			scrollTo: (x: number, y: number) => void;
     		};
+    		supports_zerogpu_headers?: boolean;
     	}
     }
@@ Expand Down @@