nexaai 1.0.6rc1__cp310-cp310-macosx_14_0_universal2.whl → 1.0.7__cp310-cp310-macosx_14_0_universal2.whl

nexaai/utils/manifest_utils.py

@@ -0,0 +1,280 @@
+"""
+Manifest and metadata utilities for handling nexa.manifest files and model metadata.
+
+This module provides utilities to:
+- Load and save nexa.manifest files
+- Create GGUF and MLX manifests
+- Process manifest metadata (handle null fields, fetch avatars, etc.)
+- Manage backward compatibility with old download_metadata.json files
+"""
+
+import os
+import json
+from datetime import datetime
+from typing import Dict, Any, List, Optional
+
+from .quantization_utils import (
+    extract_quantization_from_filename, 
+    detect_quantization_for_mlx
+)
+from .model_types import (
+    PIPELINE_TO_MODEL_TYPE,
+    MODEL_TYPE_TO_PIPELINE
+)
+
+
+def process_manifest_metadata(manifest: Dict[str, Any], repo_id: str) -> Dict[str, Any]:
+    """Process manifest metadata to handle null/missing fields."""
+    # Handle pipeline_tag
+    pipeline_tag = manifest.get('pipeline_tag')
+    if not pipeline_tag:
+        # Reverse map from ModelType if available
+        model_type = manifest.get('ModelType')
+        pipeline_tag = MODEL_TYPE_TO_PIPELINE.get(model_type) if model_type else None
+    
+    # Handle download_time - keep as null if missing
+    download_time = manifest.get('download_time')
+    
+    # Handle avatar_url - fetch on-the-fly if missing/null
+    avatar_url = manifest.get('avatar_url')
+    if not avatar_url:
+        try:
+            from .avatar_fetcher import get_avatar_url_for_repo
+            avatar_url = get_avatar_url_for_repo(repo_id)
+        except Exception:
+            # If fetching fails, leave as None
+            avatar_url = None
+    
+    # Return processed metadata
+    processed_manifest = manifest.copy()
+    processed_manifest.update({
+        'pipeline_tag': pipeline_tag,
+        'download_time': download_time,
+        'avatar_url': avatar_url
+    })
+    
+    return processed_manifest
+
+
+def load_nexa_manifest(directory_path: str) -> Dict[str, Any]:
+    """Load manifest from nexa.manifest if it exists."""
+    manifest_path = os.path.join(directory_path, 'nexa.manifest')
+    if os.path.exists(manifest_path):
+        try:
+            with open(manifest_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        except (json.JSONDecodeError, IOError):
+            pass
+    return {}
+
+
+def load_download_metadata(directory_path: str, repo_id: Optional[str] = None) -> Dict[str, Any]:
+    """Load download metadata from nexa.manifest if it exists, fallback to old format."""
+    # First try to load from new manifest format
+    manifest = load_nexa_manifest(directory_path)
+    if manifest and repo_id:
+        # Process the manifest to handle null/missing fields
+        return process_manifest_metadata(manifest, repo_id)
+    elif manifest:
+        # Return manifest as-is if no repo_id provided (for backward compatibility)
+        return manifest
+    
+    # Fallback to old format for backward compatibility
+    old_metadata_path = os.path.join(directory_path, 'download_metadata.json')
+    if os.path.exists(old_metadata_path):
+        try:
+            with open(old_metadata_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+        except (json.JSONDecodeError, IOError):
+            pass
+    return {}
+
+
+def save_download_metadata(directory_path: str, metadata: Dict[str, Any]) -> None:
+    """Save download metadata to nexa.manifest in the new format."""
+    manifest_path = os.path.join(directory_path, 'nexa.manifest')
+    try:
+        with open(manifest_path, 'w', encoding='utf-8') as f:
+            json.dump(metadata, f, indent=2)
+    except IOError:
+        # If we can't save metadata, don't fail the download
+        pass
+
+
+def create_gguf_manifest(repo_id: str, files: List[str], directory_path: str, old_metadata: Dict[str, Any]) -> Dict[str, Any]:
+    """Create GGUF format manifest."""
+    
+    # Load existing manifest to merge GGUF files if it exists
+    existing_manifest = load_nexa_manifest(directory_path)
+    
+    model_files = {}
+    if existing_manifest and "ModelFile" in existing_manifest:
+        model_files = existing_manifest["ModelFile"].copy()
+    
+    # Process GGUF files
+    for file_name in files:
+        if file_name.endswith('.gguf'):
+            # Use the new enum-based quantization extraction
+            quantization_type = extract_quantization_from_filename(file_name)
+            quant_level = quantization_type.value if quantization_type else "UNKNOWN"
+            
+            file_path = os.path.join(directory_path, file_name)
+            file_size = 0
+            if os.path.exists(file_path):
+                try:
+                    file_size = os.path.getsize(file_path)
+                except (OSError, IOError):
+                    pass
+            
+            model_files[quant_level] = {
+                "Name": file_name,
+                "Downloaded": True,
+                "Size": file_size
+            }
+    
+    manifest = {
+        "Name": repo_id,
+        "ModelType": PIPELINE_TO_MODEL_TYPE.get(old_metadata.get('pipeline_tag'), "other"),
+        "PluginId": "llama_cpp",
+        "ModelFile": model_files,
+        "MMProjFile": {
+            "Name": "",
+            "Downloaded": False,
+            "Size": 0
+        },
+        "TokenizerFile": {
+            "Name": "",
+            "Downloaded": False,
+            "Size": 0
+        },
+        "ExtraFiles": None,
+        # Preserve old metadata fields
+        "pipeline_tag": old_metadata.get('pipeline_tag'),
+        "download_time": old_metadata.get('download_time'),
+        "avatar_url": old_metadata.get('avatar_url')
+    }
+    
+    return manifest
+
+
+def create_mlx_manifest(repo_id: str, files: List[str], directory_path: str, old_metadata: Dict[str, Any]) -> Dict[str, Any]:
+    """Create MLX format manifest."""
+    
+    model_files = {}
+    extra_files = []
+    
+    # Try different methods to extract quantization for MLX models
+    quantization_type = detect_quantization_for_mlx(repo_id, directory_path)
+    
+    # Use the detected quantization or default to "DEFAULT"
+    quant_level = quantization_type.value if quantization_type else "DEFAULT"
+    
+    for file_name in files:
+        file_path = os.path.join(directory_path, file_name)
+        file_size = 0
+        if os.path.exists(file_path):
+            try:
+                file_size = os.path.getsize(file_path)
+            except (OSError, IOError):
+                pass
+        
+        # Check if this is a main model file (safetensors but not index files)
+        if (file_name.endswith('.safetensors') and not file_name.endswith('.index.json')):
+            model_files[quant_level] = {
+                "Name": file_name,
+                "Downloaded": True,
+                "Size": file_size
+            }
+        else:
+            # Add to extra files
+            extra_files.append({
+                "Name": file_name,
+                "Downloaded": True,
+                "Size": file_size
+            })
+    
+    manifest = {
+        "Name": repo_id,
+        "ModelType": PIPELINE_TO_MODEL_TYPE.get(old_metadata.get('pipeline_tag'), "other"),
+        "PluginId": "mlx",
+        "ModelFile": model_files,
+        "MMProjFile": {
+            "Name": "",
+            "Downloaded": False,
+            "Size": 0
+        },
+        "TokenizerFile": {
+            "Name": "",
+            "Downloaded": False,
+            "Size": 0
+        },
+        "ExtraFiles": extra_files if extra_files else None,
+        # Preserve old metadata fields
+        "pipeline_tag": old_metadata.get('pipeline_tag'),
+        "download_time": old_metadata.get('download_time'),
+        "avatar_url": old_metadata.get('avatar_url')
+    }
+    
+    return manifest
+
+
+def detect_model_type(files: List[str]) -> str:
+    """Detect if this is a GGUF or MLX model based on file extensions."""
+    has_gguf = any(f.endswith('.gguf') for f in files)
+    has_safetensors = any(f.endswith('.safetensors') or 'safetensors' in f for f in files)
+    
+    if has_gguf:
+        return "gguf"
+    elif has_safetensors:
+        return "mlx"
+    else:
+        # Default to mlx for other types
+        return "mlx"
+
+
+def create_manifest_from_files(repo_id: str, files: List[str], directory_path: str, old_metadata: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Create appropriate manifest format based on detected model type.
+    
+    Args:
+        repo_id: Repository ID
+        files: List of files in the model directory
+        directory_path: Path to the model directory
+        old_metadata: Existing metadata (pipeline_tag, download_time, avatar_url)
+        
+    Returns:
+        Dict containing the appropriate manifest format
+    """
+    model_type = detect_model_type(files)
+    
+    if model_type == "gguf":
+        return create_gguf_manifest(repo_id, files, directory_path, old_metadata)
+    else:  # mlx or other
+        return create_mlx_manifest(repo_id, files, directory_path, old_metadata)
+
+
+def save_manifest_with_files_metadata(repo_id: str, local_dir: str, old_metadata: Dict[str, Any]) -> None:
+    """
+    Create and save manifest based on files found in the directory.
+    
+    Args:
+        repo_id: Repository ID
+        local_dir: Local directory containing the model files
+        old_metadata: Existing metadata to preserve
+    """
+    # Get list of files in the directory
+    files = []
+    try:
+        for root, dirs, filenames in os.walk(local_dir):
+            for filename in filenames:
+                # Store relative path from the directory
+                rel_path = os.path.relpath(os.path.join(root, filename), local_dir)
+                files.append(rel_path)
+    except (OSError, IOError):
+        pass
+    
+    # Create appropriate manifest
+    manifest = create_manifest_from_files(repo_id, files, local_dir, old_metadata)
+    
+    # Save manifest
+    save_download_metadata(local_dir, manifest)

nexaai/utils/model_manager.py

@@ -5,17 +5,21 @@ from datetime import datetime
 from dataclasses import dataclass
 from typing import Optional, Callable, Dict, Any, List, Union
 import functools
+from enum import Enum
 from tqdm.auto import tqdm
 from huggingface_hub import HfApi
 from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError
 
 from .progress_tracker import CustomProgressTqdm, DownloadProgressTracker
 from .avatar_fetcher import get_avatar_url_for_repo
+from .manifest_utils import (
+    load_download_metadata,
+    save_manifest_with_files_metadata,
+)
 
 # Default path for model storage
 DEFAULT_MODEL_SAVING_PATH = "~/.cache/nexa.ai/nexa_sdk/models/"
 
-
 @dataclass
 class DownloadedModel:
     """Data class representing a downloaded model with all its metadata."""
@@ -88,30 +92,6 @@ def _check_for_incomplete_downloads(directory_path: str) -> bool:
         # If we can't access the directory, assume download is complete
         return True
 
-
-def _load_download_metadata(directory_path: str) -> Dict[str, Any]:
-    """Load download metadata from download_metadata.json if it exists."""
-    metadata_path = os.path.join(directory_path, 'download_metadata.json')
-    if os.path.exists(metadata_path):
-        try:
-            with open(metadata_path, 'r', encoding='utf-8') as f:
-                return json.load(f)
-        except (json.JSONDecodeError, IOError):
-            pass
-    return {}
-
-
-def _save_download_metadata(directory_path: str, metadata: Dict[str, Any]) -> None:
-    """Save download metadata to download_metadata.json."""
-    metadata_path = os.path.join(directory_path, 'download_metadata.json')
-    try:
-        with open(metadata_path, 'w', encoding='utf-8') as f:
-            json.dump(metadata, f, indent=2)
-    except IOError:
-        # If we can't save metadata, don't fail the download
-        pass
-
-
 def _get_directory_size_and_files(directory_path: str) -> tuple[int, List[str]]:
     """Get total size and list of files in a directory."""
     total_size = 0
@@ -162,24 +142,28 @@ def _scan_for_repo_folders(base_path: str) -> List[DownloadedModel]:
                     if os.path.isdir(subitem_path):
                         has_subdirs = True
                         # This looks like owner/repo structure
-                        size_bytes, files = _get_directory_size_and_files(subitem_path)
-                        if files:  # Only include if there are files
-                            # Check if the download is complete
-                            download_complete = _check_for_incomplete_downloads(subitem_path)
-                            # Load metadata if it exists
-                            metadata = _load_download_metadata(subitem_path)
-                            models.append(DownloadedModel(
-                                repo_id=f"{item}/{subitem}",
-                                files=files,
-                                folder_type='owner_repo',
-                                local_path=subitem_path,
-                                size_bytes=size_bytes,
-                                file_count=len(files),
-                                full_repo_download_complete=download_complete,
-                                pipeline_tag=metadata.get('pipeline_tag'),
-                                download_time=metadata.get('download_time'),
-                                avatar_url=metadata.get('avatar_url')
-                            ))
+                        # Only include if nexa.manifest exists
+                        manifest_path = os.path.join(subitem_path, 'nexa.manifest')
+                        if os.path.exists(manifest_path):
+                            size_bytes, files = _get_directory_size_and_files(subitem_path)
+                            if files:  # Only include if there are files
+                                # Check if the download is complete
+                                download_complete = _check_for_incomplete_downloads(subitem_path)
+                                # Load metadata if it exists
+                                repo_id = f"{item}/{subitem}"
+                                metadata = load_download_metadata(subitem_path, repo_id)
+                                models.append(DownloadedModel(
+                                    repo_id=repo_id,
+                                    files=files,
+                                    folder_type='owner_repo',
+                                    local_path=subitem_path,
+                                    size_bytes=size_bytes,
+                                    file_count=len(files),
+                                    full_repo_download_complete=download_complete,
+                                    pipeline_tag=metadata.get('pipeline_tag'),
+                                    download_time=metadata.get('download_time'),
+                                    avatar_url=metadata.get('avatar_url')
+                                ))
                     else:
                         direct_files.append(subitem)
             except (OSError, IOError):
@@ -188,24 +172,28 @@ def _scan_for_repo_folders(base_path: str) -> List[DownloadedModel]:
             
             # Direct repo folder (no owner structure)
             if not has_subdirs and direct_files:
-                size_bytes, files = _get_directory_size_and_files(item_path)
-                if files:  # Only include if there are files
-                    # Check if the download is complete
-                    download_complete = _check_for_incomplete_downloads(item_path)
-                    # Load metadata if it exists
-                    metadata = _load_download_metadata(item_path)
-                    models.append(DownloadedModel(
-                        repo_id=item,
-                        files=files,
-                        folder_type='direct_repo',
-                        local_path=item_path,
-                        size_bytes=size_bytes,
-                        file_count=len(files),
-                        full_repo_download_complete=download_complete,
-                        pipeline_tag=metadata.get('pipeline_tag'),
-                        download_time=metadata.get('download_time'),
-                        avatar_url=metadata.get('avatar_url')
-                    ))
+                # Only include if nexa.manifest exists
+                manifest_path = os.path.join(item_path, 'nexa.manifest')
+                if os.path.exists(manifest_path):
+                    size_bytes, files = _get_directory_size_and_files(item_path)
+                    if files:  # Only include if there are files
+                        # Check if the download is complete
+                        download_complete = _check_for_incomplete_downloads(item_path)
+                        # Load metadata if it exists
+                        repo_id = item
+                        metadata = load_download_metadata(item_path, repo_id)
+                        models.append(DownloadedModel(
+                            repo_id=repo_id,
+                            files=files,
+                            folder_type='direct_repo',
+                            local_path=item_path,
+                            size_bytes=size_bytes,
+                            file_count=len(files),
+                            full_repo_download_complete=download_complete,
+                            pipeline_tag=metadata.get('pipeline_tag'),
+                            download_time=metadata.get('download_time'),
+                            avatar_url=metadata.get('avatar_url')
+                        ))
     
     except (OSError, IOError):
         # Skip if base path can't be accessed
@@ -743,15 +731,26 @@ class HuggingFaceDownloader:
             # Get avatar URL
             avatar_url = get_avatar_url_for_repo(repo_id, custom_endpoint=self.endpoint)
             
-            # Prepare metadata
-            metadata = {
+            # Prepare old-style metadata for compatibility
+            old_metadata = {
                 'pipeline_tag': pipeline_tag,
                 'download_time': datetime.now().isoformat(),
                 'avatar_url': avatar_url
             }
             
-            # Save metadata to the repository directory
-            _save_download_metadata(local_dir, metadata)
+            # Get list of files in the directory
+            files = []
+            try:
+                for root, dirs, filenames in os.walk(local_dir):
+                    for filename in filenames:
+                        # Store relative path from the directory
+                        rel_path = os.path.relpath(os.path.join(root, filename), local_dir)
+                        files.append(rel_path)
+            except (OSError, IOError):
+                pass
+            
+            # Create and save manifest using the new utility function
+            save_manifest_with_files_metadata(repo_id, local_dir, old_metadata)
             
         except Exception:
             # Don't fail the download if metadata fetch fails
@@ -820,14 +819,6 @@ class HuggingFaceDownloader:
         # Create a subdirectory for this specific repo
         repo_local_dir = self._create_repo_directory(local_dir, repo_id)
         
-        # Check if repository already exists (basic check for directory existence)
-        if not force_download and os.path.exists(repo_local_dir) and os.listdir(repo_local_dir):
-            print(f"✓ Repository already exists, skipping: {repo_id}")
-            # Stop progress tracking
-            if progress_tracker:
-                progress_tracker.stop_tracking()
-            return repo_local_dir
-        
         try:
             download_kwargs = {
                 'repo_id': repo_id,

nexaai/utils/model_types.py

@@ -0,0 +1,47 @@
+"""
+Model type mappings for HuggingFace pipeline tags to our internal model types.
+
+This module provides centralized model type mapping functionality to avoid
+circular imports between other utility modules.
+"""
+
+from enum import Enum
+from typing import Dict
+
+
+class ModelTypeMapping(Enum):
+    """Enum for mapping HuggingFace pipeline_tag to our ModelType."""
+    TEXT_GENERATION = ("text-generation", "llm")
+    IMAGE_TEXT_TO_TEXT = ("image-text-to-text", "vlm")
+    
+    def __init__(self, pipeline_tag: str, model_type: str):
+        self.pipeline_tag = pipeline_tag
+        self.model_type = model_type
+
+
+# Create mapping dictionaries from the enum
+PIPELINE_TO_MODEL_TYPE: Dict[str, str] = {
+    mapping.pipeline_tag: mapping.model_type 
+    for mapping in ModelTypeMapping
+}
+
+MODEL_TYPE_TO_PIPELINE: Dict[str, str] = {
+    mapping.model_type: mapping.pipeline_tag 
+    for mapping in ModelTypeMapping
+}
+
+
+def map_pipeline_tag_to_model_type(pipeline_tag: str) -> str:
+    """Map HuggingFace pipeline_tag to our ModelType."""
+    if not pipeline_tag:
+        return "other"
+    
+    return PIPELINE_TO_MODEL_TYPE.get(pipeline_tag, "other")
+
+
+def map_model_type_to_pipeline_tag(model_type: str) -> str:
+    """Reverse map ModelType back to HuggingFace pipeline_tag."""
+    if not model_type:
+        return None
+    
+    return MODEL_TYPE_TO_PIPELINE.get(model_type)

nexaai/utils/progress_tracker.py

@@ -107,7 +107,7 @@ class DownloadProgressTracker:
             time_diff = current_time - self.last_time
             
             # Only calculate if we have a meaningful time difference (avoid division by very small numbers)
-            if time_diff > 0.5:  # At least 500ms between measurements
+            if time_diff > 0.1:  # At least 100ms between measurements
                 bytes_diff = current_downloaded - self.last_downloaded
                 
                 # Only calculate speed if bytes actually changed
@@ -118,6 +118,14 @@ class DownloadProgressTracker:
                     self.speed_history.append(speed)
                     if len(self.speed_history) > self.max_speed_history:
                         self.speed_history.pop(0)
+                    
+                    # Update tracking variables when we actually calculate speed
+                    self.last_downloaded = current_downloaded
+                    self.last_time = current_time
+        else:
+            # First measurement - initialize tracking variables
+            self.last_downloaded = current_downloaded
+            self.last_time = current_time
         
         # Return the average of historical speeds if we have any
         # This ensures we show the last known speed even when skipping updates
@@ -157,13 +165,9 @@ class DownloadProgressTracker:
                 total_file_sizes += data['total']
                 active_file_count += 1
         
-        # Calculate speed
+        # Calculate speed (tracking variables are updated internally)
         speed = self.calculate_speed(total_downloaded)
         
-        # Update tracking variables
-        self.last_downloaded = total_downloaded
-        self.last_time = time.time()
-        
         # Determine total size - prioritize pre-fetched repo size, then aggregate file sizes
         if self.total_repo_size > 0:
             # Use pre-fetched repository info if available