From 3ab4abcb57c6c1a6dc0e6a8aa031f6f06c5e7800 Mon Sep 17 00:00:00 2001 From: hualxie Date: Mon, 25 May 2026 14:28:19 +0800 Subject: [PATCH 1/3] add ruff_check --- .aitk/scripts/auto_formatter.py | 212 -------------------------------- .aitk/scripts/requirements.txt | 9 +- .aitk/scripts/ruff_check.py | 83 +++++++++++++ .aitk/scripts/sanitize.py | 6 +- 4 files changed, 90 insertions(+), 220 deletions(-) delete mode 100644 .aitk/scripts/auto_formatter.py create mode 100644 .aitk/scripts/ruff_check.py diff --git a/.aitk/scripts/auto_formatter.py b/.aitk/scripts/auto_formatter.py deleted file mode 100644 index 425d3ba63..000000000 --- a/.aitk/scripts/auto_formatter.py +++ /dev/null @@ -1,212 +0,0 @@ -#!/usr/bin/env python3 -""" -Auto-formatter module for Python scripts. -Provides comprehensive formatting capabilities including: -- Removing unused imports using autoflake -- Sorting and organizing imports using isort -- Formatting code using black with 120 character line length -- Checking that all imports are at the top of files -""" - -import subprocess -import sys -from pathlib import Path - -from sanitize.utils import printError, printInfo, printTip, printWarning - - -def install_formatter_tools(): - """ - Install required formatting tools if not available. - """ - tools = [ - ("black", "black"), - ("isort", "isort"), - ("autoflake", "autoflake"), # Added autoflake for removing unused imports - ] - - for tool_name, package_name in tools: - try: - subprocess.run([tool_name, "--version"], check=True, capture_output=True) - except (subprocess.CalledProcessError, FileNotFoundError): - printInfo(f"Installing {package_name} formatter...") - try: - subprocess.run([sys.executable, "-m", "pip", "install", package_name], check=True) - except subprocess.CalledProcessError as e: - printError(f"Failed to install {package_name}: {e}") - return False - return True - - -def check_imports_at_top(file_path): - """ - Check if all imports are at the top of the file (after docstring and comments). - Returns True if imports are properly placed, False otherwise. - """ - try: - with open(file_path, "r", encoding="utf-8") as f: - lines = f.readlines() - - # Skip shebang, encoding declarations, and docstrings - in_docstring = False - docstring_quotes = None - non_import_code_found = False - - for i, line in enumerate(lines): - stripped = line.strip() - - # Skip empty lines and comments - if not stripped or stripped.startswith("#"): - continue - - # Handle docstrings - if not in_docstring and (stripped.startswith('"""') or stripped.startswith("'''")): - docstring_quotes = stripped[:3] - if stripped.count(docstring_quotes) >= 2: - # Single line docstring - continue - else: - in_docstring = True - continue - elif in_docstring and docstring_quotes and docstring_quotes in stripped: - in_docstring = False - continue - elif in_docstring: - continue - - # Check for imports and from statements (including multi-line imports) - if ( - stripped.startswith("import ") - or stripped.startswith("from ") - or (not non_import_code_found and (stripped.endswith(",") or stripped.startswith(")"))) - ): - if non_import_code_found: - printWarning(f"Import found after non-import code in {file_path}:{i+1}") - return False - else: - # Non-import code found (but ignore special variables and sys.path modifications) - if ( - stripped - and not stripped.startswith("__") - and not stripped.startswith("sys.path") - and not any(special in stripped for special in ["__all__", "__version__", "__author__"]) - ): - non_import_code_found = True - - return True - except Exception as e: - printError(f"Error checking imports in {file_path}: {e}") - return True # Don't fail the entire process - - -def auto_format_scripts(target_dir=None): - """ - Auto-format all Python scripts in the target directory with comprehensive formatting: - - Remove unused imports using autoflake - - Sort and organize imports using isort - - Format code using black with 120 character line length - - Check that all imports are at the top of files - - Args: - target_dir: Path to the directory to format. If None, uses the scripts directory. - """ - if target_dir is None: - target_dir = Path(__file__).parent - else: - target_dir = Path(target_dir) - - printTip(f"Auto-formatting Python scripts in {target_dir}...") - - # Find all Python files in the target directory - python_files = [] - for py_file in target_dir.rglob("*.py"): - if py_file.is_file(): - python_files.append(str(py_file)) - - if not python_files: - printInfo("No Python files found to format.") - return - - # Install required tools - if not install_formatter_tools(): - printError("Failed to install required formatting tools.") - return - - # Step 1: Remove unused imports with autoflake - printInfo("Step 1: Removing unused imports...") - try: - autoflake_cmd = [ - "autoflake", - "--in-place", # Modify files in place - "--remove-all-unused-imports", # Remove all unused imports - "--remove-unused-variables", # Remove unused variables - "--remove-duplicate-keys", # Remove duplicate keys in dictionaries - "--ignore-init-module-imports", # Don't remove imports in __init__.py files - ] + python_files - - result = subprocess.run(autoflake_cmd, capture_output=True, text=True) - if result.returncode == 0: - printInfo(f"Successfully removed unused imports from {len(python_files)} files.") - if result.stdout: - printInfo(result.stdout) - else: - printError(f"Autoflake failed: {result.stderr}") - except Exception as e: - printError(f"Error during unused import removal: {e}") - - # Step 2: Sort imports with isort - printInfo("Step 2: Sorting and organizing imports...") - try: - isort_cmd = [ - "isort", - "--line-length", - "120", - "--multi-line", - "3", - "--trailing-comma", - "--force-grid-wrap", - "0", - "--combine-as", - "--use-parentheses", - ] + python_files - - result = subprocess.run(isort_cmd, capture_output=True, text=True) - if result.returncode == 0: - printInfo(f"Successfully sorted imports in {len(python_files)} files.") - else: - printError(f"Import sorting failed: {result.stderr}") - except Exception as e: - printError(f"Error during import sorting: {e}") - - # Step 3: Check import placement - printInfo("Step 3: Checking import placement...") - for py_file in python_files: - check_imports_at_top(py_file) - - # Step 4: Format with black - printInfo("Step 4: Formatting code with black...") - try: - black_cmd = ["black", "--line-length", "120"] + python_files - result = subprocess.run(black_cmd, capture_output=True, text=True) - - if result.returncode == 0: - printInfo(f"Successfully formatted {len(python_files)} Python files with 120 character line length.") - if result.stdout: - printInfo(result.stdout) - else: - printError(f"Black formatting failed: {result.stderr}") - except Exception as e: - printError(f"Error during black formatting: {e}") - - printInfo("Auto-formatting completed!") - - # Clear sanitize modules from cache - modules_to_clear = [name for name in sys.modules.keys() if name.startswith("sanitize")] - for module_name in modules_to_clear: - if module_name in sys.modules: - del sys.modules[module_name] - - -if __name__ == "__main__": - # Allow running this module directly for testing - auto_format_scripts() diff --git a/.aitk/scripts/requirements.txt b/.aitk/scripts/requirements.txt index 71b436aeb..be3e6caa5 100644 --- a/.aitk/scripts/requirements.txt +++ b/.aitk/scripts/requirements.txt @@ -1,4 +1,5 @@ -deepdiff -pydantic -pydash -pyyaml +deepdiff==9.1.0 +pydantic==2.13.4 +pydash==8.0.6 +pyyaml==6.0.3 +ruff==0.15.14 diff --git a/.aitk/scripts/ruff_check.py b/.aitk/scripts/ruff_check.py new file mode 100644 index 000000000..38bb3046a --- /dev/null +++ b/.aitk/scripts/ruff_check.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +""" +Ruff-based lint and format check for Python files in `.aitk/` and every `*/aitk/` folder. + +By default, runs `ruff check` and `ruff format --check` (verification only) and exits 1 on +issues. When called with `fix=True`, runs `ruff check --fix` and `ruff format` to rewrite +files in place. +""" + +import subprocess +import sys +from pathlib import Path + +from sanitize.utils import printError, printInfo, printTip + +LINE_LENGTH = "120" +SELECT_RULES = "F401,F841,I,E402" +REQUIREMENTS_HINT = ".aitk/scripts/requirements.txt" + + +def _ensure_ruff(): + try: + subprocess.run(["ruff", "--version"], check=True, capture_output=True) + return True + except (subprocess.CalledProcessError, FileNotFoundError): + printError( + f"ruff is not installed. Install dependencies first: pip install -r {REQUIREMENTS_HINT}" + ) + return False + + +def _collect_targets(): + repo_root = Path(__file__).parent.parent.parent + + targets = [] + dot_aitk = repo_root / ".aitk" + if dot_aitk.is_dir(): + targets.append(dot_aitk) + + for aitk_dir in sorted(repo_root.glob("*/aitk")): + if aitk_dir.is_dir(): + targets.append(aitk_dir) + + return targets + + +def ruff_check(fix: bool = False): + """ + Run ruff lint + format across `.aitk/` and every `*/aitk/` folder. + + Args: + fix: When True, apply autofixes and reformat files. When False, only verify. + """ + if not _ensure_ruff(): + raise SystemExit(1) + + targets = _collect_targets() + if not targets: + printInfo("No .aitk / aitk folders found for ruff to check.") + return + + target_args = [str(t) for t in targets] + printTip(f"Running ruff on {len(targets)} folder(s)...") + + check_cmd = ["ruff", "check", "--select", SELECT_RULES, "--line-length", LINE_LENGTH] + format_cmd = ["ruff", "format", "--line-length", LINE_LENGTH] + if fix: + check_cmd.append("--fix") + else: + format_cmd.append("--check") + + check_result = subprocess.run(check_cmd + target_args) + format_result = subprocess.run(format_cmd + target_args) + + if check_result.returncode != 0 or format_result.returncode != 0: + printError("Ruff reported issues. Re-run `python sanitize.py --format_scripts` to auto-fix.") + raise SystemExit(1) + + printInfo("Ruff check passed.") + + +if __name__ == "__main__": + ruff_check(fix="--fix" in sys.argv) diff --git a/.aitk/scripts/sanitize.py b/.aitk/scripts/sanitize.py index a76693bac..300d90f2a 100644 --- a/.aitk/scripts/sanitize.py +++ b/.aitk/scripts/sanitize.py @@ -9,9 +9,9 @@ import sys from pathlib import Path -from auto_formatter import auto_format_scripts from project_processor import project_processor from requirements_check import requirements_check +from ruff_check import ruff_check from sanitize.main import main from sanitize.utils import GlobalVars @@ -54,9 +54,7 @@ def run_main(): if "--fill_pipeline_tags" in sys.argv: GlobalVars.fillPipelineTags = True - # Auto-format scripts before running sanitize - if "--format_scripts" in sys.argv: - auto_format_scripts() + ruff_check(fix=True) requirements_check() project_processor() run_main() From b277230ed8a5f31d580ea78b3ad3abcbae35472d Mon Sep 17 00:00:00 2001 From: hualxie Date: Mon, 25 May 2026 14:36:45 +0800 Subject: [PATCH 2/3] format python codes --- .aitk/scripts/bump_model_versions.py | 3 +- .aitk/scripts/install_freeze.py | 2 +- .aitk/scripts/project_processor.py | 13 +- .aitk/scripts/requirements_check.py | 1 + .aitk/scripts/ruff_check.py | 11 +- .aitk/scripts/sanitize/file_validation.py | 2 +- .aitk/scripts/sanitize/generator_amd.py | 4 +- .aitk/scripts/sanitize/generator_intel.py | 2 +- .aitk/scripts/sanitize/generator_qnn.py | 4 +- .aitk/scripts/sanitize/model_parameter.py | 20 +- .aitk/scripts/sanitize/parameters.py | 4 +- .../aitk/openai_clip_ov.py | 33 ++-- .../aitk/user_script.py | 7 +- .../aitk/vit-base-patch16-224.py | 14 +- .../aitk/user_script.py | 3 +- .../aitk/laion_clip_ov.py | 5 +- .../aitk/user_script.py | 3 +- microsoft-resnet-50/aitk/imagenet-qnn-gpu.py | 3 +- microsoft-resnet-50/aitk/imagenet.py | 13 +- .../aitk/openai_clip_ov.py | 5 +- .../aitk/user_script.py | 3 +- .../aitk/openai_clip_ov.py | 5 +- .../aitk/user_script.py | 3 +- .../aitk/openai_clip_ov.py | 5 +- .../aitk/user_script.py | 3 +- .../aitk/convert_whisper_to_ovir.py | 137 ++++++++------ .../aitk/ov_evaluate.py | 14 +- .../aitk/ov_workflow.py | 53 ++++-- openai-whisper-large-v3-turbo/aitk/qnn_app.py | 7 +- .../aitk/qnn_evaluate.py | 17 +- openai-whisper-large-v3-turbo/aitk/qnn_run.py | 14 +- .../aitk/qnn_workflow.py | 132 ++++++++----- .../aitk/whisper_decoder_load.py | 7 +- .../aitk/whisper_encoder_load.py | 9 +- sam-vit-base/aitk/config.py | 1 + sam-vit-base/aitk/model_patches.py | 2 +- sam-vit-base/aitk/sam_mask_generator.py | 22 ++- sam-vit-base/aitk/sam_qnn_evaluation.py | 23 +-- sam-vit-base/aitk/sam_qnn_workflow.py | 94 +++++++--- sam-vit-base/aitk/user_script.py | 2 +- sam2.1-hiera-small/aitk/generate_model.py | 8 +- .../aitk/sam2_mask_generator.py | 39 +--- .../aitk/sam2_qnn_evaluation.py | 31 ++- sam2.1-hiera-small/aitk/sam2_qnn_workflow.py | 94 +++++++--- .../aitk/evaluation.py | 2 +- .../aitk/sd_ov_evaluation.py | 12 +- .../aitk/sd_ov_workflow.py | 76 +++++--- .../aitk/sd_qnn_evaluation.py | 12 +- .../aitk/sd_qnn_workflow.py | 165 ++++++++++------ .../aitk/sd_utils/onnx_patch.py | 8 +- .../aitk/sd_utils/ort.py | 9 +- .../aitk/sd_utils/ov.py | 10 +- .../aitk/sd_utils/qdq.py | 4 +- .../aitk/evaluation.py | 2 +- .../aitk/sd_ov_evaluation.py | 12 +- .../aitk/sd_ov_workflow.py | 90 +++++---- .../aitk/sd_qnn_evaluation.py | 12 +- .../aitk/sd_qnn_workflow.py | 177 ++++++++++++------ .../aitk/sd_utils/onnx_patch.py | 8 +- .../aitk/sd_utils/ort.py | 9 +- .../aitk/sd_utils/ov.py | 10 +- .../aitk/sd_utils/qdq.py | 4 +- 62 files changed, 898 insertions(+), 601 deletions(-) diff --git a/.aitk/scripts/bump_model_versions.py b/.aitk/scripts/bump_model_versions.py index dedb34ea1..1c5a5f4fd 100644 --- a/.aitk/scripts/bump_model_versions.py +++ b/.aitk/scripts/bump_model_versions.py @@ -6,7 +6,6 @@ from sanitize.utils import iter_aitk_info_yml - VERSION_LINE_RE = re.compile(r"^(?P\s*)version:\s*(?P\d+)\s*(?P#.*)?$") @@ -72,7 +71,7 @@ def bump_file(yml_file: Path, yaml_object: dict, delta: int, dry_run: bool) -> O print(f"Skip {yml_file}: could not locate 'version: {old}' line under aitk.modelInfo") return None m = VERSION_LINE_RE.match(lines[idx].rstrip("\n")) - newline = lines[idx][len(lines[idx].rstrip("\r\n")):] or "\n" + newline = lines[idx][len(lines[idx].rstrip("\r\n")) :] or "\n" indent = m.group("indent") comment = m.group("comment") rebuilt = f"{indent}version: {new}" diff --git a/.aitk/scripts/install_freeze.py b/.aitk/scripts/install_freeze.py index 0a2ba0fae..121134e80 100644 --- a/.aitk/scripts/install_freeze.py +++ b/.aitk/scripts/install_freeze.py @@ -219,7 +219,7 @@ def main(): # Install print(f"Installing dependencies: {temp_req}") - result = subprocess.run(["uv", "pip", "install", "-r", temp_req, "-p", args.python], text=True) + subprocess.run(["uv", "pip", "install", "-r", temp_req, "-p", args.python], text=True) # Get freeze pip_freeze = subprocess.check_output(["uv", "pip", "freeze", "-p", args.python]).decode("utf-8").splitlines() diff --git a/.aitk/scripts/project_processor.py b/.aitk/scripts/project_processor.py index 4fb5a09c6..f50b5a333 100644 --- a/.aitk/scripts/project_processor.py +++ b/.aitk/scripts/project_processor.py @@ -15,14 +15,15 @@ from sanitize.model_info import ModelInfo, ModelList from sanitize.project_config import ModelInfoProject, ModelProjectConfig, WorkflowItem from sanitize.utils import ( - GlobalVars, WINML_COPY_EXEMPT_IDS, + GlobalVars, isLLM_by_id, iter_aitk_info_yml, open_ex, winml_copy_src_for, ) + def fetch_pipeline_tags(model_link: str) -> Optional[List[str]]: """Fetch pipeline_tag from HuggingFace API for a given model link. @@ -32,7 +33,7 @@ def fetch_pipeline_tags(model_link: str) -> Optional[List[str]]: hf_prefix = "https://huggingface.co/" if not model_link.startswith(hf_prefix): return None - model_id = model_link[len(hf_prefix):].rstrip("/") + model_id = model_link[len(hf_prefix) :].rstrip("/") if not model_id: return None url = f"https://huggingface.co/api/models/{model_id}" @@ -93,7 +94,7 @@ def write_list( root_dir: Path, md_path: Path, ): - modelList.sort(key=lambda x: (x.modelName)) + modelList.sort(key=lambda x: x.modelName) f.write(f"## {title}\n\n") f.write("| Model Name | Supported Runtimes |\n") f.write("|------------|--------------------|\n") @@ -231,9 +232,7 @@ def project_processor(): modelList.models.append(modelInfo) # copy pre — auto-ensure winml.py copy entry (unless exempt), then run pre-phase copies copyConfigFile = yml_file.parent / "_copy.json.config" - copyConfig: CopyConfig | None = ( - CopyConfig.Read(copyConfigFile.as_posix()) if copyConfigFile.exists() else None - ) + copyConfig: CopyConfig | None = CopyConfig.Read(copyConfigFile.as_posix()) if copyConfigFile.exists() else None if modelInfo.id not in WINML_COPY_EXEMPT_IDS: desired_src = winml_copy_src_for(modelInfo.id) if copyConfig is None: @@ -262,7 +261,7 @@ def project_processor(): # project config and json configs convert_yaml_to_project_config(yml_file, yaml_object, modelList, model_summary) - modelList.models.sort(key=lambda x: (x.GetSortKey())) + modelList.models.sort(key=lambda x: x.GetSortKey()) modelList.writeIfChanged() all_summary.write(root_dir) diff --git a/.aitk/scripts/requirements_check.py b/.aitk/scripts/requirements_check.py index 0d4e088bf..8fe5d460a 100644 --- a/.aitk/scripts/requirements_check.py +++ b/.aitk/scripts/requirements_check.py @@ -62,5 +62,6 @@ def requirements_check(): req_is_subset(nvidia_autogptq, general_cuda_autogptq) req_is_subset(general_cuda_autogptq, nvidia_autogptq) + if __name__ == "__main__": requirements_check() diff --git a/.aitk/scripts/ruff_check.py b/.aitk/scripts/ruff_check.py index 38bb3046a..7f5e7c916 100644 --- a/.aitk/scripts/ruff_check.py +++ b/.aitk/scripts/ruff_check.py @@ -14,7 +14,7 @@ from sanitize.utils import printError, printInfo, printTip LINE_LENGTH = "120" -SELECT_RULES = "F401,F841,I,E402" +SELECT_RULES = "F401,F841,I" REQUIREMENTS_HINT = ".aitk/scripts/requirements.txt" @@ -23,9 +23,7 @@ def _ensure_ruff(): subprocess.run(["ruff", "--version"], check=True, capture_output=True) return True except (subprocess.CalledProcessError, FileNotFoundError): - printError( - f"ruff is not installed. Install dependencies first: pip install -r {REQUIREMENTS_HINT}" - ) + printError(f"ruff is not installed. Install dependencies first: pip install -r {REQUIREMENTS_HINT}") return False @@ -62,8 +60,9 @@ def ruff_check(fix: bool = False): target_args = [str(t) for t in targets] printTip(f"Running ruff on {len(targets)} folder(s)...") - check_cmd = ["ruff", "check", "--select", SELECT_RULES, "--line-length", LINE_LENGTH] - format_cmd = ["ruff", "format", "--line-length", LINE_LENGTH] + exclude_ipynb = ["--config", 'extend-exclude=["*.ipynb"]'] + check_cmd = ["ruff", "check", "--select", SELECT_RULES, "--line-length", LINE_LENGTH, *exclude_ipynb] + format_cmd = ["ruff", "format", "--line-length", LINE_LENGTH, *exclude_ipynb] if fix: check_cmd.append("--fix") else: diff --git a/.aitk/scripts/sanitize/file_validation.py b/.aitk/scripts/sanitize/file_validation.py index fa87e3fa9..02e39cc98 100644 --- a/.aitk/scripts/sanitize/file_validation.py +++ b/.aitk/scripts/sanitize/file_validation.py @@ -204,7 +204,7 @@ def readCheckIpynb(ipynbFile: str, modelItems: dict[str, ModelParameter]): printError(f"{ipynbFile} has multiple runtimes, but also has {targetStr}") foundTargetEP = True elif targetCount > 1: - printError(f'{ipynbFile} should have 1 {targetStr} but has {targetCount}') + printError(f"{ipynbFile} should have 1 {targetStr} but has {targetCount}") if not foundTargetEP: printError(f"{ipynbFile} has no runtime for it!") return True diff --git a/.aitk/scripts/sanitize/generator_amd.py b/.aitk/scripts/sanitize/generator_amd.py index 2f28c63ec..d10519f78 100644 --- a/.aitk/scripts/sanitize/generator_amd.py +++ b/.aitk/scripts/sanitize/generator_amd.py @@ -2,7 +2,7 @@ from pathlib import Path from typing import Optional -from .constants import OlivePassNames, OlivePropertyNames, PhaseTypeEnum, EPNames +from .constants import EPNames, OlivePassNames, OlivePropertyNames, PhaseTypeEnum from .generator_common import create_model_parameter, set_optimization_path from .model_info import ModelList from .model_parameter import ModelParameter, OptimizationPath, Section @@ -239,7 +239,7 @@ def generator_amd(id: str, recipe, folder: Path, modelList: ModelList): return amd_runtime_values: list[str] = recipe.get("devices", [recipe.get("device")]) - name = f"Convert to AMD {"/".join([runtime.upper() for runtime in amd_runtime_values])}" + name = f"Convert to AMD {'/'.join([runtime.upper() for runtime in amd_runtime_values])}" parameter = create_model_parameter(aitk, name, configFile) parameter.isLLM = isLLM diff --git a/.aitk/scripts/sanitize/generator_intel.py b/.aitk/scripts/sanitize/generator_intel.py index e55921096..38b315c84 100644 --- a/.aitk/scripts/sanitize/generator_intel.py +++ b/.aitk/scripts/sanitize/generator_intel.py @@ -77,7 +77,7 @@ def generator_intel(id: str, recipe, folder: Path): return intel_runtime_values: list[str] = recipe.get("devices", [recipe.get("device")]) - name = f"Convert to Intel {"/".join([runtime.upper() for runtime in intel_runtime_values])}" + name = f"Convert to Intel {'/'.join([runtime.upper() for runtime in intel_runtime_values])}" parameter = create_model_parameter(aitk, name, configFile) parameter.isLLM = isLLM diff --git a/.aitk/scripts/sanitize/generator_qnn.py b/.aitk/scripts/sanitize/generator_qnn.py index ed1cda7e3..0447028af 100644 --- a/.aitk/scripts/sanitize/generator_qnn.py +++ b/.aitk/scripts/sanitize/generator_qnn.py @@ -1,5 +1,5 @@ -from pathlib import Path import json +from pathlib import Path from .constants import OlivePassNames, OlivePropertyNames from .generator_amd import generate_quantization_config @@ -40,7 +40,7 @@ def generator_qnn(id: str, recipe, folder: Path, modelList: ModelList): return runtime_values: list[str] = recipe.get("devices", [recipe.get("device")]) - name = f"Convert to Qualcomm {"/".join([runtime.upper() for runtime in runtime_values])}" + name = f"Convert to Qualcomm {'/'.join([runtime.upper() for runtime in runtime_values])}" parameter = create_model_parameter(aitk, name, configFile) if "npu" in runtime_values: diff --git a/.aitk/scripts/sanitize/model_parameter.py b/.aitk/scripts/sanitize/model_parameter.py index e4c055cd6..0229b7f92 100644 --- a/.aitk/scripts/sanitize/model_parameter.py +++ b/.aitk/scripts/sanitize/model_parameter.py @@ -328,11 +328,7 @@ def Check(self, templates: Dict[str, Parameter], oliveJson: Any, modelList: Mode evaluateUsedInExecute=True, ) - if ( - not modelInfo.template - and not modelInfo.extension - and currentEp != EPNames.OpenVINOExecutionProvider.value - ): + if not modelInfo.template and not modelInfo.extension and currentEp != EPNames.OpenVINOExecutionProvider.value: if self.runtimeOverwrite is None: self.runtimeOverwrite = RuntimeOverwrite( autoGenerated=True, @@ -645,7 +641,11 @@ def _diffOlivePasses(self, localLabel: str, refPasses: Any, localPasses: Any, re changeds: dict[str, Any] = diff.pop("values_changed", {}) newChangeds = {} for changed in changeds: - if not (changed.endswith("['data_config']") or changed.endswith("['user_script']") or changed.endswith("['save_as_external_data']")): + if not ( + changed.endswith("['data_config']") + or changed.endswith("['user_script']") + or changed.endswith("['save_as_external_data']") + ): newChangeds[changed] = changeds[changed] if newChangeds: diff["values_changed"] = newChangeds @@ -692,7 +692,9 @@ def checkOliveFile(self, oliveJson: Any, modelInfo: ModelInfo): with open_ex(refPath, "r") as f: refJson = json.load(f) label = "/".join(Path(self._file).parts[-3:-1]) + "/" + localFile - self._diffOlivePasses(label, refJson[OlivePropertyNames.Passes], localJson[OlivePropertyNames.Passes], refFile) + self._diffOlivePasses( + label, refJson[OlivePropertyNames.Passes], localJson[OlivePropertyNames.Passes], refFile + ) return if not oliveFileRef: @@ -721,7 +723,9 @@ def checkOliveFile(self, oliveJson: Any, modelInfo: ModelInfo): with open_ex(oliveFile, "r") as file: oliveFileJson = json.load(file) label = "/".join(Path(self._file if self._file else "UNKNOWN").parts[-3:]) - self._diffOlivePasses(label, oliveFileJson[OlivePropertyNames.Passes], oliveJson[OlivePropertyNames.Passes], oliveFileRef) + self._diffOlivePasses( + label, oliveFileJson[OlivePropertyNames.Passes], oliveJson[OlivePropertyNames.Passes], oliveFileRef + ) def checkDebugInfo(self, oliveJson: Any): self.debugInfo = DebugInfo() diff --git a/.aitk/scripts/sanitize/parameters.py b/.aitk/scripts/sanitize/parameters.py index a781cc7db..34f486a4b 100644 --- a/.aitk/scripts/sanitize/parameters.py +++ b/.aitk/scripts/sanitize/parameters.py @@ -269,6 +269,8 @@ def readCheckParameterTemplate(filePath: str): for key, parameter in parameters.items(): if not parameter.Check(True): printError(f"{filePath} parameter {key} has error") - newContent = json.dumps(adapter.dump_python(parameters, exclude_none=True, mode="json"), indent=4, ensure_ascii=True) + newContent = json.dumps( + adapter.dump_python(parameters, exclude_none=True, mode="json"), indent=4, ensure_ascii=True + ) BaseModelClass.writeJsonIfChanged(newContent, filePath, fileContent) return parameters diff --git a/OFA-Sys-chinese-clip-vit-base-patch16/aitk/openai_clip_ov.py b/OFA-Sys-chinese-clip-vit-base-patch16/aitk/openai_clip_ov.py index 6a8262b14..27f32698e 100644 --- a/OFA-Sys-chinese-clip-vit-base-patch16/aitk/openai_clip_ov.py +++ b/OFA-Sys-chinese-clip-vit-base-patch16/aitk/openai_clip_ov.py @@ -1,18 +1,18 @@ +import tarfile +import zipfile from io import BytesIO from pathlib import Path +import numpy as np import requests import torch from huggingface_hub import hf_hub_download +from olive.data.registry import Registry from PIL import Image from requests.packages.urllib3.exceptions import InsecureRequestWarning +from torch.utils.data import Dataset from tqdm import tqdm from transformers import ChineseCLIPModel, ChineseCLIPProcessor -import tarfile -import zipfile -from olive.data.registry import Registry -from torch.utils.data import Dataset -import numpy as np requests.packages.urllib3.disable_warnings(InsecureRequestWarning) @@ -47,7 +47,7 @@ def open(self): """Manually open ZIP file""" if self.zip_path and Path(self.zip_path).exists() and not self.is_open: try: - self.zip_file = zipfile.ZipFile(self.zip_path, 'r') + self.zip_file = zipfile.ZipFile(self.zip_path, "r") self.is_open = True print(f"Opened ZIP: {self.zip_path}") return True @@ -74,7 +74,7 @@ def read_file(self, filename): if not self.zip_file: return None try: - with self.zip_file.open(filename, 'r') as file: + with self.zip_file.open(filename, "r") as file: return file.read() except Exception as e: print(f"Error reading {filename}: {e}") @@ -85,7 +85,7 @@ def load_image(self, filename): image_data = self.read_file(filename) if image_data: try: - return Image.open(BytesIO(image_data)).convert('RGB') + return Image.open(BytesIO(image_data)).convert("RGB") except Exception as e: print(f"Error creating image from {filename}: {e}") return None @@ -100,8 +100,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): def wrap_collate_fn(processor, max_length): def collate_fn(image, chinese_caption: str): - """Preprocess an example by loading and transforming image and text data. - """ + """Preprocess an example by loading and transforming image and text data.""" inputs = processor(text=chinese_caption, images=[image], return_tensors="pt", padding=True) if inputs["input_ids"].shape[1] > max_length: return None @@ -151,21 +150,21 @@ def get_coco_cn(target_folder, split="train") -> list[list[str]]: # Download only the tar.gz file print("Downloading coco-cn-version1805v1.1.tar.gz...") tar_path = hf_hub_download( - repo_id="AIMClab-RUC/COCO-CN", - filename="coco-cn-version1805v1.1.tar.gz", - repo_type="dataset" + repo_id="AIMClab-RUC/COCO-CN", filename="coco-cn-version1805v1.1.tar.gz", repo_type="dataset" ) # Extract the tar.gz file print(f"Extracting to {target_folder}...") - with tarfile.open(tar_path, 'r:gz') as tar: + with tarfile.open(tar_path, "r:gz") as tar: tar.extractall(path=target_folder) print("Extraction completed!") with open(target_folder / "coco-cn-version1805v1.1" / f"coco-cn_{split}.txt", "r", encoding="utf-8") as f: lines = f.readlines() - with open(target_folder / "coco-cn-version1805v1.1" / "imageid.human-written-caption.txt", "r", encoding="utf-8") as f: + with open( + target_folder / "coco-cn-version1805v1.1" / "imageid.human-written-caption.txt", "r", encoding="utf-8" + ) as f: images_lines = f.readlines() image_caption_dict = {} for line in images_lines: @@ -201,7 +200,7 @@ def conceptual_captions_dataset(data_name, opt_init_steps=200, **kwargs): model_path = kwargs.get("model_path") if not model_path: - raise ValueError( + raise ValueError( "The 'model_path' parameter is required in data_configs.load_dataset_config but was not provided." ) model = ChineseCLIPModel.from_pretrained(model_path) @@ -222,6 +221,7 @@ def custom_transform_func(data_item): # Evaluation + class CLIPDataset(Dataset): def __init__( self, @@ -259,6 +259,7 @@ def __del__(self): self.cocoVal.close() if hasattr(self, "cocoTrain") and self.cocoTrain is not None: self.cocoTrain.close() + def __len__(self): return self.length diff --git a/google-bert-bert-base-multilingual-cased/aitk/user_script.py b/google-bert-bert-base-multilingual-cased/aitk/user_script.py index f7442c2fc..e25f535e8 100644 --- a/google-bert-bert-base-multilingual-cased/aitk/user_script.py +++ b/google-bert-bert-base-multilingual-cased/aitk/user_script.py @@ -5,9 +5,8 @@ import datasets import numpy as np import torch -from transformers import BertTokenizer - from olive.data.registry import Registry +from transformers import BertTokenizer # ------------------------------------------------------------------------- # Common Dataset @@ -46,7 +45,9 @@ @Registry.register_dataset() def bert_base_multilingual_cased_dataset(data_name, split, max_samples): # load the raw wikipedia dataset for tuning. Load just 300 examples for speed. - raw_dataset = datasets.load_dataset(data_name, "20220301.en", split=f"{split}[:{max_samples}]", trust_remote_code=True) + raw_dataset = datasets.load_dataset( + data_name, "20220301.en", split=f"{split}[:{max_samples}]", trust_remote_code=True + ) def _preprocess_fn(examples): return tokenizer( diff --git a/google-vit-base-patch16-224/aitk/vit-base-patch16-224.py b/google-vit-base-patch16-224/aitk/vit-base-patch16-224.py index 09b0027c3..ccee10b79 100644 --- a/google-vit-base-patch16-224/aitk/vit-base-patch16-224.py +++ b/google-vit-base-patch16-224/aitk/vit-base-patch16-224.py @@ -6,20 +6,21 @@ from pathlib import Path import numpy as np -import torchvision.transforms as transforms import transformers +from olive.data.registry import Registry from torch import from_numpy from torch.utils.data import Dataset -from olive.data.registry import Registry - logger = getLogger(__name__) + def get_imagenet_label_map(): import json + cache_file = Path(f"./cache/data/imagenet_class_index.json") if not cache_file.exists(): import requests + imagenet_class_index_url = ( "https://raw.githubusercontent.com/pytorch/vision/main/gallery/assets/imagenet_class_index.json" ) @@ -35,10 +36,12 @@ def get_imagenet_label_map(): return {v[0]: int(k) for k, v in content.items()} + def adapt_label_for_mini_imagenet(labels: list, label_names: list): label_map = get_imagenet_label_map() return [label_map[label_names[x]] for x in labels] + class ImagenetDataset(Dataset): def __init__(self, data): self.images = from_numpy(data["images"]) @@ -59,9 +62,12 @@ def dataset_post_process(output): else output.argmax(axis=1) ) + from transformers import AutoImageProcessor + processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224", use_fast=True) + @Registry.register_pre_process() def dataset_pre_process(output_data, **kwargs): shuffle = kwargs.get("shuffle", True) @@ -89,7 +95,7 @@ def dataset_pre_process(output_data, **kwargs): images.append(image) labels.append(label) - if(output_data.info.dataset_name == "mini-imagenet"): + if output_data.info.dataset_name == "mini-imagenet": labels = adapt_label_for_mini_imagenet(labels, output_data.features["label"].names) result_data = ImagenetDataset({"images": np.array(images), "labels": np.array(labels)}) diff --git a/intel-bert-base-uncased-mrpc/aitk/user_script.py b/intel-bert-base-uncased-mrpc/aitk/user_script.py index ea05d21e0..5cb2df5db 100644 --- a/intel-bert-base-uncased-mrpc/aitk/user_script.py +++ b/intel-bert-base-uncased-mrpc/aitk/user_script.py @@ -5,9 +5,8 @@ import datasets import numpy as np import torch -from transformers import BertTokenizer - from olive.data.registry import Registry +from transformers import BertTokenizer # ------------------------------------------------------------------------- # Common Dataset diff --git a/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/laion_clip_ov.py b/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/laion_clip_ov.py index d1971b501..8dbfd3b50 100644 --- a/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/laion_clip_ov.py +++ b/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/laion_clip_ov.py @@ -3,13 +3,12 @@ import requests import torch from datasets import load_dataset +from olive.data.registry import Registry from PIL import Image from requests.packages.urllib3.exceptions import InsecureRequestWarning from tqdm import tqdm from transformers import CLIPModel, CLIPProcessor -from olive.data.registry import Registry - requests.packages.urllib3.disable_warnings(InsecureRequestWarning) # ------------------------------------------------------------------------- @@ -99,7 +98,7 @@ def prepare_calibration_data(dataloader, init_steps): @Registry.register_dataset() -def conceptual_captions_dataset(data_name,opt_init_steps=200, max_train_samples=1000, **kwargs): +def conceptual_captions_dataset(data_name, opt_init_steps=200, max_train_samples=1000, **kwargs): """Prepare a vision-text dataset for quantization.""" dataset = load_dataset(data_name, trust_remote_code=True) model_path = kwargs.get("model_path") diff --git a/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/user_script.py b/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/user_script.py index 83d75d06b..82379b592 100644 --- a/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/user_script.py +++ b/laion-CLIP-ViT-B-32-laion2B-s34B-b79K/aitk/user_script.py @@ -5,11 +5,10 @@ import numpy as np import torch from datasets import load_dataset +from olive.data.registry import Registry from torch.utils.data import Dataset from transformers import CLIPProcessor -from olive.data.registry import Registry - class CLIPDataset(Dataset): def __init__( diff --git a/microsoft-resnet-50/aitk/imagenet-qnn-gpu.py b/microsoft-resnet-50/aitk/imagenet-qnn-gpu.py index c888734ec..06b51c7dd 100644 --- a/microsoft-resnet-50/aitk/imagenet-qnn-gpu.py +++ b/microsoft-resnet-50/aitk/imagenet-qnn-gpu.py @@ -8,11 +8,10 @@ import numpy as np import torchvision.transforms as transforms import transformers +from olive.data.registry import Registry from torch import from_numpy from torch.utils.data import Dataset -from olive.data.registry import Registry - logger = getLogger(__name__) diff --git a/microsoft-resnet-50/aitk/imagenet.py b/microsoft-resnet-50/aitk/imagenet.py index e2cb70d3b..b64c02508 100644 --- a/microsoft-resnet-50/aitk/imagenet.py +++ b/microsoft-resnet-50/aitk/imagenet.py @@ -6,20 +6,21 @@ from pathlib import Path import numpy as np -import torchvision.transforms as transforms import transformers +from olive.data.registry import Registry from torch import from_numpy, permute from torch.utils.data import Dataset -from olive.data.registry import Registry - logger = getLogger(__name__) + def get_imagenet_label_map(): import json + cache_file = Path(f"./cache/data/imagenet_class_index.json") if not cache_file.exists(): import requests + imagenet_class_index_url = ( "https://raw.githubusercontent.com/pytorch/vision/main/gallery/assets/imagenet_class_index.json" ) @@ -35,10 +36,12 @@ def get_imagenet_label_map(): return {v[0]: int(k) for k, v in content.items()} + def adapt_label_for_mini_imagenet(labels: list, label_names: list): label_map = get_imagenet_label_map() return [label_map[label_names[x]] for x in labels] + class ImagenetDataset(Dataset): def __init__(self, data): self.images = from_numpy(data["images"]) @@ -61,8 +64,10 @@ def dataset_post_process(output): from transformers import AutoImageProcessor + processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50", use_fast=True) + @Registry.register_pre_process() def dataset_pre_process(output_data, **kwargs): shuffle = kwargs.get("shuffle", True) @@ -94,7 +99,7 @@ def dataset_pre_process(output_data, **kwargs): images.append(image) labels.append(label) - if(output_data.info.dataset_name == "mini-imagenet"): + if output_data.info.dataset_name == "mini-imagenet": labels = adapt_label_for_mini_imagenet(labels, output_data.features["label"].names) result_data = ImagenetDataset({"images": np.array(images), "labels": np.array(labels)}) diff --git a/openai-clip-vit-base-patch16/aitk/openai_clip_ov.py b/openai-clip-vit-base-patch16/aitk/openai_clip_ov.py index d1971b501..8dbfd3b50 100644 --- a/openai-clip-vit-base-patch16/aitk/openai_clip_ov.py +++ b/openai-clip-vit-base-patch16/aitk/openai_clip_ov.py @@ -3,13 +3,12 @@ import requests import torch from datasets import load_dataset +from olive.data.registry import Registry from PIL import Image from requests.packages.urllib3.exceptions import InsecureRequestWarning from tqdm import tqdm from transformers import CLIPModel, CLIPProcessor -from olive.data.registry import Registry - requests.packages.urllib3.disable_warnings(InsecureRequestWarning) # ------------------------------------------------------------------------- @@ -99,7 +98,7 @@ def prepare_calibration_data(dataloader, init_steps): @Registry.register_dataset() -def conceptual_captions_dataset(data_name,opt_init_steps=200, max_train_samples=1000, **kwargs): +def conceptual_captions_dataset(data_name, opt_init_steps=200, max_train_samples=1000, **kwargs): """Prepare a vision-text dataset for quantization.""" dataset = load_dataset(data_name, trust_remote_code=True) model_path = kwargs.get("model_path") diff --git a/openai-clip-vit-base-patch16/aitk/user_script.py b/openai-clip-vit-base-patch16/aitk/user_script.py index 83d75d06b..82379b592 100644 --- a/openai-clip-vit-base-patch16/aitk/user_script.py +++ b/openai-clip-vit-base-patch16/aitk/user_script.py @@ -5,11 +5,10 @@ import numpy as np import torch from datasets import load_dataset +from olive.data.registry import Registry from torch.utils.data import Dataset from transformers import CLIPProcessor -from olive.data.registry import Registry - class CLIPDataset(Dataset): def __init__( diff --git a/openai-clip-vit-base-patch32/aitk/openai_clip_ov.py b/openai-clip-vit-base-patch32/aitk/openai_clip_ov.py index d1971b501..8dbfd3b50 100644 --- a/openai-clip-vit-base-patch32/aitk/openai_clip_ov.py +++ b/openai-clip-vit-base-patch32/aitk/openai_clip_ov.py @@ -3,13 +3,12 @@ import requests import torch from datasets import load_dataset +from olive.data.registry import Registry from PIL import Image from requests.packages.urllib3.exceptions import InsecureRequestWarning from tqdm import tqdm from transformers import CLIPModel, CLIPProcessor -from olive.data.registry import Registry - requests.packages.urllib3.disable_warnings(InsecureRequestWarning) # ------------------------------------------------------------------------- @@ -99,7 +98,7 @@ def prepare_calibration_data(dataloader, init_steps): @Registry.register_dataset() -def conceptual_captions_dataset(data_name,opt_init_steps=200, max_train_samples=1000, **kwargs): +def conceptual_captions_dataset(data_name, opt_init_steps=200, max_train_samples=1000, **kwargs): """Prepare a vision-text dataset for quantization.""" dataset = load_dataset(data_name, trust_remote_code=True) model_path = kwargs.get("model_path") diff --git a/openai-clip-vit-base-patch32/aitk/user_script.py b/openai-clip-vit-base-patch32/aitk/user_script.py index 83d75d06b..82379b592 100644 --- a/openai-clip-vit-base-patch32/aitk/user_script.py +++ b/openai-clip-vit-base-patch32/aitk/user_script.py @@ -5,11 +5,10 @@ import numpy as np import torch from datasets import load_dataset +from olive.data.registry import Registry from torch.utils.data import Dataset from transformers import CLIPProcessor -from olive.data.registry import Registry - class CLIPDataset(Dataset): def __init__( diff --git a/openai-clip-vit-large-patch14/aitk/openai_clip_ov.py b/openai-clip-vit-large-patch14/aitk/openai_clip_ov.py index d1971b501..8dbfd3b50 100644 --- a/openai-clip-vit-large-patch14/aitk/openai_clip_ov.py +++ b/openai-clip-vit-large-patch14/aitk/openai_clip_ov.py @@ -3,13 +3,12 @@ import requests import torch from datasets import load_dataset +from olive.data.registry import Registry from PIL import Image from requests.packages.urllib3.exceptions import InsecureRequestWarning from tqdm import tqdm from transformers import CLIPModel, CLIPProcessor -from olive.data.registry import Registry - requests.packages.urllib3.disable_warnings(InsecureRequestWarning) # ------------------------------------------------------------------------- @@ -99,7 +98,7 @@ def prepare_calibration_data(dataloader, init_steps): @Registry.register_dataset() -def conceptual_captions_dataset(data_name,opt_init_steps=200, max_train_samples=1000, **kwargs): +def conceptual_captions_dataset(data_name, opt_init_steps=200, max_train_samples=1000, **kwargs): """Prepare a vision-text dataset for quantization.""" dataset = load_dataset(data_name, trust_remote_code=True) model_path = kwargs.get("model_path") diff --git a/openai-clip-vit-large-patch14/aitk/user_script.py b/openai-clip-vit-large-patch14/aitk/user_script.py index 83d75d06b..82379b592 100644 --- a/openai-clip-vit-large-patch14/aitk/user_script.py +++ b/openai-clip-vit-large-patch14/aitk/user_script.py @@ -5,11 +5,10 @@ import numpy as np import torch from datasets import load_dataset +from olive.data.registry import Registry from torch.utils.data import Dataset from transformers import CLIPProcessor -from olive.data.registry import Registry - class CLIPDataset(Dataset): def __init__( diff --git a/openai-whisper-large-v3-turbo/aitk/convert_whisper_to_ovir.py b/openai-whisper-large-v3-turbo/aitk/convert_whisper_to_ovir.py index cde25a402..4f99e1ba2 100644 --- a/openai-whisper-large-v3-turbo/aitk/convert_whisper_to_ovir.py +++ b/openai-whisper-large-v3-turbo/aitk/convert_whisper_to_ovir.py @@ -2,10 +2,10 @@ # Copyright (c) Intel® Corporation. All rights reserved. # Licensed under the MIT License. # -------------------------------------------------------------------------- -import os -import shutil import argparse import json +import os +import shutil from transformers import WhisperConfig @@ -17,6 +17,7 @@ def run_olive_workflow(config_json: dict): @param config_json: JSON dictionary representing Olive workflow configuration. """ from olive.workflows import run + run(config_json) @@ -30,12 +31,13 @@ def str2bool(v: str) -> bool: """ if isinstance(v, bool): return v - if v.lower() in ("yes", "true", 't', 'y', '1'): + if v.lower() in ("yes", "true", "t", "y", "1"): return True - elif v.lower() in ("no", "false", 'f', 'n', '0'): + elif v.lower() in ("no", "false", "f", "n", "0"): return False else: - raise argparse.ArgumentTypeError('Boolean value expected.') + raise argparse.ArgumentTypeError("Boolean value expected.") + def handle_arguments() -> argparse.Namespace: """ @@ -51,38 +53,38 @@ def handle_arguments() -> argparse.Namespace: parser.add_argument("--cache_dir", required=True, type=str) parser.add_argument( - '-m', + "-m", "--model", type=str, default="openai/whisper-large-v3-turbo", - help="Huggingface ID for whisper model to convert. Default: openai/whisper-large-v3-turbo" + help="Huggingface ID for whisper model to convert. Default: openai/whisper-large-v3-turbo", ) parser.add_argument( - '-w', + "-w", "--weight-format", type=str, default="fp16", choices=["fp16", "int8", "int4"], - help="The weight format used to compress the model. Default: fp16. Other options: int8, int4" + help="The weight format used to compress the model. Default: fp16. Other options: int8, int4", ) parser.add_argument( "--enable_npu_ws", type=str2bool, - nargs='?', + nargs="?", const=True, default=False, help="Enable use of NPUW weight sharing between prefill & generate compiled models. " - "Can be used as a flag (--enable_npu_ws) or with a value (--enable_npu_ws True). Default: False" + "Can be used as a flag (--enable_npu_ws) or with a value (--enable_npu_ws True). Default: False", ) parser.add_argument( "--reshape", type=str2bool, - nargs='?', + nargs="?", const=True, default=False, help="Reshape encoder & decoder models at conversion time, instead of adding reshape_input " - "provider options to genai_config.json. " - "Can be used as a flag (--reshape) or with a value (--reshape True). Default: False" + "provider options to genai_config.json. " + "Can be used as a flag (--reshape) or with a value (--reshape True). Default: False", ) parser.add_argument( "--device", @@ -109,17 +111,17 @@ def load_templates() -> list[dict]: # load default config json default_config_path = os.path.join(script_dir, "whisper_large_v3_turbo_default_ov_npu.json") - with open(default_config_path, 'r') as f: + with open(default_config_path, "r") as f: default_json = json.load(f) # load the default encapsulation json default_encapsulation_path = os.path.join(script_dir, "whisper_large_v3_turbo_encapsulate.json") - with open(default_encapsulation_path, 'r') as f: + with open(default_encapsulation_path, "r") as f: default_encapsulation_json = json.load(f) # load the default audio preprocessor config json audio_processor_config_path = os.path.join(script_dir, "audio_processor_config_default.json") - with open(audio_processor_config_path, 'r') as f: + with open(audio_processor_config_path, "r") as f: audio_processor_config_json = json.load(f) return [default_json, default_encapsulation_json, audio_processor_config_json] @@ -135,6 +137,7 @@ def reshape_and_save_to_tmp(input_ov_path: str, reshape_dict: dict, tmp_ov_path: """ try: from openvino.runtime import Core, serialize + core = Core() ov_model = core.read_model(input_ov_path) ov_model.reshape(reshape_dict) @@ -142,13 +145,21 @@ def reshape_and_save_to_tmp(input_ov_path: str, reshape_dict: dict, tmp_ov_path: except (ImportError, AttributeError): # fallback to the newer 2025+ version of openvino that doesn't have openvino.runtime import openvino as ov + core = ov.Core() ov_model = core.read_model(input_ov_path) ov_model.reshape(reshape_dict) ov.serialize(ov_model, tmp_ov_path) -def update_genai_overrides(encapsulate_json, w_config: WhisperConfig, enable_npu_ws: bool, reshape: bool, output_directory: str, device: str = "npu"): +def update_genai_overrides( + encapsulate_json, + w_config: WhisperConfig, + enable_npu_ws: bool, + reshape: bool, + output_directory: str, + device: str = "npu", +): """ Update genai_config overrides based on NPU_WS and reshape flags. @@ -182,7 +193,7 @@ def update_genai_overrides(encapsulate_json, w_config: WhisperConfig, enable_npu "NPUW_WHISPER": "YES", "NPUW_WEIGHTS_BANK": "whisper-shared", "NPUW_LLM_PREFILL_HINT": "STATIC", - "NPUW_ONLINE_PIPELINE": "NONE" + "NPUW_ONLINE_PIPELINE": "NONE", } } else: @@ -194,20 +205,16 @@ def update_genai_overrides(encapsulate_json, w_config: WhisperConfig, enable_npu "NPUW_FOLD": "NO", "NPUW_WHISPER": "YES", "NPUW_LLM_PREFILL_HINT": "STATIC", - "NPUW_ONLINE_PIPELINE": "NONE" + "NPUW_ONLINE_PIPELINE": "NONE", } } else: decoder_load_config = {} # Only stringify the load_config value, not the whole provider_options - decoder_load_config_str = json.dumps(decoder_load_config, separators=(',', ':')) + decoder_load_config_str = json.dumps(decoder_load_config, separators=(",", ":")) - provider_options_encoder = { - "OpenVINO": { - "device_type": device.upper() - } - } + provider_options_encoder = {"OpenVINO": {"device_type": device.upper()}} # This 3000 magic number is common across all official whisper models. # It could potentially be calculated by loading preprocessor_config.json and performing the calculation. But just fix it at 3000 for now. @@ -226,15 +233,21 @@ def update_genai_overrides(encapsulate_json, w_config: WhisperConfig, enable_npu if reshape: print(" [OK] Reshaping decoder model now...") - reshape_and_save_to_tmp(os.path.join(output_directory, "openvino_decoder_model.xml"), - {"encoder_hidden_states": encoder_hidden_states_shape}, "reshaped_tmp.xml") + reshape_and_save_to_tmp( + os.path.join(output_directory, "openvino_decoder_model.xml"), + {"encoder_hidden_states": encoder_hidden_states_shape}, + "reshaped_tmp.xml", + ) os.replace("reshaped_tmp.xml", os.path.join(output_directory, "openvino_decoder_model.xml")) os.replace("reshaped_tmp.bin", os.path.join(output_directory, "openvino_decoder_model.bin")) print(" [OK] Reshaping encoder model now...") - reshape_and_save_to_tmp(os.path.join(output_directory, "openvino_encoder_model.xml"), - {"input_features": input_features_shape}, "reshaped_tmp.xml") + reshape_and_save_to_tmp( + os.path.join(output_directory, "openvino_encoder_model.xml"), + {"input_features": input_features_shape}, + "reshaped_tmp.xml", + ) os.replace("reshaped_tmp.xml", os.path.join(output_directory, "openvino_encoder_model.xml")) os.replace("reshaped_tmp.bin", os.path.join(output_directory, "openvino_encoder_model.bin")) @@ -245,11 +258,7 @@ def update_genai_overrides(encapsulate_json, w_config: WhisperConfig, enable_npu # Build decoder provider_options as proper JSON object (only load_config is a string) provider_options_decoder = { - "OpenVINO": { - "device_type": device.upper(), - "enable_causallm": "True", - "load_config": decoder_load_config_str - } + "OpenVINO": {"device_type": device.upper(), "enable_causallm": "True", "load_config": decoder_load_config_str} } if encoder_hidden_states_shape_str: @@ -278,30 +287,22 @@ def post_process_genai_config(genai_config_path: str, provider_options_encoder: @param provider_options_encoder: Encoder provider options dict @param w_config: Whisper configuration object """ - with open(genai_config_path, 'r') as f: + with open(genai_config_path, "r") as f: genai_config = json.load(f) - max_length = w_config.max_target_positions # Add encoder section with all required fields genai_config["model"]["encoder"] = { - "session_options": { - "log_id": "onnxruntime-genai", - "provider_options": [provider_options_encoder] - }, + "session_options": {"log_id": "onnxruntime-genai", "provider_options": [provider_options_encoder]}, "filename": "openvino_encoder_model.onnx", "head_size": w_config.d_model // w_config.encoder_attention_heads, "hidden_size": w_config.d_model, - "inputs": { - "audio_features": "input_features" - }, - "outputs": { - "encoder_hidden_states": "last_hidden_state" - }, + "inputs": {"audio_features": "input_features"}, + "outputs": {"encoder_hidden_states": "last_hidden_state"}, "num_attention_heads": w_config.encoder_attention_heads, "num_hidden_layers": w_config.encoder_layers, - "num_key_value_heads": w_config.encoder_attention_heads + "num_key_value_heads": w_config.encoder_attention_heads, } # Fix decoder section - the genai_config may have been generated from encoder encapsulation @@ -320,15 +321,10 @@ def post_process_genai_config(genai_config_path: str, provider_options_encoder: decoder["num_key_value_heads"] = w_config.decoder_attention_heads # Fix inputs (decoder takes input_ids and encoder_hidden_states) - decoder["inputs"] = { - "input_ids": "input_ids", - "encoder_hidden_states": "encoder_hidden_states" - } + decoder["inputs"] = {"input_ids": "input_ids", "encoder_hidden_states": "encoder_hidden_states"} # Fix outputs (decoder outputs logits) - decoder["outputs"] = { - "logits": "logits" - } + decoder["outputs"] = {"logits": "logits"} # Remove graph_optimization_level if present (not needed) if "session_options" in decoder and "graph_optimization_level" in decoder["session_options"]: @@ -342,11 +338,19 @@ def post_process_genai_config(genai_config_path: str, provider_options_encoder: genai_config["model"]["context_length"] = max_length genai_config["search"]["max_length"] = max_length - with open(genai_config_path, 'w') as f: + with open(genai_config_path, "w") as f: json.dump(genai_config, f, indent=4) -def run_encapsulation(encapsulation_config_json, output_directory: str, w_config: WhisperConfig, enable_npu_ws: bool, reshape: bool, cache_dir: str = "cache", device: str = "npu"): +def run_encapsulation( + encapsulation_config_json, + output_directory: str, + w_config: WhisperConfig, + enable_npu_ws: bool, + reshape: bool, + cache_dir: str = "cache", + device: str = "npu", +): """ Run Olive encapsulation workflow for Whisper models. @@ -365,7 +369,7 @@ def run_encapsulation(encapsulation_config_json, output_directory: str, w_config for run_hash in os.listdir(cache_models_dir): run_json_path = os.path.join(cache_models_dir, run_hash, "run.json") if os.path.exists(run_json_path): - with open(run_json_path, 'r') as f: + with open(run_json_path, "r") as f: run_info = json.load(f) if not run_info.get("pass_name", "").startswith("openvino"): continue @@ -432,7 +436,9 @@ def run_encapsulation(encapsulation_config_json, output_directory: str, w_config raise FileNotFoundError(f"Could not find {model_name}.xml in output directory") # Update genai overrides based on NPU_WS and reshape flags - provider_options_encoder = update_genai_overrides(encapsulation_config_json, w_config, enable_npu_ws, reshape, ".", device=device) + provider_options_encoder = update_genai_overrides( + encapsulation_config_json, w_config, enable_npu_ws, reshape, ".", device=device + ) # Store absolute path to output directory for genai_config post-processing output_dir_abs = os.path.abspath(".") @@ -527,7 +533,9 @@ def main(): weight_format = args.weight_format enable_npu_ws = args.enable_npu_ws reshape = args.reshape - print(f"Converting model: {hf_model_id} with weight format: {weight_format}, enable_npu_ws: {enable_npu_ws}, reshape: {reshape}") + print( + f"Converting model: {hf_model_id} with weight format: {weight_format}, enable_npu_ws: {enable_npu_ws}, reshape: {reshape}" + ) # load default JSONs [default_json, default_encapsulation_json, audio_processor_config_json] = load_templates() @@ -553,7 +561,9 @@ def main(): default_json["passes"]["optimum_convert"]["ov_quant_config"]["weight_format"] = weight_format # override the mel spectrogram config in audio preprocessor config - audio_processor_config_json["feature_extraction"]["sequence"][-1]["operation"]["attrs"]["n_mel"] = w_config.num_mel_bins + audio_processor_config_json["feature_extraction"]["sequence"][-1]["operation"]["attrs"]["n_mel"] = ( + w_config.num_mel_bins + ) # Get absolute path for the model output directory before running workflows model_path_abs = os.path.abspath(model_path) @@ -564,12 +574,15 @@ def main(): # run the encapsulation workflow print("\n[2/2] Running Whisper encapsulation workflow...") - run_encapsulation(default_encapsulation_json, model_path, w_config, enable_npu_ws, reshape, cache_dir, device=args.device) + run_encapsulation( + default_encapsulation_json, model_path, w_config, enable_npu_ws, reshape, cache_dir, device=args.device + ) # JSON dump the audio preprocessor config into the output directory audio_processor_config_json = json.dumps(audio_processor_config_json, indent=4) with open(os.path.join(model_path_abs, "audio_processor_config.json"), "w") as f: f.write(audio_processor_config_json) + if __name__ == "__main__": main() diff --git a/openai-whisper-large-v3-turbo/aitk/ov_evaluate.py b/openai-whisper-large-v3-turbo/aitk/ov_evaluate.py index 9df090d72..eef9e5631 100644 --- a/openai-whisper-large-v3-turbo/aitk/ov_evaluate.py +++ b/openai-whisper-large-v3-turbo/aitk/ov_evaluate.py @@ -1,8 +1,8 @@ import argparse import json +import logging import os import time -import logging from pathlib import Path import onnxruntime_genai as og @@ -10,6 +10,7 @@ logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def test_transcript(model_path, audio_path, num_beams=0, execution_provider="OpenVINO", device_type="NPU"): config = og.Config(model_path) config.set_provider_option(execution_provider, "device_type", device_type) @@ -48,6 +49,7 @@ def test_transcript(model_path, audio_path, num_beams=0, execution_provider="Ope return latencies + def main(): parser = argparse.ArgumentParser(description="Test Whisper ONNX GenAI models.") parser.add_argument("--execution_provider", type=str, default="CPUExecutionProvider", help="ORT Execution provider") @@ -58,15 +60,19 @@ def main(): # model path model_path = args.model_path - test_audio_url = "https://storage.openvinotoolkit.org/models_contrib/speech/2021.2/librispeech_s5/how_are_you_doing_today.wav" + test_audio_url = ( + "https://storage.openvinotoolkit.org/models_contrib/speech/2021.2/librispeech_s5/how_are_you_doing_today.wav" + ) test_audio_name = "how_are_you_doing_today.wav" import requests + r = requests.get(test_audio_url) with open(test_audio_name, "wb") as audio_file: - audio_file.write(r.content) + audio_file.write(r.content) from winml import register_execution_providers_to_onnxruntime_genai + register_execution_providers_to_onnxruntime_genai() num_beams = 1 @@ -80,7 +86,7 @@ def main(): "throughput-avg": throughput_avg, } resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) diff --git a/openai-whisper-large-v3-turbo/aitk/ov_workflow.py b/openai-whisper-large-v3-turbo/aitk/ov_workflow.py index 1aa4ef14e..55fef485b 100644 --- a/openai-whisper-large-v3-turbo/aitk/ov_workflow.py +++ b/openai-whisper-large-v3-turbo/aitk/ov_workflow.py @@ -1,13 +1,14 @@ import argparse import json +import logging import os import subprocess import sys -import logging logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -19,7 +20,7 @@ def parse_arguments(): def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) # Get arguments @@ -39,23 +40,43 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "ov_evaluate.py", - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file, - "--model_path", model_path], - check=True) + subprocess.run( + [ + sys.executable, + "ov_evaluate.py", + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + "--model_path", + model_path, + ], + check=True, + ) return # Generate model - subprocess.run([sys.executable, "convert_whisper_to_ovir.py", - "--output_dir", output_dir, - "--cache_dir", cache_dir, - "--model", "openai/whisper-large-v3-turbo", - "--weight-format", weight_format, - "--enable_npu_ws", str(enable_npu_ws), - "--device", oliveJson["systems"]["target_system"]["accelerators"][0]["device"]], - check=True) + subprocess.run( + [ + sys.executable, + "convert_whisper_to_ovir.py", + "--output_dir", + output_dir, + "--cache_dir", + cache_dir, + "--model", + "openai/whisper-large-v3-turbo", + "--weight-format", + weight_format, + "--enable_npu_ws", + str(enable_npu_ws), + "--device", + oliveJson["systems"]["target_system"]["accelerators"][0]["device"], + ], + check=True, + ) if __name__ == "__main__": diff --git a/openai-whisper-large-v3-turbo/aitk/qnn_app.py b/openai-whisper-large-v3-turbo/aitk/qnn_app.py index d0c33e307..4f02ab250 100644 --- a/openai-whisper-large-v3-turbo/aitk/qnn_app.py +++ b/openai-whisper-large-v3-turbo/aitk/qnn_app.py @@ -49,6 +49,7 @@ def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): session_options.add_provider_for_devices([ep_device], {} if ep_options is None else ep_options) break + def get_device_type(device_str): if device_str.lower() == "gpu": return ort.OrtHardwareDeviceType.GPU @@ -74,12 +75,14 @@ def __init__( add_ep_for_device(options, execution_provider, device_type) self.encoder = ort.InferenceSession( - encoder, sess_options=options, + encoder, + sess_options=options, ) self.encoder_latencies = [] self.decoder = ort.InferenceSession( - decoder, sess_options=options, + decoder, + sess_options=options, ) self.decoder_latencies = [] diff --git a/openai-whisper-large-v3-turbo/aitk/qnn_evaluate.py b/openai-whisper-large-v3-turbo/aitk/qnn_evaluate.py index 87878c08d..728892a63 100644 --- a/openai-whisper-large-v3-turbo/aitk/qnn_evaluate.py +++ b/openai-whisper-large-v3-turbo/aitk/qnn_evaluate.py @@ -1,14 +1,15 @@ -import numpy as np import argparse import json -import os import logging +import os +import numpy as np from qnn_app import HfWhisperAppWithSave, get_device_type logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def main(): parser = argparse.ArgumentParser(description="Evaluate Whisper") parser.add_argument( @@ -56,8 +57,11 @@ def main(): decoder_path = args.decoder from winml import register_execution_providers_to_onnxruntime + register_execution_providers_to_onnxruntime() - app = HfWhisperAppWithSave(encoder_path, decoder_path, args.model_id, args.execution_provider, get_device_type(args.device_str)) + app = HfWhisperAppWithSave( + encoder_path, decoder_path, args.model_id, args.execution_provider, get_device_type(args.device_str) + ) encoder_latencies = [] decoder_latencies = [] @@ -77,12 +81,9 @@ def main(): encoder_latency_avg = round(sum(encoder_latencies) / len(encoder_latencies) * 1000, 5) decoder_latency_avg = round(sum(decoder_latencies) / len(decoder_latencies) * 1000, 5) - metrics = { - "encoder-latency-avg": encoder_latency_avg, - "decoder-latency-avg": decoder_latency_avg - } + metrics = {"encoder-latency-avg": encoder_latency_avg, "decoder-latency-avg": decoder_latency_avg} resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) diff --git a/openai-whisper-large-v3-turbo/aitk/qnn_run.py b/openai-whisper-large-v3-turbo/aitk/qnn_run.py index 9c042f1db..52f0898c7 100644 --- a/openai-whisper-large-v3-turbo/aitk/qnn_run.py +++ b/openai-whisper-large-v3-turbo/aitk/qnn_run.py @@ -4,14 +4,15 @@ # --------------------------------------------------------------------- import argparse -import os import logging +import os -from qnn_app import HfWhisperAppWithSave, infer_audio, get_audio_name, get_device_type +from qnn_app import HfWhisperAppWithSave, get_audio_name, get_device_type, infer_audio logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def main(): parser = argparse.ArgumentParser(description="Demo") parser.add_argument( @@ -78,12 +79,15 @@ def main(): if args.execution_provider != "CPUExecutionProvider": from winml import register_execution_providers_to_onnxruntime + register_execution_providers_to_onnxruntime() - app = HfWhisperAppWithSave(encoder_path, decoder_path, args.model_id, args.execution_provider, get_device_type(args.device_str)) + app = HfWhisperAppWithSave( + encoder_path, decoder_path, args.model_id, args.execution_provider, get_device_type(args.device_str) + ) if not os.path.exists(args.audio_path) or os.path.isdir(args.audio_path): - from datasets import load_dataset import numpy as np + from datasets import load_dataset os.makedirs(args.audio_path, exist_ok=True) streamed_dataset = load_dataset(args.dataset_name, "clean", split=args.dataset_split, streaming=True) @@ -96,7 +100,7 @@ def main(): audio_name = get_audio_name(file_path) if args.save_data and os.path.exists(os.path.join(args.save_data, audio_name)): - #print(f"Skipping {file_path} as data already exists.") + # print(f"Skipping {file_path} as data already exists.") pass else: logger.info(f"Processing data {i} in {file_path} ...") diff --git a/openai-whisper-large-v3-turbo/aitk/qnn_workflow.py b/openai-whisper-large-v3-turbo/aitk/qnn_workflow.py index 4eb67d448..4adfb114b 100644 --- a/openai-whisper-large-v3-turbo/aitk/qnn_workflow.py +++ b/openai-whisper-large-v3-turbo/aitk/qnn_workflow.py @@ -1,14 +1,16 @@ import argparse import json +import logging import os -import olive.workflows import subprocess import sys -import logging + +import olive.workflows logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -16,15 +18,17 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() + def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None, - data_path: str | None = None, - num_data: int | None = None) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + config_path: str, + cache_dir: str, + output_dir: str, + activation_type: str | None = None, + precision: str | None = None, + data_path: str | None = None, + num_data: int | None = None, +) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -43,15 +47,16 @@ def load_update_config( def generate_model( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str, - skip_existing: bool = True, - activation_type: str | None = None, - precision: str | None = None, - data_path: str | None = None, - num_data: int | None = None): + history_folder: str, + config_path: str, + cache_dir: str, + output_dir: str, + skip_existing: bool = True, + activation_type: str | None = None, + precision: str | None = None, + data_path: str | None = None, + num_data: int | None = None, +): if skip_existing and os.path.exists(os.path.join(output_dir, "model.onnx")): logger.info(f"Output model {output_dir} already exists, skipping {config_path}.") return @@ -60,7 +65,7 @@ def generate_model( # save updated config for record config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) output = olive.workflows.run(oliveJson) if output is None or not output.has_output_model(): @@ -71,7 +76,7 @@ def generate_model( def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) # Get arguments @@ -85,7 +90,7 @@ def main(): num_data: int = config_pass["length"] # The cwd is model project folder audio_path: str = os.path.join("data", dataset_name.replace("/", "_"), dataset_split) - save_data_path: str = os.path.join("data", "_data_" + dataset_name.replace("/", "_"), dataset_split) + save_data_path: str = os.path.join("data", "_data_" + dataset_name.replace("/", "_"), dataset_split) history_folder = os.path.dirname(args.config) # When we have model_config, we are in evaluation @@ -99,14 +104,25 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "qnn_evaluate.py", - "--audio-path", audio_path, - "--encoder", encoder_path, - "--decoder", decoder_path, - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "qnn_evaluate.py", + "--audio-path", + audio_path, + "--encoder", + encoder_path, + "--decoder", + decoder_path, + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return # Generate original model @@ -115,21 +131,51 @@ def main(): original_decoder = os.path.join("data", "_decoder_fp32") generate_model("data", "whisper_large_v3_turbo_decoder_fp32.json", cache_dir, original_decoder) # Generate dataset - subprocess.run([sys.executable, "qnn_run.py", - "--audio-path", audio_path, - "--encoder", os.path.join(original_encoder, "model.onnx"), - "--decoder", os.path.join(original_decoder, "model.onnx"), - "--save_data", save_data_path, - "--dataset_name", dataset_name, - "--dataset_split", dataset_split, - "--num_data", str(num_data)], - check=True) + subprocess.run( + [ + sys.executable, + "qnn_run.py", + "--audio-path", + audio_path, + "--encoder", + os.path.join(original_encoder, "model.onnx"), + "--decoder", + os.path.join(original_decoder, "model.onnx"), + "--save_data", + save_data_path, + "--dataset_name", + dataset_name, + "--dataset_split", + dataset_split, + "--num_data", + str(num_data), + ], + check=True, + ) # Generate quantized model - generate_model(history_folder, "whisper_large_v3_turbo_encoder_qdq.json", cache_dir, os.path.join(output_dir, "encoder"), - False, activation_type, precision, save_data_path, num_data) + generate_model( + history_folder, + "whisper_large_v3_turbo_encoder_qdq.json", + cache_dir, + os.path.join(output_dir, "encoder"), + False, + activation_type, + precision, + save_data_path, + num_data, + ) # decoder has more data for 1 sample, to keep variants, multiply num_data by 10 - generate_model(history_folder, "whisper_large_v3_turbo_decoder_qdq.json", cache_dir, os.path.join(output_dir, "decoder"), - False, activation_type, precision, save_data_path, num_data * 10) + generate_model( + history_folder, + "whisper_large_v3_turbo_decoder_qdq.json", + cache_dir, + os.path.join(output_dir, "decoder"), + False, + activation_type, + precision, + save_data_path, + num_data * 10, + ) if __name__ == "__main__": diff --git a/openai-whisper-large-v3-turbo/aitk/whisper_decoder_load.py b/openai-whisper-large-v3-turbo/aitk/whisper_decoder_load.py index 3665ddd26..e56bf701a 100644 --- a/openai-whisper-large-v3-turbo/aitk/whisper_decoder_load.py +++ b/openai-whisper-large-v3-turbo/aitk/whisper_decoder_load.py @@ -2,9 +2,8 @@ import os import numpy as np -from qai_hub_models.utils.input_spec import make_torch_inputs - from olive.data.registry import Registry +from qai_hub_models.utils.input_spec import make_torch_inputs def model_loader(model_name): @@ -29,7 +28,9 @@ def generate_dummy_inputs(model=None): class DecoderBaseDataLoader: def __init__(self, data_path, num_data): - self.data_files = sorted(glob.glob(os.path.join(data_path, "**", "*_decoder_input.npy"), recursive=True))[:num_data] + self.data_files = sorted(glob.glob(os.path.join(data_path, "**", "*_decoder_input.npy"), recursive=True))[ + :num_data + ] print(f"Decoder data loader loaded {len(self.data_files)} samples from {data_path}") def __len__(self): diff --git a/openai-whisper-large-v3-turbo/aitk/whisper_encoder_load.py b/openai-whisper-large-v3-turbo/aitk/whisper_encoder_load.py index e6f2b6c7f..d86926bf8 100644 --- a/openai-whisper-large-v3-turbo/aitk/whisper_encoder_load.py +++ b/openai-whisper-large-v3-turbo/aitk/whisper_encoder_load.py @@ -2,9 +2,8 @@ import os import numpy as np -from qai_hub_models.utils.input_spec import make_torch_inputs - from olive.data.registry import Registry +from qai_hub_models.utils.input_spec import make_torch_inputs def model_loader(model_name): @@ -30,9 +29,9 @@ def generate_dummy_inputs(model=None): class EncoderBaseDataLoader: def __init__(self, data_path, num_data): # Ensure deterministic ordering of input feature files - self.data_files = sorted( - glob.glob(os.path.join(data_path, "**", "*_input_features.npy"), recursive=True) - )[:num_data] + self.data_files = sorted(glob.glob(os.path.join(data_path, "**", "*_input_features.npy"), recursive=True))[ + :num_data + ] print(f"Encoder data loader loaded {len(self.data_files)} samples from {data_path}") def __len__(self): diff --git a/sam-vit-base/aitk/config.py b/sam-vit-base/aitk/config.py index 65de54136..e59959da9 100644 --- a/sam-vit-base/aitk/config.py +++ b/sam-vit-base/aitk/config.py @@ -3,6 +3,7 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- + class ModelConfig: model_name = "facebook/sam-vit-base" data_dir = "quantization_dataset" diff --git a/sam-vit-base/aitk/model_patches.py b/sam-vit-base/aitk/model_patches.py index 1d40da7c5..fe96b2aea 100644 --- a/sam-vit-base/aitk/model_patches.py +++ b/sam-vit-base/aitk/model_patches.py @@ -421,7 +421,7 @@ def __init__(self, model): self.model.prompt_encoder = ModSamPromptEncoder(self.model.prompt_encoder) def forward(self, input_points, input_labels, image_embeddings): - return self.model(input_points=input_points, input_labels = input_labels, image_embeddings=image_embeddings) + return self.model(input_points=input_points, input_labels=input_labels, image_embeddings=image_embeddings) class ModSamMaskPointDecoder(nn.Module): diff --git a/sam-vit-base/aitk/sam_mask_generator.py b/sam-vit-base/aitk/sam_mask_generator.py index b116ef0ae..ed2e7ff6a 100644 --- a/sam-vit-base/aitk/sam_mask_generator.py +++ b/sam-vit-base/aitk/sam_mask_generator.py @@ -3,16 +3,17 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -import numpy as np import argparse -import os -from PIL import Image + +import numpy as np import onnxruntime as ort +from PIL import Image from transformers import SamProcessor # Load processor processor = SamProcessor.from_pretrained("facebook/sam-vit-base") + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -23,7 +24,7 @@ def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): def get_mask_ort(sess_ve, sess_md, image, box, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs): - inputs = processor(image, input_boxes = box, return_tensors="np") + inputs = processor(image, input_boxes=box, return_tensors="np") ort_pixel_values = inputs["pixel_values"] input_boxes = inputs["input_boxes"] @@ -36,7 +37,7 @@ def get_mask_ort(sess_ve, sess_md, image, box, ve_dtype, md_dtype, sess_ve_input input_md = { sess_md_inputs[0].name: np.array(ort_input_points, dtype=md_dtype), sess_md_inputs[1].name: np.array(ort_input_labels, dtype=md_dtype), - sess_md_inputs[2].name: np.array(result_ve[0], dtype=md_dtype) + sess_md_inputs[2].name: np.array(result_ve[0], dtype=md_dtype), } result_md = sess_md.run(None, input_md) @@ -44,15 +45,14 @@ def get_mask_ort(sess_ve, sess_md, image, box, ve_dtype, md_dtype, sess_ve_input pred_masks = result_md[1] masks = processor.image_processor.post_process_masks( - [pred_masks[0]], - inputs["original_sizes"], - inputs["reshaped_input_sizes"] + [pred_masks[0]], inputs["original_sizes"], inputs["reshaped_input_sizes"] )[0].numpy() pred_max_ind = np.argmax(scores) mask = masks[0, pred_max_ind] return np.array(mask, dtype=np.uint8) + def main(): parser = argparse.ArgumentParser(description="Run SAM ONNX models and save mask.") parser.add_argument("--execution_provider", type=str, default="CPUExecutionProvider", help="ORT Execution provider") @@ -70,6 +70,7 @@ def main(): # Loading models into ORT session if args.execution_provider != "CPUExecutionProvider": from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -92,8 +93,8 @@ def main(): sess_ve_inputs = sess_ve.get_inputs() sess_md_inputs = sess_md.get_inputs() - ve_dtype = np.float32 if sess_ve_inputs[0].type == 'tensor(float)' else np.float16 - md_dtype = np.float32 if sess_md_inputs[0].type == 'tensor(float)' else np.float16 + ve_dtype = np.float32 if sess_ve_inputs[0].type == "tensor(float)" else np.float16 + md_dtype = np.float32 if sess_md_inputs[0].type == "tensor(float)" else np.float16 # Get mask mask = get_mask_ort(sess_ve, sess_md, raw_image, input_box, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs) @@ -103,5 +104,6 @@ def main(): mask_img.save(args.output_path) print(f"Mask saved to {args.output_path}") + if __name__ == "__main__": main() diff --git a/sam-vit-base/aitk/sam_qnn_evaluation.py b/sam-vit-base/aitk/sam_qnn_evaluation.py index a6d267737..0353462b5 100644 --- a/sam-vit-base/aitk/sam_qnn_evaluation.py +++ b/sam-vit-base/aitk/sam_qnn_evaluation.py @@ -3,16 +3,16 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -import numpy as np import argparse import json +import logging import os import time -import logging - from urllib import request -from PIL import Image + +import numpy as np import onnxruntime as ort +from PIL import Image from transformers import SamProcessor logger = logging.getLogger(os.path.basename(__file__)) @@ -21,6 +21,7 @@ # Load processor processor = SamProcessor.from_pretrained("facebook/sam-vit-base") + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -48,7 +49,7 @@ def test_mask_ort(sess_ve, sess_md, image, ve_dtype, md_dtype, sess_ve_inputs, s input_md = { sess_md_inputs[0].name: np.array(ort_input_points, dtype=md_dtype), sess_md_inputs[1].name: np.array(ort_input_labels, dtype=md_dtype), - sess_md_inputs[2].name: np.array(result_ve[0], dtype=md_dtype) + sess_md_inputs[2].name: np.array(result_ve[0], dtype=md_dtype), } decoder_start = time.perf_counter() @@ -69,6 +70,7 @@ def main(): # Loading models into ORT session from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -101,19 +103,18 @@ def main(): for _ in range(10): # Test mask - encoder_latency, decoder_latency = test_mask_ort(sess_ve, sess_md, raw_image, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs) + encoder_latency, decoder_latency = test_mask_ort( + sess_ve, sess_md, raw_image, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs + ) encoder_latencies.append(encoder_latency) decoder_latencies.append(decoder_latency) encoder_latency_avg = round(sum(encoder_latencies) / len(encoder_latencies) * 1000, 5) decoder_latency_avg = round(sum(decoder_latencies) / len(decoder_latencies) * 1000, 5) - metrics = { - "vision-encoder-latency-avg": encoder_latency_avg, - "mask-decoder-latency-avg": decoder_latency_avg - } + metrics = {"vision-encoder-latency-avg": encoder_latency_avg, "mask-decoder-latency-avg": decoder_latency_avg} resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) diff --git a/sam-vit-base/aitk/sam_qnn_workflow.py b/sam-vit-base/aitk/sam_qnn_workflow.py index 851e655ae..de10664ae 100644 --- a/sam-vit-base/aitk/sam_qnn_workflow.py +++ b/sam-vit-base/aitk/sam_qnn_workflow.py @@ -1,14 +1,16 @@ import argparse import json +import logging import os -import olive.workflows import subprocess import sys -import logging + +import olive.workflows logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -16,14 +18,16 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() + def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None, - num_data: int | None = None) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + config_path: str, + cache_dir: str, + output_dir: str, + activation_type: str | None = None, + precision: str | None = None, + num_data: int | None = None, +) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -41,15 +45,17 @@ def load_update_config( return oliveJson + def generate_model( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str, - skip_existing: bool = True, - activation_type: str | None = None, - precision: str | None = None, - num_data: int | None = None): + history_folder: str, + config_path: str, + cache_dir: str, + output_dir: str, + skip_existing: bool = True, + activation_type: str | None = None, + precision: str | None = None, + num_data: int | None = None, +): if skip_existing and os.path.exists(os.path.join(output_dir, "model.onnx")): logger.info(f"Output model {output_dir} already exists, skipping {config_path}.") return @@ -59,7 +65,7 @@ def generate_model( config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) print("write to: ", os.path.join(history_folder, config_name)) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) output = olive.workflows.run(oliveJson) if output is None or not output.has_output_model(): @@ -70,7 +76,7 @@ def generate_model( def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) if args.model_config: @@ -83,16 +89,25 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "sam_qnn_evaluation.py", - "--model_ve", encoder_path, - "--model_md", decoder_path, - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "sam_qnn_evaluation.py", + "--model_ve", + encoder_path, + "--model_md", + decoder_path, + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return - # Get arguments output_dir: str = oliveJson["output_dir"] cache_dir: str = oliveJson["cache_dir"] @@ -110,16 +125,33 @@ def main(): os.makedirs(os.path.join(history_folder, "model"), exist_ok=True) from config import ModelConfig + ModelConfig.image_dataset = dataset_name ModelConfig.image_dataset_split = dataset_split ModelConfig.data_dir = "quantization_dataset" # Generate encoder model - generate_model(history_folder, "sam_vision_encoder_qnn.json", cache_dir, os.path.join(output_dir, "encoder"), - False, activation_type, precision, num_data) + generate_model( + history_folder, + "sam_vision_encoder_qnn.json", + cache_dir, + os.path.join(output_dir, "encoder"), + False, + activation_type, + precision, + num_data, + ) # Generate decoder model - generate_model(history_folder, "sam_mask_decoder_qnn.json", cache_dir, os.path.join(output_dir, "decoder"), - False, activation_type, precision, num_data) + generate_model( + history_folder, + "sam_mask_decoder_qnn.json", + cache_dir, + os.path.join(output_dir, "decoder"), + False, + activation_type, + precision, + num_data, + ) if __name__ == "__main__": diff --git a/sam-vit-base/aitk/user_script.py b/sam-vit-base/aitk/user_script.py index f33392c7e..0b801633f 100644 --- a/sam-vit-base/aitk/user_script.py +++ b/sam-vit-base/aitk/user_script.py @@ -16,9 +16,9 @@ from datasets import load_dataset from model_patches import ( ModSamMaskBoxDecoder, + ModSamMaskdecoder, ModSamMaskPointDecoder, ModSamVisionEncoder, - ModSamMaskdecoder, ) from olive.data.registry import Registry from transformers import SamModel, SamProcessor diff --git a/sam2.1-hiera-small/aitk/generate_model.py b/sam2.1-hiera-small/aitk/generate_model.py index cb10d4e7c..9a69db0bf 100644 --- a/sam2.1-hiera-small/aitk/generate_model.py +++ b/sam2.1-hiera-small/aitk/generate_model.py @@ -412,9 +412,13 @@ def main(): } with torch.no_grad(): - torch.onnx.export(encoder, en_inputs, "sam21_vision_encoder.onnx", opset_version=20, do_constant_folding=True, dynamo = False) + torch.onnx.export( + encoder, en_inputs, "sam21_vision_encoder.onnx", opset_version=20, do_constant_folding=True, dynamo=False + ) with torch.no_grad(): - torch.onnx.export(decoder, de_inputs, "sam21_mask_decoder.onnx", opset_version=20, do_constant_folding=True, dynamo = False) + torch.onnx.export( + decoder, de_inputs, "sam21_mask_decoder.onnx", opset_version=20, do_constant_folding=True, dynamo=False + ) encoder_onnx_model = onnx.load("sam21_vision_encoder.onnx") simplified_encoder_onnx_model, check = simplify(encoder_onnx_model) diff --git a/sam2.1-hiera-small/aitk/sam2_mask_generator.py b/sam2.1-hiera-small/aitk/sam2_mask_generator.py index d6513b9df..02f1a99c3 100644 --- a/sam2.1-hiera-small/aitk/sam2_mask_generator.py +++ b/sam2.1-hiera-small/aitk/sam2_mask_generator.py @@ -4,7 +4,6 @@ # -------------------------------------------------------------------------- import argparse -import os import numpy as np import onnxruntime as ort @@ -18,12 +17,11 @@ [ transforms.Resize((1024, 1024)), # Resize to 1024x1024 transforms.ToTensor(), # Convert to tensor [C,H,W] and scale to [0,1] - transforms.Normalize( - mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] - ), # Normalize + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # Normalize ] ) + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -33,9 +31,7 @@ def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): break -def get_mask_ort( - sess_ve, sess_md, image, box, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs -): +def get_mask_ort(sess_ve, sess_md, image, box, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs): w, h = image.size processed_image = sam2_transform(image) inputs = processed_image.float().numpy()[None, :] @@ -53,9 +49,7 @@ def get_mask_ort( ort_input_labels = np.concatenate([blank_labels, box_labels], axis=0)[None, :] input_ve = {sess_ve_inputs[0].name: np.array(ort_pixel_values, dtype=ve_dtype)} - image_embedding, high_res_features1, high_res_features2 = sess_ve.run( - None, input_ve - ) + image_embedding, high_res_features1, high_res_features2 = sess_ve.run(None, input_ve) input_md = { sess_md_inputs[0].name: image_embedding.astype(md_dtype), @@ -69,13 +63,7 @@ def get_mask_ort( pred_masks = result_md[0] scores = result_md[1] - masks = ( - F.interpolate( - torch.Tensor(pred_masks), size=(h, w), mode="bilinear", align_corners=False - ) - .detach() - .numpy() - ) + masks = F.interpolate(torch.Tensor(pred_masks), size=(h, w), mode="bilinear", align_corners=False).detach().numpy() pred_max_ind = np.argmax(scores) mask = masks[0, pred_max_ind] @@ -86,24 +74,16 @@ def main(): parser = argparse.ArgumentParser(description="Run SAM ONNX models and save mask.") parser.add_argument("--execution_provider", type=str, default="CPUExecutionProvider", help="ORT Execution provider") parser.add_argument("--device_str", type=str, default="cpu") - parser.add_argument( - "--model_ve", required=True, help="Path to vision encoder ONNX model" - ) - parser.add_argument( - "--model_md", required=True, help="Path to mask decoder ONNX model" - ) + parser.add_argument("--model_ve", required=True, help="Path to vision encoder ONNX model") + parser.add_argument("--model_md", required=True, help="Path to mask decoder ONNX model") parser.add_argument("--image_path", required=True, help="Path to input image") parser.add_argument( "--output_path", default="mask_output.png", help="Path to save the output mask image", ) - parser.add_argument( - "--box_x", type=int, default=40, help="Top-Left X coordinate of input box" - ) - parser.add_argument( - "--box_y", type=int, default=235, help="Top-Left Y coordinate of input box" - ) + parser.add_argument("--box_x", type=int, default=40, help="Top-Left X coordinate of input box") + parser.add_argument("--box_y", type=int, default=235, help="Top-Left Y coordinate of input box") parser.add_argument("--box_w", type=int, default=940, help="Width of input box") parser.add_argument("--box_h", type=int, default=490, help="Height of input box") args = parser.parse_args() @@ -111,6 +91,7 @@ def main(): # Loading models into ORT session if args.execution_provider != "CPUExecutionProvider": from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() diff --git a/sam2.1-hiera-small/aitk/sam2_qnn_evaluation.py b/sam2.1-hiera-small/aitk/sam2_qnn_evaluation.py index 50578fb37..9de9f9cb4 100644 --- a/sam2.1-hiera-small/aitk/sam2_qnn_evaluation.py +++ b/sam2.1-hiera-small/aitk/sam2_qnn_evaluation.py @@ -3,14 +3,14 @@ # Licensed under the MIT License. # -------------------------------------------------------------------------- -import numpy as np import argparse import json +import logging import os import time -import logging - from urllib import request + +import numpy as np import onnxruntime as ort from PIL import Image from torchvision import transforms @@ -23,12 +23,11 @@ [ transforms.Resize((1024, 1024)), # Resize to 1024x1024 transforms.ToTensor(), # Convert to tensor [C,H,W] and scale to [0,1] - transforms.Normalize( - mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225] - ), # Normalize + transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # Normalize ] ) + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -38,9 +37,7 @@ def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): break -def test_mask_ort( - sess_ve, sess_md, image, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs -): +def test_mask_ort(sess_ve, sess_md, image, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs): w, h = image.size processed_image = sam2_transform(image) inputs = processed_image.float().numpy()[None, :] @@ -60,9 +57,7 @@ def test_mask_ort( input_ve = {sess_ve_inputs[0].name: np.array(ort_pixel_values, dtype=ve_dtype)} encoder_start = time.perf_counter() - image_embedding, high_res_features1, high_res_features2 = sess_ve.run( - None, input_ve - ) + image_embedding, high_res_features1, high_res_features2 = sess_ve.run(None, input_ve) encoder_latency = time.perf_counter() - encoder_start input_md = { @@ -91,6 +86,7 @@ def main(): # Loading models into ORT session from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -123,19 +119,18 @@ def main(): for _ in range(10): # Test mask - encoder_latency, decoder_latency = test_mask_ort(sess_ve, sess_md, raw_image, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs) + encoder_latency, decoder_latency = test_mask_ort( + sess_ve, sess_md, raw_image, ve_dtype, md_dtype, sess_ve_inputs, sess_md_inputs + ) encoder_latencies.append(encoder_latency) decoder_latencies.append(decoder_latency) encoder_latency_avg = round(sum(encoder_latencies) / len(encoder_latencies) * 1000, 5) decoder_latency_avg = round(sum(decoder_latencies) / len(decoder_latencies) * 1000, 5) - metrics = { - "vision-encoder-latency-avg": encoder_latency_avg, - "mask-decoder-latency-avg": decoder_latency_avg - } + metrics = {"vision-encoder-latency-avg": encoder_latency_avg, "mask-decoder-latency-avg": decoder_latency_avg} resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) diff --git a/sam2.1-hiera-small/aitk/sam2_qnn_workflow.py b/sam2.1-hiera-small/aitk/sam2_qnn_workflow.py index 0f23cf28a..a3db22646 100644 --- a/sam2.1-hiera-small/aitk/sam2_qnn_workflow.py +++ b/sam2.1-hiera-small/aitk/sam2_qnn_workflow.py @@ -1,14 +1,16 @@ import argparse import json +import logging import os -import olive.workflows import subprocess import sys -import logging + +import olive.workflows logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -16,14 +18,16 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() + def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None, - num_data: int | None = None) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + config_path: str, + cache_dir: str, + output_dir: str, + activation_type: str | None = None, + precision: str | None = None, + num_data: int | None = None, +) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -41,15 +45,17 @@ def load_update_config( return oliveJson + def generate_model( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str, - skip_existing: bool = True, - activation_type: str | None = None, - precision: str | None = None, - num_data: int | None = None): + history_folder: str, + config_path: str, + cache_dir: str, + output_dir: str, + skip_existing: bool = True, + activation_type: str | None = None, + precision: str | None = None, + num_data: int | None = None, +): if skip_existing and os.path.exists(os.path.join(output_dir, "model.onnx")): logger.info(f"Output model {output_dir} already exists, skipping {config_path}.") return @@ -59,7 +65,7 @@ def generate_model( config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) print("write to: ", os.path.join(history_folder, config_name)) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) output = olive.workflows.run(oliveJson) if output is None or not output.has_output_model(): @@ -70,7 +76,7 @@ def generate_model( def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) if args.model_config: @@ -83,16 +89,25 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "sam2_qnn_evaluation.py", - "--model_ve", encoder_path, - "--model_md", decoder_path, - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "sam2_qnn_evaluation.py", + "--model_ve", + encoder_path, + "--model_md", + decoder_path, + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return - # Get arguments output_dir: str = oliveJson["output_dir"] cache_dir: str = oliveJson["cache_dir"] @@ -110,6 +125,7 @@ def main(): os.makedirs(os.path.join(history_folder, "model"), exist_ok=True) from config import ModelConfig + ModelConfig.image_dataset = dataset_name ModelConfig.image_dataset_split = dataset_split ModelConfig.data_dir = "quantization_dataset" @@ -117,11 +133,27 @@ def main(): subprocess.run([sys.executable, "generate_model.py"], check=True) # Generate encoder model - generate_model(history_folder, "sam21_vision_encoder_qnn.json", cache_dir, os.path.join(output_dir, "encoder"), - False, activation_type, precision, num_data) + generate_model( + history_folder, + "sam21_vision_encoder_qnn.json", + cache_dir, + os.path.join(output_dir, "encoder"), + False, + activation_type, + precision, + num_data, + ) # Generate decoder model - generate_model(history_folder, "sam21_mask_decoder_qnn.json", cache_dir, os.path.join(output_dir, "decoder"), - False, activation_type, precision, num_data) + generate_model( + history_folder, + "sam21_mask_decoder_qnn.json", + cache_dir, + os.path.join(output_dir, "decoder"), + False, + activation_type, + precision, + num_data, + ) if __name__ == "__main__": diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/evaluation.py b/sd-legacy-stable-diffusion-v1-5/aitk/evaluation.py index 41cdd924d..6dee878fe 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/evaluation.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/evaluation.py @@ -249,7 +249,7 @@ def main(raw_args=None): train_data = dataset[args.dataset_split] prompts = [sanitize_path(example["captions"][0]) for i, example in enumerate(train_data) if i < args.num_data] # TODO unused data for evaluation - #real_images = get_real_images(train_data, args.num_data) + # real_images = get_real_images(train_data, args.num_data) train_num = math.floor(len(prompts) * args.train_ratio) clip_score_fn = partial(clip_score, model_name_or_path="openai/clip-vit-base-patch16") diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_evaluation.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_evaluation.py index 66e5380ba..3fd6b4737 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_evaluation.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_evaluation.py @@ -4,18 +4,18 @@ # -------------------------------------------------------------------------- import json -import os import logging +import os from pathlib import Path import numpy as np import onnxruntime as ort - from sd_utils.ov import OVStableDiffusionPipeline logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_args(raw_args): import argparse @@ -56,6 +56,7 @@ def parse_args(raw_args): ) return parser.parse_args(raw_args) + def get_device_type(device_str): if device_str.lower() == "gpu": return ort.OrtHardwareDeviceType.GPU @@ -64,6 +65,7 @@ def get_device_type(device_str): else: return ort.OrtHardwareDeviceType.CPU + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -80,6 +82,7 @@ def main(raw_args=None): model_dir = Path(args.script_dir) / "model" / args.model_dir / args.model_id from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -121,12 +124,13 @@ def main(raw_args=None): metrics = { "text-encoder-latency-avg": text_encoder_latency_avg, "unet-latency-avg": unet_latency_avg, - "vae-decoder-latency-avg": vae_decoder_latency_avg + "vae-decoder-latency-avg": vae_decoder_latency_avg, } resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) + if __name__ == "__main__": main() diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_workflow.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_workflow.py index 51c6c5282..f0f3d4fef 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_workflow.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_ov_workflow.py @@ -1,13 +1,14 @@ import argparse import json +import logging import os import subprocess import sys -import logging logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -15,11 +16,9 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() -def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + +def load_update_config(config_path: str, cache_dir: str, output_dir: str) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -27,28 +26,26 @@ def load_update_config( return oliveJson -def copy_olive_config( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str): + +def copy_olive_config(history_folder: str, config_path: str, cache_dir: str, output_dir: str): logger.info(f"Copying {config_path} to {history_folder}...") oliveJson = load_update_config(config_path, cache_dir, output_dir) # save updated config for record config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) + def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) # For static quantization, the QDQ data should match the target scenario. - guidance_scale=str(7.5) - num_inference_steps=str(25) + guidance_scale = str(7.5) + num_inference_steps = str(25) if args.model_config: model_path: str = os.path.dirname(args.model_config) @@ -57,19 +54,31 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "sd_ov_evaluation.py", - "--script_dir", os.path.dirname(model_path), - "--model_dir", "optimized", - "--model_id", "stable-diffusion-v1-5/stable-diffusion-v1-5", - "--guidance_scale", guidance_scale, - "--num_inference_steps", num_inference_steps, - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "sd_ov_evaluation.py", + "--script_dir", + os.path.dirname(model_path), + "--model_dir", + "optimized", + "--model_id", + "stable-diffusion-v1-5/stable-diffusion-v1-5", + "--guidance_scale", + guidance_scale, + "--num_inference_steps", + num_inference_steps, + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return - # Get arguments output_dir: str = oliveJson["output_dir"] cache_dir: str = oliveJson["cache_dir"] @@ -87,11 +96,15 @@ def main(): static_shape = oliveJson["passes"]["aitkpython"]["static_shape"] cmd = [ - sys.executable, "stable_diffusion.py", - "--script_dir", history_folder, - "--model_id", "stable-diffusion-v1-5/stable-diffusion-v1-5", - "--provider", "openvino", - "--optimize" + sys.executable, + "stable_diffusion.py", + "--script_dir", + history_folder, + "--model_id", + "stable-diffusion-v1-5/stable-diffusion-v1-5", + "--provider", + "openvino", + "--optimize", ] if static_shape: cmd.append("--static_shape") @@ -99,5 +112,6 @@ def main(): # run stable_diffusion.py to generate onnx model subprocess.run(cmd, check=True) + if __name__ == "__main__": main() diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_evaluation.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_evaluation.py index fed43c4ae..9f5b02ee0 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_evaluation.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_evaluation.py @@ -4,18 +4,18 @@ # -------------------------------------------------------------------------- import json -import os import logging +import os from pathlib import Path import numpy as np import onnxruntime as ort - from sd_utils.qdq import OnnxStableDiffusionPipelineWithSave logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_args(raw_args): import argparse @@ -56,6 +56,7 @@ def parse_args(raw_args): ) return parser.parse_args(raw_args) + def get_device_type(device_str): if device_str.lower() == "gpu": return ort.OrtHardwareDeviceType.GPU @@ -64,6 +65,7 @@ def get_device_type(device_str): else: return ort.OrtHardwareDeviceType.CPU + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -80,6 +82,7 @@ def main(raw_args=None): model_dir = Path(args.script_dir) / "model" / args.model_dir / args.model_id from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -122,12 +125,13 @@ def main(raw_args=None): metrics = { "text-encoder-latency-avg": text_encoder_latency_avg, "unet-latency-avg": unet_latency_avg, - "vae-decoder-latency-avg": vae_decoder_latency_avg + "vae-decoder-latency-avg": vae_decoder_latency_avg, } resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) + if __name__ == "__main__": main() diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_workflow.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_workflow.py index ce5ee42c0..025f70151 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_workflow.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_qnn_workflow.py @@ -1,13 +1,14 @@ import argparse import json +import logging import os import subprocess import sys -import logging logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -15,13 +16,11 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() + def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + config_path: str, cache_dir: str, output_dir: str, activation_type: str | None = None, precision: str | None = None +) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -35,30 +34,33 @@ def load_update_config( return oliveJson + def copy_olive_config( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None): + history_folder: str, + config_path: str, + cache_dir: str, + output_dir: str, + activation_type: str | None = None, + precision: str | None = None, +): logger.info(f"Copying {config_path} to {history_folder}...") oliveJson = load_update_config(config_path, cache_dir, output_dir, activation_type, precision) # save updated config for record config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) + def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) # For static quantization, the QDQ data should match the target scenario. - guidance_scale=str(7.5) - num_inference_steps=str(25) + guidance_scale = str(7.5) + num_inference_steps = str(25) if args.model_config: model_path: str = os.path.dirname(args.model_config) @@ -67,19 +69,31 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "sd_qnn_evaluation.py", - "--script_dir", os.path.dirname(model_path), - "--model_dir", "optimized", - "--model_id", "stable-diffusion-v1-5/stable-diffusion-v1-5", - "--guidance_scale", guidance_scale, - "--num_inference_steps", num_inference_steps, - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "sd_qnn_evaluation.py", + "--script_dir", + os.path.dirname(model_path), + "--model_dir", + "optimized", + "--model_id", + "stable-diffusion-v1-5/stable-diffusion-v1-5", + "--guidance_scale", + guidance_scale, + "--num_inference_steps", + num_inference_steps, + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return - # Get arguments output_dir: str = oliveJson["output_dir"] cache_dir: str = oliveJson["cache_dir"] @@ -103,41 +117,76 @@ def main(): copy_olive_config(history_folder, config_name, cache_dir, output_dir, activation_type, precision) # run stable_diffusion.py to generate onnx unoptimized model - subprocess.run([sys.executable, "stable_diffusion.py", - "--script_dir", history_folder, - "--model_id", "stable-diffusion-v1-5/stable-diffusion-v1-5", - "--provider", "cpu", - "--format", "qdq", - "--optimize", - "--only_conversion"], - check=True) + subprocess.run( + [ + sys.executable, + "stable_diffusion.py", + "--script_dir", + history_folder, + "--model_id", + "stable-diffusion-v1-5/stable-diffusion-v1-5", + "--provider", + "cpu", + "--format", + "qdq", + "--optimize", + "--only_conversion", + ], + check=True, + ) # # run evaluation.py to generate data data_dir = f"../../quantize_data/{guidance_scale}_{num_inference_steps}" - subprocess.run([sys.executable, "evaluation.py", - "--script_dir", history_folder, - "--save_data", - "--model_id", "stable-diffusion-v1-5/stable-diffusion-v1-5", - "--num_inference_steps", num_inference_steps, - "--seed", "0", - "--data_dir", data_dir, - "--dataset_name", dataset_name, - "--dataset_split", dataset_split, - "--num_data", str(num_data), - # generate data for quantization calibration, so use all data as training data to match num_data - "--train_ratio", "1", - "--guidance_scale", guidance_scale], - check=True) + subprocess.run( + [ + sys.executable, + "evaluation.py", + "--script_dir", + history_folder, + "--save_data", + "--model_id", + "stable-diffusion-v1-5/stable-diffusion-v1-5", + "--num_inference_steps", + num_inference_steps, + "--seed", + "0", + "--data_dir", + data_dir, + "--dataset_name", + dataset_name, + "--dataset_split", + dataset_split, + "--num_data", + str(num_data), + # generate data for quantization calibration, so use all data as training data to match num_data + "--train_ratio", + "1", + "--guidance_scale", + guidance_scale, + ], + check=True, + ) # run stable_diffusion.py to generate onnx quantized model - subprocess.run([sys.executable, "stable_diffusion.py", - "--script_dir", history_folder, - "--model_id", "stable-diffusion-v1-5/stable-diffusion-v1-5", - "--provider", "cpu", - "--format", "qdq", - "--data_dir", data_dir + "/data", - "--optimize"], - check=True) + subprocess.run( + [ + sys.executable, + "stable_diffusion.py", + "--script_dir", + history_folder, + "--model_id", + "stable-diffusion-v1-5/stable-diffusion-v1-5", + "--provider", + "cpu", + "--format", + "qdq", + "--data_dir", + data_dir + "/data", + "--optimize", + ], + check=True, + ) + if __name__ == "__main__": main() diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/onnx_patch.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/onnx_patch.py index 0570ef67f..253c4dcef 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/onnx_patch.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/onnx_patch.py @@ -4,16 +4,15 @@ # -------------------------------------------------------------------------- import os -import time import shutil +import time from pathlib import Path from typing import Optional, Union -from diffusers import OnnxRuntimeModel -from diffusers.utils import ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME - import numpy as np import onnxruntime as ort +from diffusers import OnnxRuntimeModel +from diffusers.utils import ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME class PatchedOnnxRuntimeModel(OnnxRuntimeModel): @@ -78,7 +77,6 @@ def save_pretrained( self._save_pretrained(save_directory, **kwargs) - @classmethod def _from_pretrained( cls, diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ort.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ort.py index 3f9c8e9ba..0bb70acb0 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ort.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ort.py @@ -9,7 +9,6 @@ from pathlib import Path import onnxruntime as ort - from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from olive.model import ONNXModelHandler from onnxruntime import __version__ as OrtVersion @@ -103,7 +102,9 @@ def save_onnx_pipeline( from sd_utils.onnx_patch import PatchedOnnxRuntimeModel if has_safety_checker: - safety_checker = PatchedOnnxRuntimeModel.from_pretrained(model_info["safety_checker"]["unoptimized"]["path"].parent) + safety_checker = PatchedOnnxRuntimeModel.from_pretrained( + model_info["safety_checker"]["unoptimized"]["path"].parent + ) else: safety_checker = None @@ -112,7 +113,9 @@ def save_onnx_pipeline( vae_decoder=PatchedOnnxRuntimeModel.from_pretrained(model_info["vae_decoder"]["unoptimized"]["path"].parent), text_encoder=PatchedOnnxRuntimeModel.from_pretrained(model_info["text_encoder"]["unoptimized"]["path"].parent), tokenizer=pipeline.tokenizer, - unet=PatchedOnnxRuntimeModel.from_pretrained(model_info["unet"]["unoptimized"]["path"].parent, provider_options=[]), + unet=PatchedOnnxRuntimeModel.from_pretrained( + model_info["unet"]["unoptimized"]["path"].parent, provider_options=[] + ), scheduler=pipeline.scheduler, safety_checker=safety_checker, feature_extractor=pipeline.feature_extractor, diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ov.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ov.py index a84c05f75..f53fe31cc 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ov.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/ov.py @@ -4,19 +4,16 @@ # Example modified from: https://docs.openvino.ai/2023.3/notebooks/225-stable-diffusion-text-to-image-with-output.html # -------------------------------------------------------------------------- import inspect -import os from pathlib import Path from typing import Callable, Optional, Union import numpy as np import onnxruntime as ort import torch - from diffusers import StableDiffusionPipeline -from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput - +from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from sd_utils.onnx_patch import PatchedOnnxRuntimeModel @@ -161,8 +158,8 @@ def update_ov_config(config: dict, static_shape: bool): config["passes"] = { "ov_convert": config["passes"]["ov_convert"], "ov_io_update": config["passes"]["ov_io_update"], - "ov_encapsulation": config["passes"]["ov_encapsulation"] - } + "ov_encapsulation": config["passes"]["ov_encapsulation"], + } if static_shape == False: config["passes"]["ov_io_update"]["static"] = False @@ -206,6 +203,7 @@ def get_ov_pipeline(common_args, ov_args, optimized_model_dir): return StableDiffusionPipeline.from_pretrained(common_args.model_id) from winml import register_execution_providers + register_execution_providers() print("Loading models into ORT session...") diff --git a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/qdq.py b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/qdq.py index 6847dcf0e..bf0fbc1af 100644 --- a/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/qdq.py +++ b/sd-legacy-stable-diffusion-v1-5/aitk/sd_utils/qdq.py @@ -12,10 +12,9 @@ import sd_utils import sd_utils.config import torch - -from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput +from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline # ruff: noqa: T201 @@ -266,6 +265,7 @@ def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): def get_qdq_pipeline(model_dir, common_args, qdq_args, script_dir): if common_args.provider != "cpu": from winml import register_execution_providers + register_execution_providers() ort.set_default_logger_severity(3) diff --git a/sd2-community-stable-diffusion-2-1/aitk/evaluation.py b/sd2-community-stable-diffusion-2-1/aitk/evaluation.py index a74ef524d..74b33d307 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/evaluation.py +++ b/sd2-community-stable-diffusion-2-1/aitk/evaluation.py @@ -249,7 +249,7 @@ def main(raw_args=None): train_data = dataset[args.dataset_split] prompts = [sanitize_path(example["captions"][0]) for i, example in enumerate(train_data) if i < args.num_data] # TODO unused data for evaluation - #real_images = get_real_images(train_data, args.num_data) + # real_images = get_real_images(train_data, args.num_data) train_num = math.floor(len(prompts) * args.train_ratio) clip_score_fn = partial(clip_score, model_name_or_path="openai/clip-vit-base-patch16") diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_ov_evaluation.py b/sd2-community-stable-diffusion-2-1/aitk/sd_ov_evaluation.py index 6720d520b..ac56f7269 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_ov_evaluation.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_ov_evaluation.py @@ -4,18 +4,18 @@ # -------------------------------------------------------------------------- import json -import os import logging +import os from pathlib import Path import numpy as np import onnxruntime as ort - from sd_utils.ov import OVStableDiffusionPipeline logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_args(raw_args): import argparse @@ -56,6 +56,7 @@ def parse_args(raw_args): ) return parser.parse_args(raw_args) + def get_device_type(device_str): if device_str.lower() == "gpu": return ort.OrtHardwareDeviceType.GPU @@ -64,6 +65,7 @@ def get_device_type(device_str): else: return ort.OrtHardwareDeviceType.CPU + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -80,6 +82,7 @@ def main(raw_args=None): model_dir = Path(args.script_dir) / "model" / args.model_dir / args.model_id from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -121,12 +124,13 @@ def main(raw_args=None): metrics = { "text-encoder-latency-avg": text_encoder_latency_avg, "unet-latency-avg": unet_latency_avg, - "vae-decoder-latency-avg": vae_decoder_latency_avg + "vae-decoder-latency-avg": vae_decoder_latency_avg, } resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) + if __name__ == "__main__": main() diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_ov_workflow.py b/sd2-community-stable-diffusion-2-1/aitk/sd_ov_workflow.py index 0cca41f59..56f44de15 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_ov_workflow.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_ov_workflow.py @@ -1,13 +1,14 @@ import argparse import json +import logging import os import subprocess import sys -import logging logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -15,11 +16,9 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() -def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + +def load_update_config(config_path: str, cache_dir: str, output_dir: str) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -27,29 +26,27 @@ def load_update_config( return oliveJson -def copy_olive_config( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str): + +def copy_olive_config(history_folder: str, config_path: str, cache_dir: str, output_dir: str): logger.info(f"Copying {config_path} to {history_folder}...") oliveJson = load_update_config(config_path, cache_dir, output_dir) # save updated config for record config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) + def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) # For static quantization, data should match the target scenario. static_shape = oliveJson["passes"]["aitkpython"]["static_shape"] - guidance_scale=str(7.5) - num_inference_steps=str(25) + guidance_scale = str(7.5) + num_inference_steps = str(25) if args.model_config: model_path: str = os.path.dirname(args.model_config) @@ -58,20 +55,33 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "sd_ov_evaluation.py", - "--script_dir", os.path.dirname(model_path), - "--model_dir", "optimized", - "--model_id", "sd2-community/stable-diffusion-2-1", - "--guidance_scale", guidance_scale, - "--num_inference_steps", num_inference_steps, - "--image_size", "512" if static_shape else "768", - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "sd_ov_evaluation.py", + "--script_dir", + os.path.dirname(model_path), + "--model_dir", + "optimized", + "--model_id", + "sd2-community/stable-diffusion-2-1", + "--guidance_scale", + guidance_scale, + "--num_inference_steps", + num_inference_steps, + "--image_size", + "512" if static_shape else "768", + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return - # Get arguments output_dir: str = oliveJson["output_dir"] cache_dir: str = oliveJson["cache_dir"] @@ -88,20 +98,28 @@ def main(): copy_olive_config(history_folder, config_name, cache_dir, output_dir) cmd = [ - sys.executable, "stable_diffusion.py", - "--script_dir", history_folder, - "--model_id", "sd2-community/stable-diffusion-2-1", - "--provider", "openvino", - "--optimize" + sys.executable, + "stable_diffusion.py", + "--script_dir", + history_folder, + "--model_id", + "sd2-community/stable-diffusion-2-1", + "--provider", + "openvino", + "--optimize", ] if static_shape: - cmd.extend([ - "--static_shape", - "--image_size", "512", - ]) + cmd.extend( + [ + "--static_shape", + "--image_size", + "512", + ] + ) # run stable_diffusion.py to generate onnx model subprocess.run(cmd, check=True) + if __name__ == "__main__": main() diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_evaluation.py b/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_evaluation.py index 777f14821..df7fc1d56 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_evaluation.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_evaluation.py @@ -4,18 +4,18 @@ # -------------------------------------------------------------------------- import json -import os import logging +import os from pathlib import Path import numpy as np import onnxruntime as ort - from sd_utils.qdq import OnnxStableDiffusionPipelineWithSave logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_args(raw_args): import argparse @@ -56,6 +56,7 @@ def parse_args(raw_args): ) return parser.parse_args(raw_args) + def get_device_type(device_str): if device_str.lower() == "gpu": return ort.OrtHardwareDeviceType.GPU @@ -64,6 +65,7 @@ def get_device_type(device_str): else: return ort.OrtHardwareDeviceType.CPU + def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): ep_devices = ort.get_ep_devices() for ep_device in ep_devices: @@ -80,6 +82,7 @@ def main(raw_args=None): model_dir = Path(args.script_dir) / "model" / args.model_dir / args.model_id from winml import register_execution_providers + register_execution_providers() sess_options = ort.SessionOptions() @@ -122,12 +125,13 @@ def main(raw_args=None): metrics = { "text-encoder-latency-avg": text_encoder_latency_avg, "unet-latency-avg": unet_latency_avg, - "vae-decoder-latency-avg": vae_decoder_latency_avg + "vae-decoder-latency-avg": vae_decoder_latency_avg, } resultStr = json.dumps(metrics, indent=4) - with open(args.output_file, 'w') as file: + with open(args.output_file, "w") as file: file.write(resultStr) logger.info("Model lab succeeded for evaluation.\n%s", resultStr) + if __name__ == "__main__": main() diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_workflow.py b/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_workflow.py index ea6186464..b702177d6 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_workflow.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_qnn_workflow.py @@ -1,13 +1,14 @@ import argparse import json +import logging import os import subprocess import sys -import logging logger = logging.getLogger(os.path.basename(__file__)) logging.basicConfig(level=logging.INFO) + def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument("--config", required=True, help="path to input config file") @@ -15,13 +16,11 @@ def parse_arguments(): parser.add_argument("--runtime", required=True, help="runtime") return parser.parse_args() + def load_update_config( - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None) -> dict: - with open(config_path, 'r', encoding='utf-8') as file: + config_path: str, cache_dir: str, output_dir: str, activation_type: str | None = None, precision: str | None = None +) -> dict: + with open(config_path, "r", encoding="utf-8") as file: oliveJson = json.load(file) oliveJson["cache_dir"] = cache_dir @@ -35,30 +34,33 @@ def load_update_config( return oliveJson + def copy_olive_config( - history_folder: str, - config_path: str, - cache_dir: str, - output_dir: str, - activation_type: str | None = None, - precision: str | None = None): + history_folder: str, + config_path: str, + cache_dir: str, + output_dir: str, + activation_type: str | None = None, + precision: str | None = None, +): logger.info(f"Copying {config_path} to {history_folder}...") oliveJson = load_update_config(config_path, cache_dir, output_dir, activation_type, precision) # save updated config for record config_name = os.path.basename(config_path) os.makedirs(history_folder, exist_ok=True) - with open(os.path.join(history_folder, config_name), 'w', encoding='utf-8') as file: + with open(os.path.join(history_folder, config_name), "w", encoding="utf-8") as file: json.dump(oliveJson, file, indent=4) + def main(): args = parse_arguments() - with open(args.config, 'r', encoding='utf-8') as file: + with open(args.config, "r", encoding="utf-8") as file: oliveJson = json.load(file) # For static quantization, the QDQ data should match the target scenario. - guidance_scale=str(7.5) - num_inference_steps=str(25) + guidance_scale = str(7.5) + num_inference_steps = str(25) if args.model_config: model_path: str = os.path.dirname(args.model_config) @@ -67,20 +69,33 @@ def main(): output_file = os.path.join(os.path.dirname(args.config), "metrics.json") # Run evaluator - subprocess.run([sys.executable, "sd_qnn_evaluation.py", - "--script_dir", os.path.dirname(model_path), - "--model_dir", "optimized", - "--model_id", "sd2-community/stable-diffusion-2-1", - "--guidance_scale", guidance_scale, - "--num_inference_steps", num_inference_steps, - "--image_size", "512", - "--execution_provider", execution_provider, - "--device_str", device_str, - "--output_file", output_file], - check=True) + subprocess.run( + [ + sys.executable, + "sd_qnn_evaluation.py", + "--script_dir", + os.path.dirname(model_path), + "--model_dir", + "optimized", + "--model_id", + "sd2-community/stable-diffusion-2-1", + "--guidance_scale", + guidance_scale, + "--num_inference_steps", + num_inference_steps, + "--image_size", + "512", + "--execution_provider", + execution_provider, + "--device_str", + device_str, + "--output_file", + output_file, + ], + check=True, + ) return - # Get arguments output_dir: str = oliveJson["output_dir"] cache_dir: str = oliveJson["cache_dir"] @@ -104,44 +119,82 @@ def main(): copy_olive_config(history_folder, config_name, cache_dir, output_dir, activation_type, precision) # run stable_diffusion.py to generate onnx unoptimized model - subprocess.run([sys.executable, "stable_diffusion.py", - "--script_dir", history_folder, - "--model_id", "sd2-community/stable-diffusion-2-1", - "--provider", "cpu", - "--format", "qdq", - "--optimize", - "--image_size", "512", - "--only_conversion"], - check=True) + subprocess.run( + [ + sys.executable, + "stable_diffusion.py", + "--script_dir", + history_folder, + "--model_id", + "sd2-community/stable-diffusion-2-1", + "--provider", + "cpu", + "--format", + "qdq", + "--optimize", + "--image_size", + "512", + "--only_conversion", + ], + check=True, + ) # # run evaluation.py to generate data data_dir = f"../../quantize_data/{guidance_scale}_{num_inference_steps}" - subprocess.run([sys.executable, "evaluation.py", - "--script_dir", history_folder, - "--save_data", - "--model_id", "sd2-community/stable-diffusion-2-1", - "--num_inference_steps", num_inference_steps, - "--seed", "0", - "--data_dir", data_dir, - "--dataset_name", dataset_name, - "--dataset_split", dataset_split, - "--num_data", str(num_data), - # generate data for quantization calibration, so use all data as training data to match num_data - "--train_ratio", "1", - "--guidance_scale", guidance_scale, - "--image_size", "512"], - check=True) + subprocess.run( + [ + sys.executable, + "evaluation.py", + "--script_dir", + history_folder, + "--save_data", + "--model_id", + "sd2-community/stable-diffusion-2-1", + "--num_inference_steps", + num_inference_steps, + "--seed", + "0", + "--data_dir", + data_dir, + "--dataset_name", + dataset_name, + "--dataset_split", + dataset_split, + "--num_data", + str(num_data), + # generate data for quantization calibration, so use all data as training data to match num_data + "--train_ratio", + "1", + "--guidance_scale", + guidance_scale, + "--image_size", + "512", + ], + check=True, + ) # run stable_diffusion.py to generate onnx quantized model - subprocess.run([sys.executable, "stable_diffusion.py", - "--script_dir", history_folder, - "--model_id", "sd2-community/stable-diffusion-2-1", - "--provider", "cpu", - "--format", "qdq", - "--data_dir", data_dir + "/data", - "--optimize", - "--image_size", "512"], - check=True) + subprocess.run( + [ + sys.executable, + "stable_diffusion.py", + "--script_dir", + history_folder, + "--model_id", + "sd2-community/stable-diffusion-2-1", + "--provider", + "cpu", + "--format", + "qdq", + "--data_dir", + data_dir + "/data", + "--optimize", + "--image_size", + "512", + ], + check=True, + ) + if __name__ == "__main__": main() diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/onnx_patch.py b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/onnx_patch.py index 0570ef67f..253c4dcef 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/onnx_patch.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/onnx_patch.py @@ -4,16 +4,15 @@ # -------------------------------------------------------------------------- import os -import time import shutil +import time from pathlib import Path from typing import Optional, Union -from diffusers import OnnxRuntimeModel -from diffusers.utils import ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME - import numpy as np import onnxruntime as ort +from diffusers import OnnxRuntimeModel +from diffusers.utils import ONNX_EXTERNAL_WEIGHTS_NAME, ONNX_WEIGHTS_NAME class PatchedOnnxRuntimeModel(OnnxRuntimeModel): @@ -78,7 +77,6 @@ def save_pretrained( self._save_pretrained(save_directory, **kwargs) - @classmethod def _from_pretrained( cls, diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ort.py b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ort.py index 3f9c8e9ba..0bb70acb0 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ort.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ort.py @@ -9,7 +9,6 @@ from pathlib import Path import onnxruntime as ort - from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from olive.model import ONNXModelHandler from onnxruntime import __version__ as OrtVersion @@ -103,7 +102,9 @@ def save_onnx_pipeline( from sd_utils.onnx_patch import PatchedOnnxRuntimeModel if has_safety_checker: - safety_checker = PatchedOnnxRuntimeModel.from_pretrained(model_info["safety_checker"]["unoptimized"]["path"].parent) + safety_checker = PatchedOnnxRuntimeModel.from_pretrained( + model_info["safety_checker"]["unoptimized"]["path"].parent + ) else: safety_checker = None @@ -112,7 +113,9 @@ def save_onnx_pipeline( vae_decoder=PatchedOnnxRuntimeModel.from_pretrained(model_info["vae_decoder"]["unoptimized"]["path"].parent), text_encoder=PatchedOnnxRuntimeModel.from_pretrained(model_info["text_encoder"]["unoptimized"]["path"].parent), tokenizer=pipeline.tokenizer, - unet=PatchedOnnxRuntimeModel.from_pretrained(model_info["unet"]["unoptimized"]["path"].parent, provider_options=[]), + unet=PatchedOnnxRuntimeModel.from_pretrained( + model_info["unet"]["unoptimized"]["path"].parent, provider_options=[] + ), scheduler=pipeline.scheduler, safety_checker=safety_checker, feature_extractor=pipeline.feature_extractor, diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ov.py b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ov.py index 3608f3b10..72ef94e45 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ov.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/ov.py @@ -4,19 +4,16 @@ # Example modified from: https://docs.openvino.ai/2023.3/notebooks/225-stable-diffusion-text-to-image-with-output.html # -------------------------------------------------------------------------- import inspect -import os from pathlib import Path from typing import Callable, Optional, Union import numpy as np import onnxruntime as ort import torch - from diffusers import StableDiffusionPipeline -from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput - +from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from sd_utils.onnx_patch import PatchedOnnxRuntimeModel @@ -161,8 +158,8 @@ def update_ov_config(config: dict, static_shape: bool): config["passes"] = { "ov_convert": config["passes"]["ov_convert"], "ov_io_update": config["passes"]["ov_io_update"], - "ov_encapsulation": config["passes"]["ov_encapsulation"] - } + "ov_encapsulation": config["passes"]["ov_encapsulation"], + } if static_shape == False: config["passes"]["ov_io_update"]["static"] = False @@ -206,6 +203,7 @@ def get_ov_pipeline(common_args, ov_args, optimized_model_dir): return StableDiffusionPipeline.from_pretrained(common_args.model_id) from winml import register_execution_providers + register_execution_providers() print("Loading models into ORT session...") diff --git a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/qdq.py b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/qdq.py index 13537ea19..c670eaa95 100644 --- a/sd2-community-stable-diffusion-2-1/aitk/sd_utils/qdq.py +++ b/sd2-community-stable-diffusion-2-1/aitk/sd_utils/qdq.py @@ -12,10 +12,9 @@ import sd_utils import sd_utils.config import torch - -from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput +from diffusers.pipelines.stable_diffusion.pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline # ruff: noqa: T201 @@ -266,6 +265,7 @@ def add_ep_for_device(session_options, ep_name, device_type, ep_options=None): def get_qdq_pipeline(model_dir, common_args, qdq_args, script_dir): if common_args.provider != "cpu": from winml import register_execution_providers + register_execution_providers() ort.set_default_logger_severity(3) From 4e4894074e273e0e9a2ad037044a07c9150ae6e0 Mon Sep 17 00:00:00 2001 From: xieofxie Date: Mon, 25 May 2026 14:45:25 +0800 Subject: [PATCH 3/3] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .aitk/scripts/ruff_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.aitk/scripts/ruff_check.py b/.aitk/scripts/ruff_check.py index 7f5e7c916..c2af1dc34 100644 --- a/.aitk/scripts/ruff_check.py +++ b/.aitk/scripts/ruff_check.py @@ -72,7 +72,7 @@ def ruff_check(fix: bool = False): format_result = subprocess.run(format_cmd + target_args) if check_result.returncode != 0 or format_result.returncode != 0: - printError("Ruff reported issues. Re-run `python sanitize.py --format_scripts` to auto-fix.") + printError("Ruff reported issues. Re-run `python .aitk/scripts/ruff_check.py --fix` to auto-fix.") raise SystemExit(1) printInfo("Ruff check passed.")