nexaai 0.0.8.6__tar.gz → 0.0.8.6.1__tar.gz

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nexaai
-Version: 0.0.8.6
+Version: 0.0.8.6.1
 Summary: Nexa AI SDK
 Author-Email: Nexa AI <octopus@nexa4ai.com>
 License: MIT

nexaai-0.0.8.6.1/nexa/__init__.py

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

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/cli/entry.py

@@ -108,6 +108,7 @@ def main():
     text_group.add_argument("-k", "--top_k", type=int, help="Top-k sampling parameter")
     text_group.add_argument("-p", "--top_p", type=float, help="Top-p sampling parameter")
     text_group.add_argument("-sw", "--stop_words", nargs="*", help="List of stop words for early stopping")
+    text_group.add_argument("--lora_path", type=str, help="Path to a LoRA file to apply to the model.")
     text_group.add_argument("-hf", "--huggingface", action="store_true", help="Load model from Hugging Face Hub")
 
     # Image generation arguments
@@ -123,7 +124,7 @@ def main():
     image_group.add_argument("--wtype", type=str, help="Weight type (f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)")
     image_group.add_argument("--control_net_path", type=str, help="Path to control net model")
     image_group.add_argument("--control_image_path", type=str, help="Path to image condition for Control Net")
-    image_group.add_argument("--control_strength", type=str, help="Strength to apply Control Net")
+    image_group.add_argument("--control_strength", type=float, help="Strength to apply Control Net")
 
     # ASR arguments
     asr_group = run_parser.add_argument_group('Automatic Speech Recognition options')

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/general.py

@@ -8,6 +8,7 @@ import concurrent.futures
 import time
 import os
 from tqdm import tqdm
+import platform
 
 from nexa.constants import (
     NEXA_API_URL,
@@ -105,7 +106,7 @@ def get_user_info(token):
         return None
 
 
-def pull_model(model_path, hf = False):
+def pull_model(model_path, hf = False, **kwargs):
     model_path = NEXA_RUN_MODEL_MAP.get(model_path, model_path)
 
     try:
@@ -118,9 +119,9 @@ def pull_model(model_path, hf = False):
                 return location, run_type
 
             if "/" in model_path:
-                result = pull_model_from_hub(model_path)
+                result = pull_model_from_hub(model_path, **kwargs)
             else:
-                result = pull_model_from_official(model_path)
+                result = pull_model_from_official(model_path, **kwargs)
 
         if result["success"]:
             add_model_to_list(model_path, result["local_path"], result["model_type"], result["run_type"])
@@ -134,7 +135,7 @@ def pull_model(model_path, hf = False):
         return None, "NLP"
 
 
-def pull_model_from_hub(model_path):
+def pull_model_from_hub(model_path, **kwargs):
     NEXA_MODELS_HUB_DIR.mkdir(parents=True, exist_ok=True)
 
     token = ""
@@ -174,7 +175,7 @@ def pull_model_from_hub(model_path):
     for file_path, presigned_link in presigned_links.items():
         try:
             download_path = NEXA_MODELS_HUB_DIR / file_path
-            download_file_with_progress(presigned_link, download_path, use_processes=True)
+            download_file_with_progress(presigned_link, download_path, **kwargs)
 
             if local_path is None:
                 if model_type == "onnx" or model_type == "bin":
@@ -196,7 +197,7 @@ def pull_model_from_hub(model_path):
     }
 
 
-def pull_model_from_official(model_path):
+def pull_model_from_official(model_path, **kwargs):
     NEXA_MODELS_HUB_OFFICIAL_DIR.mkdir(parents=True, exist_ok=True)
 
     if "onnx" in model_path:
@@ -208,7 +209,7 @@ def pull_model_from_official(model_path):
 
     run_type = get_run_type_from_model_path(model_path)
     run_type_str = run_type.value if isinstance(run_type, ModelType) else str(run_type)
