lemonade-sdk 8.1.11__py3-none-any.whl → 8.2.2__py3-none-any.whl

lemonade/tools/oga/load.py

@@ -38,6 +38,17 @@ execution_providers = {
 }
 
 
+def find_onnx_files_recursively(directory):
+    """
+    Recursively search for ONNX files in a directory and its subdirectories.
+    """
+    for _, _, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".onnx"):
+                return True
+    return False
+
+
 def _get_npu_driver_version():
     """
     Get the NPU driver version using PowerShell directly.
@@ -321,6 +332,7 @@ class OgaLoad(FirstTool):
 
     @staticmethod
     def _setup_model_dependencies(full_model_path, device, ryzenai_version, oga_path):
+        # pylint: disable=unused-argument
         """
         Sets up model dependencies for hybrid and NPU inference by:
         1. Configuring the custom_ops_library path in genai_config.json.
@@ -328,76 +340,45 @@ class OgaLoad(FirstTool):
         3. Check NPU driver version if required for device and ryzenai_version.
         """
 
-        env_path = sys.prefix
+        # For RyzenAI 1.6.0, check NPU driver version for NPU and hybrid devices
+        if device in ["npu", "hybrid"]:
+            required_driver_version = REQUIRED_NPU_DRIVER_VERSION
 
-        if "1.4.0" in ryzenai_version:
-            if device == "npu":
-                custom_ops_path = os.path.join(
-                    oga_path, "libs", "onnxruntime_vitis_ai_custom_ops.dll"
-                )
-            else:
-                custom_ops_path = os.path.join(oga_path, "libs", "onnx_custom_ops.dll")
-        else:
-            # For 1.5.0+, check NPU driver version for NPU and hybrid devices
-            if device in ["npu", "hybrid"]:
-                required_driver_version = REQUIRED_NPU_DRIVER_VERSION
-
-                current_driver_version = _get_npu_driver_version()
-
-                if not current_driver_version:
-                    printing.log_warning(
-                        f"NPU driver not found. {device.upper()} inference requires NPU driver "
-                        f"version {required_driver_version}.\n"
-                        "Please download and install the NPU Driver from:\n"
-                        f"{NPU_DRIVER_DOWNLOAD_URL}\n"
-                        "NPU functionality may not work properly."
-                    )
-                    _open_driver_install_page()
-
-                elif not _compare_driver_versions(
-                    current_driver_version, required_driver_version
-                ):
-                    printing.log_warning(
-                        f"Incorrect NPU driver version detected: {current_driver_version}\n"
-                        f"{device.upper()} inference with RyzenAI 1.5.0 requires driver "
-                        f"version {required_driver_version} or higher.\n"
-                        "Please download and install the correct NPU Driver from:\n"
-                        f"{NPU_DRIVER_DOWNLOAD_URL}\n"
-                        "NPU functionality may not work properly."
-                    )
-                    _open_driver_install_page()
-
-            if device == "npu":
-                # For 1.5.0, custom ops are in the conda environment's onnxruntime package
-                custom_ops_path = os.path.join(
-                    env_path,
-                    "Lib",
-                    "site-packages",
-                    "onnxruntime",
-                    "capi",
-                    "onnxruntime_vitis_ai_custom_ops.dll",
-                )
-                dll_source_path = os.path.join(
-                    env_path, "Lib", "site-packages", "onnxruntime", "capi"
-                )
-                required_dlls = ["dyn_dispatch_core.dll", "xaiengine.dll"]
-            else:
-                custom_ops_path = os.path.join(
-                    env_path,
-                    "Lib",
-                    "site-packages",
-                    "onnxruntime_genai",
-                    "onnx_custom_ops.dll",
+            current_driver_version = _get_npu_driver_version()
+            rai_version, _ = _get_ryzenai_version_info(device)
+
+            if not current_driver_version:
+                printing.log_warning(
+                    f"NPU driver not found. {device.upper()} inference requires NPU driver "
+                    f"version {required_driver_version}.\n"
+                    "Please download and install the NPU Driver from:\n"
+                    f"{NPU_DRIVER_DOWNLOAD_URL}\n"
+                    "NPU functionality may not work properly."
                 )
-                dll_source_path = os.path.join(
-                    env_path, "Lib", "site-packages", "onnxruntime_genai"
+                _open_driver_install_page()
+
+            elif not _compare_driver_versions(
+                current_driver_version, required_driver_version
+            ):
+                printing.log_warning(
+                    f"Incorrect NPU driver version detected: {current_driver_version}\n"
+                    f"{device.upper()} inference with RyzenAI {rai_version} requires driver "
+                    f"version {required_driver_version} or higher.\n"
+                    "Please download and install the correct NPU Driver from:\n"
+                    f"{NPU_DRIVER_DOWNLOAD_URL}\n"
+                    "NPU functionality may not work properly."
                 )
-                required_dlls = ["libutf8_validity.dll", "abseil_dll.dll"]
+                _open_driver_install_page()
+
+            # Setup DLL paths for NPU/hybrid inference
+            env_path = os.path.dirname(sys.executable)
+            dll_source_path = os.path.join(
+                env_path, "Lib", "site-packages", "onnxruntime_genai"
+            )
+            required_dlls = ["libutf8_validity.dll", "abseil_dll.dll"]
 
             # Validate that all required DLLs exist in the source directory
             missing_dlls = []
-            if not os.path.exists(custom_ops_path):
-                missing_dlls.append(custom_ops_path)
 
             for dll_name in required_dlls:
                 dll_source = os.path.join(dll_source_path, dll_name)
@@ -408,7 +389,9 @@ class OgaLoad(FirstTool):
                 dll_list = "\n  - ".join(missing_dlls)
                 raise RuntimeError(
                     f"Required DLLs not found for {device} inference:\n  - {dll_list}\n"
-                    f"Please ensure your RyzenAI installation is complete and supports {device}."
+                    f"Please ensure your RyzenAI installation is complete and supports {device}.\n"
+                    "Please reinstall the RyzenAI Software for your platform. Run:\n"
+                    "    pip install lemonade-sdk[oga-ryzenai]\n"
                 )
 
             # Add the DLL source directory to PATH
@@ -416,29 +399,6 @@ class OgaLoad(FirstTool):
             if dll_source_path not in current_path:
                 os.environ["PATH"] = dll_source_path + os.pathsep + current_path
 
-        # Update the model config with custom_ops_library path
-        config_path = os.path.join(full_model_path, "genai_config.json")
-        if os.path.exists(config_path):
-            with open(config_path, "r", encoding="utf-8") as f:
-                config = json.load(f)
-
-            if (
-                "model" in config
-                and "decoder" in config["model"]
-                and "session_options" in config["model"]["decoder"]
-            ):
-                config["model"]["decoder"]["session_options"][
-                    "custom_ops_library"
-                ] = custom_ops_path
-
-            with open(config_path, "w", encoding="utf-8") as f:
-                json.dump(config, f, indent=4)
-
-        else:
-            printing.log_info(
-                f"Model's `genai_config.json` not found in {full_model_path}"
-            )
-
     @staticmethod
     def _is_preoptimized_model(input_model_path):
         """
