nexaai 0.0.8__tar.gz → 0.0.8.2__tar.gz

{nexaai-0.0.8 → nexaai-0.0.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nexaai
-Version: 0.0.8
+Version: 0.0.8.2
 Summary: Nexa AI SDK
 Author-Email: Nexa AI <octopus@nexa4ai.com>
 License: MIT
@@ -28,10 +28,9 @@ Requires-Dist: pillow
 Requires-Dist: prompt_toolkit
 Requires-Dist: tqdm
 Requires-Dist: tabulate
-Requires-Dist: streamlit
+Requires-Dist: streamlit>=1.37.1
 Requires-Dist: streamlit-audiorec
 Requires-Dist: python-multipart
-Requires-Dist: streamlit-audiorec
 Requires-Dist: cmake
 Requires-Dist: librosa; extra == "onnx"
 Requires-Dist: optimum[onnxruntime]; extra == "onnx"
@@ -70,7 +69,7 @@ Description-Content-Type: text/markdown
 
 </div>
 
-Nexa SDK is a comprehensive toolkit for supporting **ONNX** and **GGML** models. It supports text generation, image generation, vision-language models (VLM), and text-to-speech (TTS) capabilities. Additionally, it offers an OpenAI-compatible API server with JSON schema mode for function calling and streaming support, and a user-friendly Streamlit UI.
+Nexa SDK is a comprehensive toolkit for supporting **ONNX** and **GGML** models. It supports text generation, image generation, vision-language models (VLM), and text-to-speech (TTS) capabilities. Additionally, it offers an OpenAI-compatible API server with JSON schema mode for function calling and streaming support, and a user-friendly Streamlit UI. Users can run Nexa SDK in any device with Python environment, and GPU acceleration is supported.
 
 ## Features
 
@@ -111,6 +110,17 @@ Below is our differentiation from other similar tools:
 
 We have released pre-built wheels for various Python versions, platforms, and backends for convenient installation on our [index page](https://nexaai.github.io/nexa-sdk/whl/).
 
+How to clone this repo
+```bash
+git clone --recursive https://github.com/NexaAI/nexa-sdk
+```
+
+If you forget to use `--recursive`, you can use below command to add submodule
+```bash
+git submodule update --init --recursive
+```
+
+
 #### CPU
 
 ```bash
@@ -142,14 +152,31 @@ CMAKE_ARGS="-DGGML_METAL=ON -DSD_METAL=ON" pip install nexaai --prefer-binary --
 
 #### GPU (CUDA)
 
-For the GPU version supporting **CUDA (Linux/Windows)**:
+For **Linux**:
 
 ```bash
 CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON" pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
 ```
 
+For **Windows PowerShell**:
+
+```bash
+$env:CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON"; pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
+```
+
+For **Windows Command Prompt**:
+
+```bash
+set CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON" & pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
+```
+
+For **Windows Git Bash**:
+
+```bash
+CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON" pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
+```
 > [!NOTE]
-> The CUDA wheels are built with CUDA 12.4, but should be compatible with all CUDA 12.X
+> If you want to use ONNX model, just replace `pip install nexaai` with `pip install nexaai[onnx]` in above commands
 
 
 <details>
@@ -259,4 +286,4 @@ For detailed information on server setup, API endpoints, and usage examples, ple
 We would like to thank the following projects:
 - [llama.cpp](https://github.com/ggerganov/llama.cpp)
 - [stable-diffusion.cpp](https://github.com/leejet/stable-diffusion.cpp)
-- [optimum](https://github.com/huggingface/optimum)
+- [optimum](https://github.com/huggingface/optimum)

{nexaai-0.0.8 → nexaai-0.0.8.2}/README.md

@@ -21,7 +21,7 @@
 
 </div>
 
-Nexa SDK is a comprehensive toolkit for supporting **ONNX** and **GGML** models. It supports text generation, image generation, vision-language models (VLM), and text-to-speech (TTS) capabilities. Additionally, it offers an OpenAI-compatible API server with JSON schema mode for function calling and streaming support, and a user-friendly Streamlit UI.
+Nexa SDK is a comprehensive toolkit for supporting **ONNX** and **GGML** models. It supports text generation, image generation, vision-language models (VLM), and text-to-speech (TTS) capabilities. Additionally, it offers an OpenAI-compatible API server with JSON schema mode for function calling and streaming support, and a user-friendly Streamlit UI. Users can run Nexa SDK in any device with Python environment, and GPU acceleration is supported.
 
 ## Features
 
@@ -62,6 +62,17 @@ Below is our differentiation from other similar tools:
 
 We have released pre-built wheels for various Python versions, platforms, and backends for convenient installation on our [index page](https://nexaai.github.io/nexa-sdk/whl/).
 
+How to clone this repo
+```bash
+git clone --recursive https://github.com/NexaAI/nexa-sdk
+```
+
+If you forget to use `--recursive`, you can use below command to add submodule
+```bash
+git submodule update --init --recursive
+```
+
+
 #### CPU
 
 ```bash
@@ -93,14 +104,31 @@ CMAKE_ARGS="-DGGML_METAL=ON -DSD_METAL=ON" pip install nexaai --prefer-binary --
 
 #### GPU (CUDA)
 
-For the GPU version supporting **CUDA (Linux/Windows)**:
+For **Linux**:
 
 ```bash
 CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON" pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
 ```
 
+For **Windows PowerShell**:
+
+```bash
+$env:CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON"; pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
+```
+
+For **Windows Command Prompt**:
+
+```bash
+set CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON" & pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
+```
+
+For **Windows Git Bash**:
+
+```bash
+CMAKE_ARGS="-DGGML_CUDA=ON -DSD_CUBLAS=ON" pip install nexaai --prefer-binary --index-url https://nexaai.github.io/nexa-sdk/whl/cu124 --extra-index-url https://pypi.org/simple --no-cache-dir
+```
 > [!NOTE]
-> The CUDA wheels are built with CUDA 12.4, but should be compatible with all CUDA 12.X
+> If you want to use ONNX model, just replace `pip install nexaai` with `pip install nexaai[onnx]` in above commands
 
 
 <details>
@@ -210,4 +238,4 @@ For detailed information on server setup, API endpoints, and usage examples, ple
 We would like to thank the following projects:
 - [llama.cpp](https://github.com/ggerganov/llama.cpp)
 - [stable-diffusion.cpp](https://github.com/leejet/stable-diffusion.cpp)
-- [optimum](https://github.com/huggingface/optimum)
+- [optimum](https://github.com/huggingface/optimum)

{nexaai-0.0.8 → nexaai-0.0.8.2}/examples/ai_soulmate/README.md

@@ -26,7 +26,7 @@ This project is an AI chatbot that interacts with users via text and voice. The
 ### Technical Architecture
 
 <p align="center">
-  <img src="tech_architecture.png" alt="Technical Architecture" width="50%">
+  <img src="https://public-storage.nexa4ai.com/ai_soulmate.png" alt="Technical Architecture" width="50%">
 </p>
 
 ### Setup:

{nexaai-0.0.8 → nexaai-0.0.8.2}/examples/financial-advisor/README.md

@@ -15,6 +15,12 @@
   - `utils/text_generator.py`: handles similarity search and text generation
   - `assets/fake_bank_statements`: fake bank statement for testing purpose
 
+### Technical Architecture
+
+<p align="center">
+  <img src="https://public-storage.nexa4ai.com/financial_advisor.JPG" alt="Technical Architecture" width="70%">
+</p>
+
 ### Setup:
 
 1. Install required packages:

nexaai-0.0.8.2/nexa/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.0.8.2"

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/cli/entry.py

@@ -15,27 +15,32 @@ def run_ggml_inference(args):
     
     stop_words = kwargs.pop("stop_words", [])
 
-    if run_type == "NLP":
-        from nexa.gguf.nexa_inference_text import NexaTextInference
-        inference = NexaTextInference(model_path=model_path, local_path=local_path, stop_words=stop_words, **kwargs)
-    elif run_type == "Computer Vision":
-        from nexa.gguf.nexa_inference_image import NexaImageInference
-        inference = NexaImageInference(model_path=model_path, local_path=local_path, **kwargs)
-        if hasattr(args, 'streamlit') and args.streamlit:
-            inference.run_streamlit(model_path)
-        elif args.img2img:
-            inference.run_img2img()
+    try:
+        if run_type == "NLP":
+            from nexa.gguf.nexa_inference_text import NexaTextInference
+            inference = NexaTextInference(model_path=model_path, local_path=local_path, stop_words=stop_words, **kwargs)
+        elif run_type == "Computer Vision":
+            from nexa.gguf.nexa_inference_image import NexaImageInference
+            inference = NexaImageInference(model_path=model_path, local_path=local_path, **kwargs)
+            if hasattr(args, 'streamlit') and args.streamlit:
+                inference.run_streamlit(model_path)
+            elif args.img2img:
+                inference.run_img2img()
+            else:
+                inference.run_txt2img()
+            return
+        elif run_type == "Multimodal":
+            from nexa.gguf.nexa_inference_vlm import NexaVLMInference
+            inference = NexaVLMInference(model_path=model_path, local_path=local_path, stop_words=stop_words, **kwargs)
+        elif run_type == "Audio":
+            from nexa.gguf.nexa_inference_voice import NexaVoiceInference
+            inference = NexaVoiceInference(model_path=model_path, local_path=local_path, **kwargs)
         else:
-            inference.run_txt2img()
+            print(f"Unknown task: {run_type}. Skipping inference.")
+            return
+    except Exception as e:
+        print(f"Error loading GGUF models, please refer to our docs to install nexaai package: https://docs.nexaai.com/getting-started/installation ")
         return
-    elif run_type == "Multimodal":
-        from nexa.gguf.nexa_inference_vlm import NexaVLMInference
-        inference = NexaVLMInference(model_path=model_path, local_path=local_path, stop_words=stop_words, **kwargs)
-    elif run_type == "Audio":
-        from nexa.gguf.nexa_inference_voice import NexaVoiceInference
-        inference = NexaVoiceInference(model_path=model_path, local_path=local_path, **kwargs)
-    else:
-        raise ValueError(f"Unknown task: {run_type}")
 
     if hasattr(args, 'streamlit') and args.streamlit:
         inference.run_streamlit(model_path)
@@ -49,20 +54,25 @@ def run_onnx_inference(args):
     from nexa.general import pull_model
     local_path, run_type = pull_model(model_path)
 
-    if run_type == "NLP":
-        from nexa.onnx.nexa_inference_text import NexaTextInference as NexaTextOnnxInference
-        inference = NexaTextOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
-    elif run_type == "Computer Vision":
-        from nexa.onnx.nexa_inference_image import NexaImageInference as NexaImageOnnxInference
-        inference = NexaImageOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
-    elif run_type == "Audio":
-        from nexa.onnx.nexa_inference_voice import NexaVoiceInference as NexaVoiceOnnxInference
-        inference = NexaVoiceOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
-    elif run_type == "TTS":
-        from nexa.onnx.nexa_inference_tts import NexaTTSInference as NexaTTSOnnxInference
-        inference = NexaTTSOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
-    else:
-        raise ValueError(f"Unknown task: {run_type}")
+    try:
+        if run_type == "NLP":
+            from nexa.onnx.nexa_inference_text import NexaTextInference as NexaTextOnnxInference
+            inference = NexaTextOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
+        elif run_type == "Computer Vision":
+            from nexa.onnx.nexa_inference_image import NexaImageInference as NexaImageOnnxInference
+            inference = NexaImageOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
+        elif run_type == "Audio":
+            from nexa.onnx.nexa_inference_voice import NexaVoiceInference as NexaVoiceOnnxInference
+            inference = NexaVoiceOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
+        elif run_type == "TTS":
+            from nexa.onnx.nexa_inference_tts import NexaTTSInference as NexaTTSOnnxInference
+            inference = NexaTTSOnnxInference(model_path=model_path, local_path=local_path, **kwargs)
+        else:
+            print(f"Unknown task: {run_type}. Skipping inference.")
+            return
+    except Exception as e:
+        print(f"Error loading ONNX models, please refer to our docs to install nexaai[onnx] package: https://docs.nexaai.com/getting-started/installation ")
+        return
 
     if hasattr(args, 'streamlit') and args.streamlit:
         inference.run_streamlit(model_path)

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/constants.py

@@ -103,13 +103,24 @@ NEXA_RUN_MODEL_MAP_VOICE = {
 
 NEXA_RUN_MODEL_MAP_FUNCTION_CALLING = {
   "llama2-function-calling": "Llama2-7b-function-calling:q4_K_M",
+  "Llama2-7b-function-calling:fp16": "Llama2-7b-function-calling:fp16",
+  "Llama2-7b-function-calling:q2_K": "Llama2-7b-function-calling:q2_K",
+  "Llama2-7b-function-calling:q3_K_L": "Llama2-7b-function-calling:q3_K_L",
+  "Llama2-7b-function-calling:q3_K_M": "Llama2-7b-function-calling:q3_K_M",
+  "Llama2-7b-function-calling:q3_K_S": "Llama2-7b-function-calling:q3_K_S",
+  "Llama2-7b-function-calling:q4_K_M": "Llama2-7b-function-calling:q4_K_M",
+  "Llama2-7b-function-calling:q4_K_S": "Llama2-7b-function-calling:q4_K_S",
+  "Llama2-7b-function-calling:q5_K_M": "Llama2-7b-function-calling:q5_K_M",
+  "Llama2-7b-function-calling:q5_K_S": "Llama2-7b-function-calling:q5_K_S",
+  "Llama2-7b-function-calling:q6_K": "Llama2-7b-function-calling:q6_K",
+  "Llama2-7b-function-calling:q8_0": "Llama2-7b-function-calling:q8_0",
 }
 
 
 
 NEXA_RUN_PROJECTOR_MAP = {
     "nanollava": "nanoLLaVA:projector-fp16",
-    "nanoLLaVA:fp16": "nanoLLaVA:project-fp16",
+    "nanoLLaVA:fp16": "nanoLLaVA:projector-fp16",
     "llava-phi3": "llava-phi-3-mini:projector-q4_0",
     "llava-phi-3-mini:q4_0": "llava-phi-3-mini:projector-q4_0",
     "llava-phi-3-mini:fp16": "llava-phi-3-mini:projector-fp16",

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/gguf/llama/_utils_transformers.py

@@ -17,7 +17,7 @@ class suppress_stdout_stderr(object):
     sys = sys
     os = os
 
-    def __init__(self, disable: bool = True):
+    def __init__(self, disable: bool = False):
         self.disable = disable
 
     # Oddly enough this works better than the contextlib version

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/gguf/llama/llama.py

@@ -1337,22 +1337,25 @@ class Llama:
                         returned_tokens += 1
                         yield {
                             "id": completion_id,
-                            "object": "text_completion",
-                            "created": created,
                             "model": model_name,
+                            "object": "text_completion.chunk",
+                            "created": created,
                             "choices": [
                                 {
-                                    "text": self.detokenize(
-                                        [token],
-                                        prev_tokens=prompt_tokens
-                                        + completion_tokens[:returned_tokens],
-                                    ).decode("utf-8", errors="ignore"),
                                     "index": 0,
-                                    "logprobs": logprobs_or_none,
+                                    "delta": {
+                                        "content": self.detokenize(
+                                            [token],
+                                            prev_tokens=prompt_tokens
+                                            + completion_tokens[:returned_tokens],
+                                        ).decode("utf-8", errors="ignore")
+                                    },
                                     "finish_reason": None,
+                                    "logprobs": logprobs_or_none,
                                 }
                             ],
                         }
+
                 else:
                     while len(remaining_tokens) > 0:
                         decode_success = False
@@ -1383,19 +1386,22 @@ class Llama:
 
                         yield {
                             "id": completion_id,
-                            "object": "text_completion",
-                            "created": created,
                             "model": model_name,
+                            "object": "text_completion.chunk",
+                            "created": created,
                             "choices": [
                                 {
-                                    "text": ts,
                                     "index": 0,
                                     "logprobs": None,
+                                    "delta": {
+                                        "content": ts,
+                                    },
                                     "finish_reason": None,
                                 }
                             ],
                         }
 
+
             if len(completion_tokens) >= max_tokens:
                 text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens)
                 finish_reason = "length"
@@ -1475,52 +1481,59 @@ class Llama:
                     returned_tokens += 1
                     yield {
                         "id": completion_id,
-                        "object": "text_completion",
-                        "created": created,
                         "model": model_name,
+                        "object": "text_completion.chunk",
+                        "created": created,
                         "choices": [
                             {
-                                "text": last_text[
-                                    : len(last_text) - (token_end_position - end)
-                                ].decode("utf-8", errors="ignore"),
                                 "index": 0,
+                                "delta": {
+                                    "content": last_text[: len(last_text) - (token_end_position - end)].decode(
+                                        "utf-8", errors="ignore"
+                                    )
+                                },
                                 "logprobs": logprobs_or_none,
                                 "finish_reason": None,
                             }
                         ],
                     }
+
                     break
                 returned_tokens += 1
                 yield {
                     "id": completion_id,
-                    "object": "text_completion",
-                    "created": created,
                     "model": model_name,
+                    "object": "text_completion.chunk",
+                    "created": created,
                     "choices": [
                         {
-                            "text": self.detokenize([token]).decode(
-                                "utf-8", errors="ignore"
-                            ),
                             "index": 0,
+                            "delta": {
+                                "content": self.detokenize([token]).decode("utf-8", errors="ignore")
+                            },
                             "logprobs": logprobs_or_none,
                             "finish_reason": None,
                         }
                     ],
                 }
+
             yield {
                 "id": completion_id,
-                "object": "text_completion",
-                "created": created,
                 "model": model_name,
+                "object": "text_completion.chunk",
+                "created": created,
                 "choices": [
                     {
-                        "text": "",
                         "index": 0,
+                        "delta": {
+                            "content": "",
+                        },
                         "logprobs": None,
                         "finish_reason": finish_reason,
                     }
                 ],
             }
+
             if self.cache:
                 if self.verbose:
                     print("Llama._create_completion: cache save", file=sys.stderr)

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/gguf/llama/llama_chat_format.py

@@ -304,6 +304,7 @@ def _convert_text_completion_chunks_to_chat(
                     }
                 ],
             }
+
         yield {
             "id": "chat" + chunk["id"],
             "model": chunk["model"],
@@ -314,7 +315,7 @@ def _convert_text_completion_chunks_to_chat(
                     "index": 0,
                     "delta": (
                         {
-                            "content": chunk["choices"][0]["text"],
+                            "content": chunk["choices"][0]["delta"]["content"],
                         }
                         if chunk["choices"][0]["finish_reason"] is None
                         else {}

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/gguf/llama/llama_types.py

@@ -278,6 +278,24 @@ ChatCompletionToolChoiceOption = Union[
     Literal["none", "auto", "required"], ChatCompletionNamedToolChoice
 ]
 
+class CompletionStreamResponseDelta(TypedDict):
+    content: Optional[str]
+
+
+class CompletionStreamResponseChoice(TypedDict):
+    index: int
+    delta: CompletionStreamResponseDelta
+    finish_reason: Optional[Literal["stop", "length"]]
+    logprobs: NotRequired[Optional[CompletionLogprobs]]
+
+
+class CreateCompletionStreamResponse(TypedDict):
+    id: str
+    model: str
+    object: Literal["text_completion.chunk"]
+    created: int
+    choices: List[CompletionStreamResponseChoice]
+
 
 # NOTE: The following type names are not part of the OpenAI OpenAPI specification
 # and will be removed in a future major release.
@@ -285,7 +303,6 @@ ChatCompletionToolChoiceOption = Union[
 EmbeddingData = Embedding
 CompletionChunk = CreateCompletionResponse
 Completion = CreateCompletionResponse
-CreateCompletionStreamResponse = CreateCompletionResponse
 ChatCompletionMessage = ChatCompletionResponseMessage
 ChatCompletionChoice = ChatCompletionResponseChoice
 ChatCompletion = CreateChatCompletionResponse

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/gguf/nexa_inference_image.py

@@ -115,7 +115,7 @@ class NexaImageInference:
             except Exception as e:
                 logging.error(f"Attempt {attempt + 1} failed with error: {e}")
                 time.sleep(1)
-        logging.error("All retry attempts failed.")
+        print("All retry attempts failed becase of Out of Memory error, Try to use smaller models...")
         return None
 
     def txt2img(

{nexaai-0.0.8 → nexaai-0.0.8.2}/nexa/gguf/server/nexa_service.py

@@ -62,68 +62,64 @@ class GenerationRequest(BaseModel):
 
 async def load_model():
     global model, chat_format, completion_template, model_path
-    if model_path in NEXA_RUN_MODEL_MAP_TEXT:
-        chat_format = NEXA_RUN_CHAT_TEMPLATE_MAP.get(model_path, None)
-        completion_template = NEXA_RUN_COMPLETION_TEMPLATE_MAP.get(model_path, None)
-        model_path = NEXA_RUN_MODEL_MAP_TEXT.get(model_path)
-        downloaded_path, run_type = pull_model(model_path)
-        with suppress_stdout_stderr():
-            try:
-                model = Llama(
-                    model_path=downloaded_path,
-                    verbose=False,
-                    chat_format=chat_format,
-                    n_gpu_layers=-1 if is_gpu_available() else 0,
-                )
-            except Exception as e:
-                logging.error(
-                    f"Failed to load model: {e}. Falling back to CPU.", exc_info=True
-                )
-                model = Llama(
-                    model_path=downloaded_path,
-                    verbose=False,
-                    chat_format=chat_format,
-                    n_gpu_layers=0,  # hardcode to use CPU
-                )
-            logging.info(f"model loaded as {model}")
-    elif model_path in NEXA_RUN_MODEL_MAP_FUNCTION_CALLING:
-        chat_format = "chatml-function-calling"
-        model_path = NEXA_RUN_MODEL_MAP_FUNCTION_CALLING.get(model_path)
-        downloaded_path, run_type = pull_model(model_path)
-        with suppress_stdout_stderr():
-            try:
-                model = Llama(
-                    model_path=downloaded_path,
-                    verbose=False,
-                    chat_format=chat_format,
-                    n_gpu_layers=-1 if is_gpu_available() else 0,
-                )
-            except Exception as e:
-                logging.error(
-                    f"Failed to load model: {e}. Falling back to CPU.", exc_info=True
-                )
-                model = Llama(
-                    model_path=downloaded_path,
-                    verbose=False,
-                    chat_format=chat_format,
-                    n_gpu_layers=0,  # hardcode to use CPU
-                )
-
-            logging.info(f"model loaded as {model}")
-    elif model_path in NEXA_RUN_MODEL_MAP_IMAGE:
-        downloaded_path, run_type = pull_model(model_path)
+    downloaded_path, run_type = pull_model(model_path)
+
+    if run_type == "NLP":
+        if model_path in NEXA_RUN_MODEL_MAP_FUNCTION_CALLING:
+            chat_format = "chatml-function-calling"
+            with suppress_stdout_stderr():
+                try:
+                    model = Llama(
+                        model_path=downloaded_path,
+                        verbose=False,
+                        chat_format=chat_format,
+                        n_gpu_layers=-1 if is_gpu_available() else 0,
+                    )
+                except Exception as e:
+                    logging.error(
+                        f"Failed to load model: {e}. Falling back to CPU.", exc_info=True
+                    )
+                    model = Llama(
+                        model_path=downloaded_path,
+                        verbose=False,
+                        chat_format=chat_format,
+                        n_gpu_layers=0,  # hardcode to use CPU
+                    )
+
+                logging.info(f"model loaded as {model}")
+        else:
+            chat_format = NEXA_RUN_CHAT_TEMPLATE_MAP.get(model_path, None)
+            completion_template = NEXA_RUN_COMPLETION_TEMPLATE_MAP.get(model_path, None)
+            with suppress_stdout_stderr():
+                try:
+                    model = Llama(
+                        model_path=downloaded_path,
+                        verbose=False,
+                        chat_format=chat_format,
+                        n_gpu_layers=-1 if is_gpu_available() else 0,
+                    )
+                except Exception as e:
+                    logging.error(
+                        f"Failed to load model: {e}. Falling back to CPU.", exc_info=True
+                    )
+                    model = Llama(
+                        model_path=downloaded_path,
+                        verbose=False,
+                        chat_format=chat_format,
+                        n_gpu_layers=0,  # hardcode to use CPU
+                    )
+                logging.info(f"model loaded as {model}")
+    elif run_type == "Computer Vision":
         with suppress_stdout_stderr():
             model = StableDiffusion(
                 model_path=downloaded_path,
                 wtype=NEXA_RUN_MODEL_PRECISION_MAP.get(
-                    model_path, "default"
+                    model_path, "f32"
                 ),  # Weight type (options: default, f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
                 n_threads=multiprocessing.cpu_count(),
             )
         logging.info(f"model loaded as {model}")
-    elif model_path in NEXA_RUN_MODEL_MAP_VOICE:
-        model_path = NEXA_RUN_MODEL_MAP_VOICE.get(model_path)
-        downloaded_path, run_type = pull_model(model_path)
+    elif run_type == "Audio":
         with suppress_stdout_stderr():
             model = WhisperModel(
                 downloaded_path,
@@ -132,7 +128,7 @@ async def load_model():
             )
         logging.info(f"model loaded as {model}")
     else:
-        raise ValueError(f"Model {model_path} not found in NEXA_RUN_MODEL_MAP")
+        raise ValueError(f"Model {model_path} not found in Model Hub")
 
 
 @app.on_event("startup")