-    success, location = download_model_from_official(model_path, model_type)
+    success, location = download_model_from_official(model_path, model_type, **kwargs)
     
     return {
         "success": success,
@@ -288,6 +289,21 @@ def download_chunk(url, start, end, output_file, chunk_number):
             if attempt == max_retries - 1:
                 raise
             time.sleep(2 ** attempt)  # Exponential backoff
+            
+            
+def default_use_processes():
+    """
+    Distinct operating systems may have different default behaviors for threading vs multiprocessing.
+    """
+    platform_name = platform.system()
+    if platform_name == "Linux":
+        return True
+    elif platform_name == "Windows":
+        return False
+    elif platform_name == "Darwin":
+        return False
+    else:
+        return False
 
 
 def download_file_with_progress(
@@ -295,7 +311,8 @@ def download_file_with_progress(
     file_path: Path,
     chunk_size: int = 40 * 1024 * 1024,
     max_workers: int = 20,
-    use_processes: bool = False
+    use_processes: bool = default_use_processes(),
+    **kwargs
 ):
     file_path.parent.mkdir(parents=True, exist_ok=True)
 
@@ -376,7 +393,7 @@ def download_file_with_progress(
             os.remove(file_path)
 
 
-def download_model_from_official(model_path, model_type):
+def download_model_from_official(model_path, model_type, **kwargs):
     try:
         model_name, model_version = model_path.split(":")
         file_extension = ".zip" if model_type == "onnx" or model_type == "bin" else ".gguf"
@@ -388,7 +405,7 @@ def download_model_from_official(model_path, model_type):
         download_url = f"{NEXA_OFFICIAL_BUCKET}{filepath}"
 
         full_path.parent.mkdir(parents=True, exist_ok=True)
-        download_file_with_progress(download_url, full_path, use_processes=True)
+        download_file_with_progress(download_url, full_path, **kwargs)
 
         if model_type == "onnx" or model_type == "bin":
             unzipped_folder = full_path.parent / model_version

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/gguf/nexa_inference_image.py

@@ -53,8 +53,14 @@ class NexaImageInference:
     guidance_scale (float): Guidance scale for diffusion.
     output_path (str): Output path for the generated image.
     random_seed (int): Random seed for image generation.
+    lora_dir (str): Path to directory containing LoRA files.
+    lora_path (str): Path to a LoRA file to apply to the model.
+    wtype (str): Weight type (options: default, f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0).
+    control_net_path (str): Path to control net model.
+    control_image_path (str): Path to image condition for Control Net.
+    control_strength (float): Strength to apply Control Net.
     streamlit (bool): Run the inference in Streamlit UI.
-
+    profiling (bool): Enable profiling logs for the inference process.
     """
 
     def __init__(self, model_path, local_path=None, **kwargs):
@@ -71,7 +77,7 @@ class NexaImageInference:
 
         # Download base model if not provided
         if self.downloaded_path is None:
-            self.downloaded_path, _ = pull_model(self.model_path)
+            self.downloaded_path, _ = pull_model(self.model_path, **kwargs)
             if self.downloaded_path is None:
                 logging.error(
                     f"Model ({model_path}) is not applicable. Please refer to our docs for proper usage.",
@@ -86,11 +92,11 @@ class NexaImageInference:
             self.clip_l_path = FLUX_CLIP_L_PATH
 
             if self.t5xxl_path:
-                self.t5xxl_downloaded_path, _ = pull_model(self.t5xxl_path)
+                self.t5xxl_downloaded_path, _ = pull_model(self.t5xxl_path, **kwargs)
             if self.ae_path:
-                self.ae_downloaded_path, _ = pull_model(self.ae_path)
+                self.ae_downloaded_path, _ = pull_model(self.ae_path, **kwargs)
             if self.clip_l_path:
-                self.clip_l_downloaded_path, _ = pull_model(self.clip_l_path)
+                self.clip_l_downloaded_path, _ = pull_model(self.clip_l_path, **kwargs)
         if "lcm-dreamshaper" in self.model_path or "flux" in self.model_path:
             self.params = DEFAULT_IMG_GEN_PARAMS_LCM.copy() # both lcm-dreamshaper and flux use the same params
         elif "sdxl-turbo" in self.model_path:
@@ -98,6 +104,7 @@ class NexaImageInference:
         else:
             self.params = DEFAULT_IMG_GEN_PARAMS.copy()
 
+        self.profiling = kwargs.get("profiling", False)
         self.params.update({k: v for k, v in kwargs.items() if v is not None})
         if not kwargs.get("streamlit", False):
             self._load_model(model_path)
@@ -119,7 +126,7 @@ class NexaImageInference:
                     wtype=self.params.get(
                         "wtype", NEXA_RUN_MODEL_PRECISION_MAP.get(model_path, "default")
                     ),  # Weight type (options: default, f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
-                    verbose=False,
+                    verbose=self.profiling,
                 )
             else:
                 self.model = StableDiffusion(
@@ -130,7 +137,7 @@ class NexaImageInference:
                         "wtype", NEXA_RUN_MODEL_PRECISION_MAP.get(model_path, "default")
                     ),  # Weight type (options: default, f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)
                     control_net_path=self.params.get("control_net_path", ""),
-                    verbose=False,
+                    verbose=self.profiling,
                 )
 
     def _save_images(self, images):
@@ -352,13 +359,43 @@ if __name__ == "__main__":
         default=0,
         help="Random seed for image generation",
     )
-    # parser.add_argument("--device", type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help="Device to run the model on (default: cuda if available, else cpu)")
+    parser.add_argument(
+        "--lora_dir",
+        type=str,
+        help="Path to directory containing LoRA files.",
+    )
+    parser.add_argument(
+        "--wtype",
+        type=str,
+        help="Weight type (options: default, f32, f16, q4_0, q4_1, q5_0, q5_1, q8_0)",
+    )
+    parser.add_argument(
+        "--control_net_path",
+        type=str,
+        help="Path to control net model.",
+    )
+    parser.add_argument(
+        "--control_image_path",
+        type=str,
+        help="Path to image condition for Control Net.",
+    )
+    parser.add_argument(
+        "--control_strength",
+        type=float,
+        help="Strength to apply Control Net.",
+    )
     parser.add_argument(
         "-st",
         "--streamlit",
         action="store_true",
         help="Run the inference in Streamlit UI",
     )
+    parser.add_argument(
+        "-pf",
+        "--profiling",
+        action="store_true",
+        help="Enable profiling logs for the inference process",
+    )
     args = parser.parse_args()
     kwargs = {k: v for k, v in vars(args).items() if v is not None}
     model_path = kwargs.pop("model_path")

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/gguf/nexa_inference_text.py

@@ -56,7 +56,7 @@ class NexaTextInference:
         self.top_logprobs = kwargs.get('top_logprobs', None)
 
         if self.downloaded_path is None:
-            self.downloaded_path, _ = pull_model(self.model_path)
+            self.downloaded_path, _ = pull_model(self.model_path, **kwargs)
 
         if self.downloaded_path is None:
             logging.error(
@@ -107,6 +107,7 @@ class NexaTextInference:
                     chat_format=self.chat_format,
                     n_ctx=2048,
                     n_gpu_layers=-1 if is_gpu_available() else 0,
+                    lora_path=self.params.get("lora_path", ""),
                 )
             except Exception as e:
                 logging.error(f"Failed to load model: {e}. Falling back to CPU.", exc_info=True)
@@ -116,6 +117,7 @@ class NexaTextInference:
                     chat_format=self.chat_format,
                     n_ctx=2048,
                     n_gpu_layers=0,  # hardcode to use CPU
+                    lora_path=self.params.get("lora_path", ""),
                 )
 
         load_time = time.time() - start_time
@@ -331,13 +333,11 @@ if __name__ == "__main__":
         action="store_true",
         help="Run the inference in Streamlit UI",
     )
-    # parser.add_argument(
-    #     "-tlps",
-    #     "--top_logprobs",
-    #     type=int,
-    #     default=None,  # -tlps 5
-    #     help="Number of most likely tokens to return at each token position",
-    # )
+    parser.add_argument(
+        "--lora_path",
+        type=str,
+        help="Path to a LoRA file to apply to the model.",
+    )
     args = parser.parse_args()
     kwargs = {k: v for k, v in vars(args).items() if v is not None}
     model_path = kwargs.pop("model_path")

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/gguf/nexa_inference_vlm.py

@@ -101,12 +101,12 @@ class NexaVLMInference:
         if self.downloaded_path is not None:
             if model_path in NEXA_RUN_MODEL_MAP_VLM:
                 self.projector_path = NEXA_RUN_PROJECTOR_MAP[model_path]
-                self.projector_downloaded_path, _ = pull_model(self.projector_path)
+                self.projector_downloaded_path, _ = pull_model(self.projector_path, **kwargs)
         elif model_path in NEXA_RUN_MODEL_MAP_VLM:
             self.model_path = NEXA_RUN_MODEL_MAP_VLM[model_path]
             self.projector_path = NEXA_RUN_PROJECTOR_MAP[model_path]
-            self.downloaded_path, _ = pull_model(self.model_path)
-            self.projector_downloaded_path, _ = pull_model(self.projector_path)
+            self.downloaded_path, _ = pull_model(self.model_path, **kwargs)
+            self.projector_downloaded_path, _ = pull_model(self.projector_path, **kwargs)
         elif Path(model_path).parent.exists():
             local_dir = Path(model_path).parent
             model_name = Path(model_path).name

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/gguf/nexa_inference_voice.py

@@ -45,7 +45,7 @@ class NexaVoiceInference:
         self.params = DEFAULT_VOICE_GEN_PARAMS
 
         if self.downloaded_path is None:
-            self.downloaded_path, _ = pull_model(self.model_path)
+            self.downloaded_path, _ = pull_model(self.model_path, **kwargs)
 
         if self.downloaded_path is None:
             logging.error(

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/onnx/nexa_inference_image.py

@@ -64,7 +64,7 @@ class NexaImageInference:
     def run(self):
 
         if self.download_onnx_folder is None:
-            self.download_onnx_folder, run_type = pull_model(self.model_path)
+            self.download_onnx_folder, run_type = pull_model(self.model_path, **kwargs)
 
         if self.download_onnx_folder is None:
             logging.error(

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/onnx/nexa_inference_text.py

@@ -146,7 +146,7 @@ class NexaTextInference:
             self.run_streamlit()
         else:
             if self.downloaded_onnx_folder is None:
-                self.downloaded_onnx_folder, run_type = pull_model(self.model_path)
+                self.downloaded_onnx_folder, run_type = pull_model(self.model_path, **kwargs)
 
             if self.downloaded_onnx_folder is None:
                 logging.error(

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/onnx/nexa_inference_tts.py

@@ -47,7 +47,7 @@ class NexaTTSInference:
         self.downloaded_onnx_folder = local_path
 
         if self.downloaded_onnx_folder is None:
-            self.downloaded_onnx_folder, run_type = pull_model(self.model_path)
+            self.downloaded_onnx_folder, run_type = pull_model(self.model_path, **kwargs)
         
         if self.downloaded_onnx_folder is None:
             logging.error(

{nexaai-0.0.8.6 → nexaai-0.0.8.6.1}/nexa/onnx/nexa_inference_voice.py

@@ -42,7 +42,7 @@ class NexaVoiceInference:
 
     def run(self):
         if self.downloaded_onnx_folder is None:
-            self.downloaded_onnx_folder, run_type = pull_model(self.model_path)
+            self.downloaded_onnx_folder, run_type = pull_model(self.model_path, **kwargs)
 
         if self.downloaded_onnx_folder is None:
             logging.error(

nexaai-0.0.8.6/nexa/__init__.py

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