@@ -502,34 +462,6 @@ class OgaLoad(FirstTool):
 
         return full_model_path
 
-    @staticmethod
-    def _setup_npu_environment(ryzenai_version, oga_path):
-        """
-        Sets up environment for NPU flow of ONNX model and returns saved state to be restored
-        later in cleanup.
-        """
-        if "1.5.0" in ryzenai_version:
-            # For PyPI installation (1.5.0+), no environment setup needed
-            return None
-        elif "1.4.0" in ryzenai_version:
-            # Legacy lemonade-install approach for 1.4.0
-            if not os.path.exists(os.path.join(oga_path, "libs", "onnxruntime.dll")):
-                raise RuntimeError(
-                    f"Cannot find libs/onnxruntime.dll in lib folder: {oga_path}"
-                )
-
-            # Save current state so they can be restored after inference.
-            saved_state = {"cwd": os.getcwd(), "path": os.environ["PATH"]}
-
-            # Setup NPU environment (cwd and path will be restored later)
-            os.chdir(oga_path)
-            os.environ["PATH"] = (
-                os.path.join(oga_path, "libs") + os.pathsep + os.environ["PATH"]
-            )
-            return saved_state
-        else:
-            raise ValueError(f"Unsupported RyzenAI version: {ryzenai_version}")
-
     @staticmethod
     def _load_model_and_setup_state(
         state, full_model_path, checkpoint, trust_remote_code
@@ -702,8 +634,7 @@ class OgaLoad(FirstTool):
             state.save_stat(Keys.CHECKPOINT, checkpoint)
             state.save_stat(Keys.LOCAL_MODEL_FOLDER, full_model_path)
             # See if there is a file ending in ".onnx" in this folder
-            dir = os.listdir(input)
-            has_onnx_file = any([filename.endswith(".onnx") for filename in dir])
+            has_onnx_file = find_onnx_files_recursively(input)
             if not has_onnx_file:
                 raise ValueError(
                     f"The folder {input} does not contain an ONNX model file."
@@ -852,15 +783,10 @@ class OgaLoad(FirstTool):
 
             try:
                 if device == "npu":
-                    saved_env_state = self._setup_npu_environment(
-                        ryzenai_version, oga_path
-                    )
                     # Set USE_AIE_RoPE based on model type
                     os.environ["USE_AIE_RoPE"] = (
                         "0" if "phi-" in checkpoint.lower() else "1"
                     )
-                elif device == "hybrid":
-                    saved_env_state = None
 
                 self._load_model_and_setup_state(
                     state, full_model_path, checkpoint, trust_remote_code

lemonade/tools/oga/migration.py

@@ -0,0 +1,403 @@
+"""
+Migration utilities for handling RyzenAI version upgrades.
+
+This module provides functionality to detect and clean up incompatible RyzenAI models
+when upgrading between major versions (e.g., 1.4/1.5 -> 1.6).
+"""
+
+import os
+import json
+import shutil
+import logging
+from typing import List, Dict, Optional, Tuple
+
+
+def get_directory_size(path: str) -> int:
+    """
+    Calculate the total size of a directory in bytes.
+
+    Args:
+        path: Path to the directory
+
+    Returns:
+        Total size in bytes
+    """
+    total_size = 0
+    try:
+        for dirpath, _, filenames in os.walk(path):
+            for filename in filenames:
+                filepath = os.path.join(dirpath, filename)
+                try:
+                    total_size += os.path.getsize(filepath)
+                except (OSError, FileNotFoundError):
+                    # Skip files that can't be accessed
+                    pass
+    except (OSError, FileNotFoundError):
+        pass
+    return total_size
+
+
+def format_size(size_bytes: int) -> str:
+    """
+    Format byte size to human-readable string.
+
+    Args:
+        size_bytes: Size in bytes
+
+    Returns:
+        Formatted string (e.g., "1.5 GB", "450 MB")
+    """
+    for unit in ["B", "KB", "MB"]:
+        if size_bytes < 1024.0:
+            return f"{size_bytes:.1f} {unit}"
+        size_bytes /= 1024.0
+    return f"{size_bytes:.1f} GB"
+
+
+def check_rai_config_version(model_path: str, required_version: str = "1.6.0") -> bool:
+    """
+    Check if a model's rai_config.json contains the required version.
+
+    Args:
+        model_path: Path to the model directory
+        required_version: Version string to check for (default: "1.6.0")
+
+    Returns:
+        True if model is compatible (has required version), False otherwise
+    """
+    rai_config_path = os.path.join(model_path, "rai_config.json")
+
+    # If no rai_config.json exists, it's not a RyzenAI model
+    if not os.path.exists(rai_config_path):
+        return True
+
+    try:
+        with open(rai_config_path, "r", encoding="utf-8") as f:
+            config = json.load(f)
+
+        # Check if max_prompt_length exists and has the required version
+        if "max_prompt_length" in config:
+            max_prompt_length = config["max_prompt_length"]
+            if isinstance(max_prompt_length, dict):
+                # If it's a dict with version keys, check for required version
+                return required_version in max_prompt_length
+            # Fallback to True to avoid deleting models if format changes
+            return True
+
+        return True
+
+    except (json.JSONDecodeError, OSError) as e:
+        logging.warning(f"Could not read rai_config.json from {model_path}: {e}")
+        # If we can't read it, assume it's compatible to avoid false positives
+        return True
+
+
+def scan_oga_models_cache(cache_dir: str) -> List[Dict[str, any]]:
+    """
+    Scan the Lemonade OGA models cache for incompatible models.
+
+    Args:
+        cache_dir: Path to the Lemonade cache directory
+
+    Returns:
+        List of dicts with model info (path, name, size, compatible)
+    """
+    oga_models_path = os.path.join(cache_dir, "oga_models")
+    incompatible_models = []
+
+    if not os.path.exists(oga_models_path):
+        return incompatible_models
+
+    try:
+        # Iterate through model directories in oga_models
+        for model_name in os.listdir(oga_models_path):
+            model_dir = os.path.join(oga_models_path, model_name)
+
+            if not os.path.isdir(model_dir):
+                continue
+
+            # Check all subdirectories (e.g., npu-int4, hybrid-int4)
+            for subdir in os.listdir(model_dir):
+                subdir_path = os.path.join(model_dir, subdir)
+
+                if not os.path.isdir(subdir_path):
+                    continue
+
+                # Check if this model version is compatible
+                if not check_rai_config_version(subdir_path):
+                    size = get_directory_size(subdir_path)
+                    incompatible_models.append(
+                        {
+                            "path": subdir_path,
+                            "name": f"{model_name}/{subdir}",
+                            "size": size,
+                            "size_formatted": format_size(size),
+                            "cache_type": "lemonade",
+                        }
+                    )
+
+    except (OSError, PermissionError) as e:
+        logging.warning(f"Error scanning oga_models cache: {e}")
+
+    return incompatible_models
+
+
+def scan_huggingface_cache(hf_home: Optional[str] = None) -> List[Dict[str, any]]:
+    """
+    Scan the HuggingFace cache for incompatible RyzenAI models.
+
+    Args:
+        hf_home: Path to HuggingFace home directory (default: from env or ~/.cache/huggingface)
+
+    Returns:
+        List of dicts with model info (path, name, size, compatible)
+    """
+    if hf_home is None:
+        hf_home = os.environ.get(
+            "HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
+        )
+
+    hub_path = os.path.join(hf_home, "hub")
+    incompatible_models = []
+
+    if not os.path.exists(hub_path):
+        return incompatible_models
+
+    try:
+        # Iterate through model directories in HuggingFace cache
+        for item in os.listdir(hub_path):
+            if not item.startswith("models--"):
+                continue
+
+            model_dir = os.path.join(hub_path, item)
+            if not os.path.isdir(model_dir):
+                continue
+
+            # Look in snapshots subdirectory
+            snapshots_dir = os.path.join(model_dir, "snapshots")
+            if not os.path.exists(snapshots_dir):
+                continue
+
+            # Check each snapshot
+            for snapshot_hash in os.listdir(snapshots_dir):
+                snapshot_path = os.path.join(snapshots_dir, snapshot_hash)
+
+                if not os.path.isdir(snapshot_path):
+                    continue
+
+                # Check if this snapshot has incompatible RyzenAI model
+                if not check_rai_config_version(snapshot_path):
+                    # Extract readable model name from directory
+                    model_name = item.replace("models--", "").replace("--", "/")
+                    size = get_directory_size(
+                        model_dir
+                    )  # Size of entire model directory
+                    incompatible_models.append(
+                        {
+                            "path": model_dir,
+                            "name": model_name,
+                            "size": size,
+                            "size_formatted": format_size(size),
+                            "cache_type": "huggingface",
+                        }
+                    )
+                    break
+
+    except (OSError, PermissionError) as e:
+        logging.warning(f"Error scanning HuggingFace cache: {e}")
+
+    return incompatible_models
+
+
+def detect_incompatible_ryzenai_models(
+    cache_dir: str, hf_home: Optional[str] = None
+) -> Tuple[List[Dict[str, any]], int]:
+    """
+    Detect all incompatible RyzenAI models in both Lemonade and HuggingFace caches.
+
+    Args:
+        cache_dir: Path to the Lemonade cache directory
+        hf_home: Path to HuggingFace home directory (optional)
+
+    Returns:
+        Tuple of (list of incompatible models, total size in bytes)
+    """
+    incompatible_models = []
+
+    # Scan Lemonade cache
+    oga_models = scan_oga_models_cache(cache_dir)
+    incompatible_models.extend(oga_models)
+
+    # Scan HuggingFace cache
+    hf_models = scan_huggingface_cache(hf_home)
+    incompatible_models.extend(hf_models)
+
+    # Calculate total size
+    total_size = sum(model["size"] for model in incompatible_models)
+
+    logging.info(
+        f"Found {len(incompatible_models)} incompatible RyzenAI models "
+        f"({format_size(total_size)} total)"
+    )
+
+    return incompatible_models, total_size
+
+
+def delete_model_directory(model_path: str) -> bool:
+    """
+    Safely delete a model directory.
+
+    Args:
+        model_path: Path to the model directory to delete
+
+    Returns:
+        True if deletion successful, False otherwise
+    """
+    try:
+        if os.path.exists(model_path):
+            shutil.rmtree(model_path)
+            logging.info(f"Deleted model directory: {model_path}")
+            return True
+        else:
+            logging.warning(f"Model directory not found: {model_path}")
+            return False
+    except (OSError, PermissionError) as e:
+        logging.error(f"Failed to delete model directory {model_path}: {e}")
+        return False
+
+
+def _extract_checkpoint_from_path(path: str) -> Optional[str]:
+    """
+    Extract the checkpoint name from a model path.
+
+    Args:
+        path: Model directory path (either Lemonade cache or HuggingFace cache)
+
+    Returns:
+        Checkpoint name (e.g., "amd/Qwen2.5-1.5B-Instruct-awq") or None if not extractable
+    """
+    # Normalize path separators to handle both Unix and Windows paths
+    normalized_path = path.replace("\\", "/")
+    parts = normalized_path.split("/")
+
+    # Handle HuggingFace cache paths: models--{org}--{repo}
+    if "models--" in normalized_path:
+        for part in parts:
+            if part.startswith("models--"):
+                # Convert models--org--repo to org/repo
+                # Replace first two occurrences of -- with /
+                checkpoint = part.replace("models--", "", 1).replace("--", "/", 1)
+                return checkpoint
+        return None
+
+    # Handle Lemonade cache paths: oga_models/{model_name}/{device}-{dtype}
+    if "oga_models" in normalized_path:
+        try:
+            oga_models_idx = parts.index("oga_models")
+            if oga_models_idx + 1 < len(parts):
+                model_name = parts[oga_models_idx + 1]
+                # Convert model_name back to checkpoint (e.g., amd_model -> amd/model)
+                # This is a heuristic - we look for the pattern {org}_{model}
+                checkpoint = model_name.replace("_", "/", 1)
+                return checkpoint
+        except (ValueError, IndexError):
+            return None
+
+    return None
+
+
+def _cleanup_user_models_json(deleted_checkpoints: List[str], user_models_file: str):
+    """
+    Remove entries from user_models.json for models that have been deleted.
+
+    Args:
+        deleted_checkpoints: List of checkpoint names that were deleted
+        user_models_file: Path to user_models.json
+    """
+    if not deleted_checkpoints or not os.path.exists(user_models_file):
+        return
+
+    try:
+        with open(user_models_file, "r", encoding="utf-8") as f:
+            user_models = json.load(f)
+
+        # Track which models to remove
+        models_to_remove = []
+        for model_name, model_info in user_models.items():
+            checkpoint = model_info.get("checkpoint", "")
+            # Check if this checkpoint matches any deleted checkpoints
+            # We do a case-insensitive comparison since paths may have been lowercased
+            for deleted_checkpoint in deleted_checkpoints:
+                if checkpoint.lower() == deleted_checkpoint.lower():
+                    models_to_remove.append(model_name)
+                    break
+
+        # Remove the models
+        for model_name in models_to_remove:
+            del user_models[model_name]
+            logging.info(f"Removed {model_name} from user_models.json")
+
+        # Save the updated file only if we removed something
+        if models_to_remove:
+            with open(user_models_file, "w", encoding="utf-8") as f:
+                json.dump(user_models, f, indent=2)
+            logging.info(
+                f"Updated user_models.json - removed {len(models_to_remove)} entries"
+            )
+
+    except (json.JSONDecodeError, OSError) as e:
+        logging.warning(f"Could not update user_models.json: {e}")
+
+
+def delete_incompatible_models(
+    model_paths: List[str], user_models_file: Optional[str] = None
+) -> Dict[str, any]:
+    """
+    Delete multiple incompatible model directories and clean up user_models.json.
+
+    Args:
+        model_paths: List of paths to delete
+        user_models_file: Path to user_models.json (optional, will use default if not provided)
+
+    Returns:
+        Dict with deletion results (success_count, failed_count, freed_size, cleaned_user_models)
+    """
+    success_count = 0
+    failed_count = 0
+    freed_size = 0
+    deleted_checkpoints = []
+
+    for path in model_paths:
+        # Calculate size before deletion
+        size = get_directory_size(path)
+
+        # Extract checkpoint name before deleting
+        checkpoint = _extract_checkpoint_from_path(path)
+        if checkpoint:
+            deleted_checkpoints.append(checkpoint)
+
+        if delete_model_directory(path):
+            success_count += 1
+            freed_size += size
+        else:
+            failed_count += 1
+
+    # Clean up user_models.json if we deleted any models
+    cleaned_user_models = False
+    if deleted_checkpoints:
+        # Use default path if not provided
+        if user_models_file is None:
+            from lemonade.cache import DEFAULT_CACHE_DIR
+
+            user_models_file = os.path.join(DEFAULT_CACHE_DIR, "user_models.json")
+
+        _cleanup_user_models_json(deleted_checkpoints, user_models_file)
+        cleaned_user_models = True
+
+    return {
+        "success_count": success_count,
+        "failed_count": failed_count,
+        "freed_size": freed_size,
+        "freed_size_formatted": format_size(freed_size),
+        "cleaned_user_models": cleaned_user_models,
+    }