add warning when status=staging

huggingface · Feb 5, 2025 · 0568aa5 · 0568aa5
1 parent 181b6a7
commit 0568aa5
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 0 deletions.
diff --git a/src/huggingface_hub/inference/_providers/fal_ai.py b/src/huggingface_hub/inference/_providers/fal_ai.py
@@ -91,11 +91,14 @@ def _map_model(self, model: Optional[str]) -> str:
         provider_mapping = _get_provider_mapping(model, "fal-ai")
         if provider_mapping:
             provider_task = provider_mapping.get("task")
+            status = provider_mapping.get("status")
             if provider_task != self.task:
                 raise ValueError(
                     f"Model {model} is not supported for task {self.task} and provider fal.ai. "
                     f"Supported task: {provider_task}."
                 )
+            if status == "staging":
+                logger.warning(f"Model {model} is in staging mode for provider fal.ai and may not be ready for use.")
             return provider_mapping["providerId"]
         if self.task not in SUPPORTED_MODELS:
             raise ValueError(f"Task {self.task} not supported with fal.ai.")

diff --git a/src/huggingface_hub/inference/_providers/replicate.py b/src/huggingface_hub/inference/_providers/replicate.py
@@ -97,12 +97,17 @@ def _map_model(self, model: Optional[str]) -> str:
         provider_mapping = _get_provider_mapping(model, "replicate")
         if provider_mapping:
             provider_task = provider_mapping.get("task")
+            status = provider_mapping.get("status")
 
             if provider_task != self.task:
                 raise ValueError(
                     f"Model {model} is not supported for task {self.task} and provider Replicate. "
                     f"Supported task: {provider_task}."
                 )
+            if status == "staging":
+                logger.warning(
+                    f"Model {model} is in staging mode for provider Replicate and may not be ready for use."
+                )
             return provider_mapping["providerId"]
         if self.task not in SUPPORTED_MODELS:
             raise ValueError(f"Task {self.task} not supported with Replicate.")

diff --git a/src/huggingface_hub/inference/_providers/sambanova.py b/src/huggingface_hub/inference/_providers/sambanova.py
@@ -85,11 +85,16 @@ def _map_model(self, model: Optional[str]) -> str:
         provider_mapping = _get_provider_mapping(model, "sambanova")
         if provider_mapping:
             provider_task = provider_mapping.get("task")
+            status = provider_mapping.get("status")
             if provider_task != self.task:
                 raise ValueError(
                     f"Model {model} is not supported for task {self.task} and provider Sambanova. "
                     f"Supported task: {provider_task}."
                 )
+            if status == "staging":
+                logger.warning(
+                    f"Model {model} is in staging mode for provider Sambanova and may not be ready for use."
+                )
             return provider_mapping["providerId"]
 
         if self.task not in SUPPORTED_MODELS:

diff --git a/src/huggingface_hub/inference/_providers/together.py b/src/huggingface_hub/inference/_providers/together.py
@@ -126,11 +126,14 @@ def _map_model(self, model: Optional[str]) -> str:
 
         if provider_mapping:
             provider_task = provider_mapping.get("task")
+            status = provider_mapping.get("status")
             if provider_task != self.task:
                 raise ValueError(
                     f"Model {model} is not supported for task {self.task} and provider Together. "
                     f"Supported task: {provider_task}."
                 )
+            if status == "staging":
+                logger.warning(f"Model {model} is in staging mode for provider Together and may not be ready for use.")
             return provider_mapping["providerId"]
 
         if self.task not in SUPPORTED_MODELS: