adding support for nscale inference provider

Author: nbarr07

src/huggingface_hub/inference/_client.py

@@ -133,7 +133,7 @@ class InferenceClient:
             path will be appended to the base URL (see the [TGI Messages API](https://huggingface.co/docs/text-generation-inference/en/messages_api)
             documentation for details). When passing a URL as `model`, the client will not append any suffix path to it.
         provider (`str`, *optional*):
-            Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"fireworks-ai"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"openai"`, `"replicate"`, "sambanova"` or `"together"`.
+            Name of the provider to use for inference. Can be `"black-forest-labs"`, `"cerebras"`, `"cohere"`, `"fal-ai"`, `"fireworks-ai"`, `"hf-inference"`, `"hyperbolic"`, `"nebius"`, `"novita"`, `"nscale"`, `"openai"`, `"replicate"`, "sambanova"` or `"together"`.
             Defaults to "auto" i.e. the first of the providers available for the model, sorted by the user's order in https://hf.co/settings/inference-providers.
             If model is a URL or `base_url` is passed, then `provider` is not used.
         token (`str`, *optional*):

src/huggingface_hub/inference/_providers/__init__.py

@@ -22,6 +22,7 @@
 from .hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
 from .nebius import NebiusConversationalTask, NebiusTextGenerationTask, NebiusTextToImageTask
 from .novita import NovitaConversationalTask, NovitaTextGenerationTask, NovitaTextToVideoTask
+from .nscale import NscaleChatCompletion, NscaleTextToImageTask
 from .openai import OpenAIConversationalTask
 from .replicate import ReplicateTask, ReplicateTextToImageTask, ReplicateTextToSpeechTask
 from .sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
@@ -41,6 +42,7 @@
     "hyperbolic",
     "nebius",
     "novita",
+    "nscale",
     "openai",
     "replicate",
     "sambanova",
@@ -111,6 +113,10 @@
         "conversational": NovitaConversationalTask(),
         "text-to-video": NovitaTextToVideoTask(),
     },
+    "nscale": {
+        "conversational": NscaleChatCompletion(),
+        "text-to-image": NscaleTextToImageTask(),
+    },
     "openai": {
         "conversational": OpenAIConversationalTask(),
     },

src/huggingface_hub/inference/_providers/_common.py

@@ -27,6 +27,7 @@
     "hf-inference": {},
     "hyperbolic": {},
     "nebius": {},
+    "nscale": {},
     "replicate": {},
     "sambanova": {},
     "together": {},

src/huggingface_hub/inference/_providers/nscale.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, Optional, Union
+import base64
+
+from huggingface_hub.inference._common import RequestParameters, _as_dict
+from ._common import (
+    TaskProviderHelper,
+    BaseConversationalTask,
+    filter_none,
+)
+
+class NscaleTask(TaskProviderHelper):
+    def __init__(self, task: str):
+        super().__init__(provider="nscale", base_url="https://inference.api.nscale.com", task=task)
+
+    def _prepare_route(self, mapped_model: str, api_key: str) -> str:
+        if self.task == "text-to-image":
+            return "/v1/images/generations"
+        elif self.task == "conversational":
+            return "/v1/chat/completions"
+        raise ValueError(f"Unsupported task '{self.task}' for Nscale API.")
+
+class NscaleChatCompletion(BaseConversationalTask):
+    def __init__(self):
+        super().__init__(provider="nscale", base_url="https://inference.api.nscale.com")
+
+class NscaleTextToImageTask(NscaleTask):
+    def __init__(self):
+        super().__init__("text-to-image")
+        
+    def _prepare_payload_as_dict(self, inputs: Any, parameters: Dict, mapped_model: str) -> Optional[Dict]:
+        # Combine all parameters except inputs and parameters
+        parameters = filter_none(parameters)
+        if "width" in parameters and "height" in parameters:
+            parameters["size"] = f"{parameters.pop('width')}x{parameters.pop('height')}"
+        if "num_inference_steps" in parameters:
+            parameters.pop("num_inference_steps")
+        if "cfg_scale" in parameters:
+            parameters.pop("cfg_scale")
+        payload = {
+            "response_format": "b64_json",
+            "prompt": inputs,
+            "model": mapped_model,
+            **parameters,
+        }
+        return payload
+        
+    def get_response(self, response: Union[bytes, Dict], request_params: Optional[RequestParameters] = None) -> Any:
+        response_dict = _as_dict(response)
+        return base64.b64decode(response_dict["data"][0]["b64_json"])
+        

tests/test_inference_providers.py

@@ -33,6 +33,7 @@
 from huggingface_hub.inference._providers.hyperbolic import HyperbolicTextGenerationTask, HyperbolicTextToImageTask
 from huggingface_hub.inference._providers.nebius import NebiusTextToImageTask
 from huggingface_hub.inference._providers.novita import NovitaConversationalTask, NovitaTextGenerationTask
+from huggingface_hub.inference._providers.nscale import NscaleChatCompletion, NscaleTextToImageTask
 from huggingface_hub.inference._providers.openai import OpenAIConversationalTask
 from huggingface_hub.inference._providers.replicate import ReplicateTask, ReplicateTextToSpeechTask
 from huggingface_hub.inference._providers.sambanova import SambanovaConversationalTask, SambanovaFeatureExtractionTask
@@ -829,6 +830,56 @@ def test_prepare_url_conversational(self):
         url = helper._prepare_url("novita_token", "username/repo_name")
         assert url == "https://api.novita.ai/v3/openai/chat/completions"
 
+class TestNscaleProvider:
+    def test_prepare_route_text_to_image(self):
+        helper = NscaleTextToImageTask()
+        assert helper._prepare_route("model_name", "api_key") == "/v1/images/generations"
+
+    def test_prepare_route_chat_completion(self):
+        helper = NscaleChatCompletion()
+        assert helper._prepare_route("model_name", "api_key") == "/v1/chat/completions"    
+    
+    def test_prepare_payload_with_size_conversion(self):
+        helper = NscaleTextToImageTask()
+        payload = helper._prepare_payload_as_dict(
+            "a beautiful landscape",
+            {
+                "width": 512,
+                "height": 512,
+            },
+            "stabilityai/stable-diffusion-xl-base-1.0",
+        )
+        assert payload == {
+            "prompt": "a beautiful landscape",
+            "size": "512x512",
+            "response_format": "b64_json",
+            "model": "stabilityai/stable-diffusion-xl-base-1.0",
+        }
+        
+    def test_prepare_payload_as_dict(self):
+        helper = NscaleTextToImageTask()
+        payload = helper._prepare_payload_as_dict(
+            "a beautiful landscape",
+            {
+                "width": 1024,
+                "height": 768,
+                "cfg_scale": 7.5,
+                "num_inference_steps": 50,
+            },
+            "stabilityai/stable-diffusion-xl-base-1.0",
+        )
+        assert "width" not in payload
+        assert "height" not in payload
+        assert "num_inference_steps" not in payload
+        assert "cfg_scale" not in payload
+        assert payload["size"] == "1024x768"
+        assert payload["model"] == "stabilityai/stable-diffusion-xl-base-1.0"
+        
+    def test_text_to_image_get_response(self):
+        helper = NscaleTextToImageTask()
+        response = helper.get_response({"data": [{"b64_json": base64.b64encode(b"image_bytes").decode()}]})
+        assert response == b"image_bytes"
+
 
 class TestOpenAIProvider:
     def test_prepare_url(self):