Add X-wait-for-model only after first call

huggingface · Jun 7, 2024 · 8ea3f1d · 8ea3f1d
1 parent bf57460
commit 8ea3f1d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
@@ -250,8 +250,6 @@ def post(
         headers = self.headers.copy()
         if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
             headers["Accept"] = "image/png"
-        if "X-wait-for-model" not in headers and url.startswith(INFERENCE_ENDPOINT):
-            headers["X-wait-for-model"] = "1"
 
         t0 = time.time()
         timeout = self.timeout
@@ -291,6 +289,8 @@ def post(
                     # ...or wait 1s and retry
                     logger.info(f"Waiting for model to be loaded on the server: {error}")
                     time.sleep(1)
+                    if "X-wait-for-model" not in headers and url.startswith(INFERENCE_ENDPOINT):
+                        headers["X-wait-for-model"] = "1"
                     if timeout is not None:
                         timeout = max(self.timeout - (time.time() - t0), 1)  # type: ignore
                     continue

diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -238,8 +238,6 @@ async def post(
         headers = self.headers.copy()
         if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
             headers["Accept"] = "image/png"
-        if "X-wait-for-model" not in headers and url.startswith(INFERENCE_ENDPOINT):
-            headers["X-wait-for-model"] = "1"
 
         t0 = time.time()
         timeout = self.timeout
@@ -286,6 +284,8 @@ async def post(
                             ) from error
                         # ...or wait 1s and retry
                         logger.info(f"Waiting for model to be loaded on the server: {error}")
+                        if "X-wait-for-model" not in headers and url.startswith(INFERENCE_ENDPOINT):
+                            headers["X-wait-for-model"] = "1"
                         time.sleep(1)
                         if timeout is not None:
                             timeout = max(self.timeout - (time.time() - t0), 1)  # type: ignore

diff --git a/utils/generate_async_inference_client.py b/utils/generate_async_inference_client.py
@@ -186,8 +186,6 @@ def _rename_to_AsyncInferenceClient(code: str) -> str:
         headers = self.headers.copy()
         if task in TASKS_EXPECTING_IMAGES and "Accept" not in headers:
             headers["Accept"] = "image/png"
-        if "X-wait-for-model" not in headers and url.startswith(INFERENCE_ENDPOINT):
-            headers["X-wait-for-model"] = "1"
 
         t0 = time.time()
         timeout = self.timeout
@@ -234,6 +232,8 @@ def _rename_to_AsyncInferenceClient(code: str) -> str:
                             ) from error
                         # ...or wait 1s and retry
                         logger.info(f"Waiting for model to be loaded on the server: {error}")
+                        if "X-wait-for-model" not in headers and url.startswith(INFERENCE_ENDPOINT):
+                            headers["X-wait-for-model"] = "1"
                         time.sleep(1)
                         if timeout is not None:
                             timeout = max(self.timeout - (time.time() - t0), 1)  # type: ignore