diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f8248c37..92c762a7 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,8 +23,8 @@ jobs: echo "$HOME/.local/bin" >> "$GITHUB_PATH" # Run the ruff pinned in pyproject.toml's lint group (also used by pre-commit), # so CI doesn't change behaviour whenever a new ruff is released. - - run: uv run --frozen --only-group lint ruff check python/src/ python/tests/ - - run: uv run --frozen --only-group lint ruff format --check python/src/ python/tests/ + - run: uv run --frozen --only-group lint ruff check python/src/ python/tests/ models/ + - run: uv run --frozen --only-group lint ruff format --check python/src/ python/tests/ models/ swift-format: if: github.repository == 'apple/coreai-models' diff --git a/models/clap/export.py b/models/clap/export.py index 58dab373..086af38a 100644 --- a/models/clap/export.py +++ b/models/clap/export.py @@ -97,18 +97,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -122,7 +118,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "Y. Wu et al." metadata.license = "Apache-2.0" - metadata.model_description = "CLAP (Contrastive Language-Audio Pretraining) learns joint representations of audio and text, enabling zero-shot audio classification with natural language labels. Source: https://huggingface.co/laion/clap-htsat-unfused" + metadata.model_description = ( + "CLAP (Contrastive Language-Audio Pretraining) learns joint representations of audio and " + "text, enabling zero-shot audio classification with natural language labels. Source: " + "https://huggingface.co/laion/clap-htsat-unfused" + ) metadata.creation_date = int(time.time()) return metadata @@ -145,15 +145,11 @@ def create_clap( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["input_ids", "attention_mask", "input_features", "is_longer"], @@ -173,9 +169,7 @@ def create_clap( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for CLAP." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for CLAP.") parser.add_argument( "--model", choices=["laion/clap-htsat-unfused"], @@ -207,7 +201,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/clip/export.py b/models/clip/export.py index c881bde0..e05b94c7 100644 --- a/models/clip/export.py +++ b/models/clip/export.py @@ -103,18 +103,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -128,7 +124,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "A. Radford et al." metadata.license = "MIT" - metadata.model_description = "CLIP (Contrastive Language-Image Pretraining) learns joint representations of images and text, enabling zero-shot image classification with natural language labels. Source: https://huggingface.co/openai/clip-vit-base-patch32" + metadata.model_description = ( + "CLIP (Contrastive Language-Image Pretraining) learns joint representations of images and " + "text, enabling zero-shot image classification with natural language labels. Source: " + "https://huggingface.co/openai/clip-vit-base-patch32" + ) metadata.creation_date = int(time.time()) return metadata @@ -151,15 +151,11 @@ def create_clip( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["pixel_values", "input_ids", "attention_mask"], @@ -179,9 +175,7 @@ def create_clip( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for CLIP." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for CLIP.") parser.add_argument( "--model", choices=["openai/clip-vit-base-patch32"], @@ -213,7 +207,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/depth-anything/export.py b/models/depth-anything/export.py index 8644576c..bf141575 100644 --- a/models/depth-anything/export.py +++ b/models/depth-anything/export.py @@ -47,9 +47,7 @@ def _patch_depth_anything_for_export() -> None: # `RotaryPositionEmbedding2D.forward` uses `int(positions.max()) + 1`, a # data-dependent guard that breaks torch.export. Swap in a shape-based # equivalent. - def _rope_forward( - self, tokens: torch.Tensor, positions: torch.Tensor - ) -> torch.Tensor: + def _rope_forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor: assert tokens.size(-1) % 2 == 0 assert positions.ndim == 3 and positions.shape[-1] == 2 feature_dim = tokens.size(-1) // 2 @@ -59,9 +57,7 @@ def _rope_forward( ) vertical, horizontal = tokens.chunk(2, dim=-1) vertical = self._apply_1d_rope(vertical, positions[..., 0], cos_comp, sin_comp) - horizontal = self._apply_1d_rope( - horizontal, positions[..., 1], cos_comp, sin_comp - ) + horizontal = self._apply_1d_rope(horizontal, positions[..., 1], cos_comp, sin_comp) return torch.cat((vertical, horizontal), dim=-1) _rope.RotaryPositionEmbedding2D.forward = _rope_forward @@ -110,9 +106,7 @@ def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype) -> Path: def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -126,7 +120,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "H. Lin et al." metadata.license = "Apache-2.0" - metadata.model_description = "Depth Anything v3 is a monocular depth estimation model that predicts depth, confidence, camera intrinsics, and extrinsics from a batch of image views. Source: https://github.com/ByteDance-Seed/Depth-Anything-3" + metadata.model_description = ( + "Depth Anything v3 is a monocular depth estimation model that predicts depth, confidence, " + "camera intrinsics, and extrinsics from a batch of image views. Source: " + "https://github.com/ByteDance-Seed/Depth-Anything-3" + ) metadata.creation_date = int(time.time()) return metadata @@ -154,9 +152,7 @@ def create_depth_anything( exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["image"], @@ -200,7 +196,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/diffusion_gemma/export.py b/models/diffusion_gemma/export.py index 57fba7b3..fc68227d 100644 --- a/models/diffusion_gemma/export.py +++ b/models/diffusion_gemma/export.py @@ -109,9 +109,7 @@ def export_diffusion_gemma( # encoder_len and surfaced as keyCache/valueCache state. ids0 = torch.zeros(1, encoder_len, dtype=torch.int32) pos0 = torch.arange(encoder_len, dtype=torch.int32).unsqueeze(0) - kc = torch.zeros( - n_layers_eff, 1, cache_n_kv, encoder_len, cache_hd, dtype=dtype - ) + kc = torch.zeros(n_layers_eff, 1, cache_n_kv, encoder_len, cache_hd, dtype=dtype) vc = torch.zeros_like(kc) enc_prog = export_to_coreai( encoder, @@ -156,9 +154,7 @@ def export_diffusion_gemma( # -- Decoder (bidirectional canvas denoiser, self-conditioning folded in) - logger.info("Loading + exporting DECODER ...") - decoder = load_diffusion_gemma_decoder( - hf_model_id, target_dtype=dtype, num_layers=num_layers - ) + decoder = load_diffusion_gemma_decoder(hf_model_id, target_dtype=dtype, num_layers=num_layers) h = text_cfg.hidden_size decoder_input_ids = torch.zeros(1, canvas_length, dtype=torch.int32) prev_soft_embeds = torch.zeros(1, canvas_length, h, dtype=dtype) @@ -274,9 +270,7 @@ def _quantize_encoder(encoder, compression, dtype): preset = get_preset(compression) quant_cfg = preset.get("torch_quantization_config") if quant_cfg is None: - logger.warning( - "Compression '%s' has no torch_quantization_config; skipping", compression - ) + logger.warning("Compression '%s' has no torch_quantization_config; skipping", compression) return encoder quant_cfg = dict(quant_cfg) @@ -332,9 +326,7 @@ def _quantize_decoder(decoder, compression, dec_inputs): preset = get_preset(compression) quant_cfg = preset.get("torch_quantization_config") if quant_cfg is None: - logger.warning( - "Compression '%s' has no torch_quantization_config; skipping", compression - ) + logger.warning("Compression '%s' has no torch_quantization_config; skipping", compression) return decoder quant_cfg = dict(quant_cfg) @@ -405,9 +397,7 @@ def _build_parser() -> argparse.ArgumentParser: help="Encoder KV-cache max context", ) parser.add_argument("--canvas-length", type=int, default=256, help="Canvas length") - parser.add_argument( - "--compression", default="none", help="Compression preset or 'none'" - ) + parser.add_argument("--compression", default="none", help="Compression preset or 'none'") parser.add_argument( "--compute-precision", choices=["float16", "bfloat16", "float32"], @@ -415,9 +405,7 @@ def _build_parser() -> argparse.ArgumentParser: ) parser.add_argument("--output-dir", default="./exports") parser.add_argument("--output-name", default=None) - parser.add_argument( - "--num-layers", type=int, default=None, help="Truncate to N layers (smoke)" - ) + parser.add_argument("--num-layers", type=int, default=None, help="Truncate to N layers (smoke)") parser.add_argument( "--encoder-len", type=int, diff --git a/models/edsr/export.py b/models/edsr/export.py index fc6c1bde..cff5c088 100644 --- a/models/edsr/export.py +++ b/models/edsr/export.py @@ -27,9 +27,7 @@ from coreai_torch import TorchConverter, get_decomp_table -def reference_inputs( - dtype: torch.dtype, dynamic: bool = False -) -> dict[str, torch.Tensor]: +def reference_inputs(dtype: torch.dtype, dynamic: bool = False) -> dict[str, torch.Tensor]: B = 2 if dynamic else 1 return {"x": torch.rand(B, 3, 16, 16).to(dtype)} @@ -58,18 +56,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -83,7 +77,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "B. Lim et al." metadata.license = "MIT" - metadata.model_description = "EDSR (Enhanced Deep Residual Networks for Single Image Super-Resolution) upscales low-resolution images by a fixed integer factor (2x, 3x, 4x). Source: https://github.com/Coloquinte/torchSR" + metadata.model_description = ( + "EDSR (Enhanced Deep Residual Networks for Single Image Super-Resolution) upscales " + "low-resolution images by a fixed integer factor (2x, 3x, 4x). Source: " + "https://github.com/Coloquinte/torchSR" + ) metadata.creation_date = int(time.time()) return metadata @@ -106,15 +104,11 @@ def create_edsr( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["x"], @@ -129,9 +123,7 @@ def create_edsr( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for EDSR." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for EDSR.") parser.add_argument( "--model", choices=["edsr_r16f64_x2"], @@ -163,7 +155,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/efficient-sam/export.py b/models/efficient-sam/export.py index 21f9f666..4bf88599 100644 --- a/models/efficient-sam/export.py +++ b/models/efficient-sam/export.py @@ -88,14 +88,10 @@ def _bundle_paths( return bundle_dir, bundle_dir / f"{variant}.aimodel" -def _save_asset( - coreai_program, bundle_dir: Path, model_path: Path, overwrite: bool -) -> None: +def _save_asset(coreai_program, bundle_dir: Path, model_path: Path, overwrite: bool) -> None: if bundle_dir.exists(): if not overwrite: - raise FileExistsError( - f"{bundle_dir} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{bundle_dir} already exists. Pass --overwrite to replace it.") shutil.rmtree(bundle_dir) bundle_dir.mkdir(parents=True, exist_ok=True) coreai_program.save_asset(model_path, _build_aimodel_metadata()) @@ -106,7 +102,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "Y. Xiong et al." metadata.license = "Apache-2.0" - metadata.model_description = "EfficientSAM is a lightweight, promptable image segmentation model that uses a masked autoencoder pretrained ViT-Tiny encoder to reduce compute while preserving accuracy. Source: https://github.com/yformer/EfficientSAM" + metadata.model_description = ( + "EfficientSAM is a lightweight, promptable image segmentation model that uses a masked " + "autoencoder pretrained ViT-Tiny encoder to reduce compute while preserving accuracy. " + "Source: https://github.com/yformer/EfficientSAM" + ) metadata.creation_date = int(time.time()) return metadata @@ -143,9 +143,7 @@ def create_efficient_sam( ) print("[INFO] Downloading weights and sourcing model...") - checkpoint = ( - "https://huggingface.co/merve/EfficientSAM/resolve/main/efficient_sam_vitt.pt" - ) + checkpoint = "https://huggingface.co/merve/EfficientSAM/resolve/main/efficient_sam_vitt.pt" state_dict = torch.hub.load_state_dict_from_url( checkpoint, map_location="cpu", progress=True, weights_only=True ) @@ -165,15 +163,11 @@ def create_efficient_sam( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["batched_images", "batched_points", "batched_point_labels"], @@ -188,9 +182,7 @@ def create_efficient_sam( _save_asset(coreai_program, bundle_dir, model_path, overwrite) print(f"[INFO] Successfully created and saved Core AI model to {model_path}.") - _write_metadata( - bundle_dir, _variant_name(model_name, dtype, dynamic, num_queries, num_pts) - ) + _write_metadata(bundle_dir, _variant_name(model_name, dtype, dynamic, num_queries, num_pts)) def main(): @@ -248,7 +240,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/gemma4/export.py b/models/gemma4/export.py index f3549478..eac044e9 100644 --- a/models/gemma4/export.py +++ b/models/gemma4/export.py @@ -32,7 +32,7 @@ import shutil import tempfile from pathlib import Path -from typing import Any, Optional +from typing import Any import torch from coreai.authoring import AIProgram @@ -67,9 +67,7 @@ logger = logging.getLogger("gemma4.export") REPO_ROOT = Path(__file__).resolve().parents[2] -DEFAULT_IOS_COMPRESSION_CONFIG = ( - Path(__file__).resolve().parent / "4bit_palettized.yaml" -) +DEFAULT_IOS_COMPRESSION_CONFIG = Path(__file__).resolve().parent / "4bit_palettized.yaml" # Checkpoints whose exports have been accuracy- and performance-verified. These # are the ones the model card advertises as supported. Other Gemma 4 checkpoints @@ -142,9 +140,7 @@ def _ios_metadata_extras(text_config: Any) -> dict[str, Any]: "global_head_dim": text_config.global_head_dim, "sliding_rope_theta": rope["sliding_attention"]["rope_theta"], "global_rope_theta": rope["full_attention"]["rope_theta"], - "partial_rotary_factor": rope["full_attention"][ - "partial_rotary_factor" - ], + "partial_rotary_factor": rope["full_attention"]["partial_rotary_factor"], }, } except (AttributeError, KeyError, TypeError) as exc: @@ -162,8 +158,8 @@ def _patch_language_metadata( bundle_path: Path, hf_model_id: str, text_config: Any, - extras: Optional[dict[str, Any]] = None, - auxiliary_assets: Optional[dict[str, Any]] = None, + extras: dict[str, Any] | None = None, + auxiliary_assets: dict[str, Any] | None = None, ) -> None: """Add Gemma4-specific keys to a written bundle. @@ -200,18 +196,14 @@ def _text_config(hf_model_id: str) -> Any: return getattr(raw_config, HF_CONFIG_ATTR, raw_config) -def _resolve_bundle_paths( - output_dir: str, output_name: str, overwrite: bool -) -> tuple[Path, Path]: +def _resolve_bundle_paths(output_dir: str, output_name: str, overwrite: bool) -> tuple[Path, Path]: """Create the bundle directory, failing fast on an existing asset without ``--overwrite``. Returns (bundle, asset). The existing asset is only removed just before the new one is saved, so a failed export leaves it in place.""" bundle_path = Path(output_dir) / output_name aimodel_path = bundle_path / f"{output_name}.aimodel" if aimodel_path.exists() and not overwrite: - raise SystemExit( - f"{aimodel_path} already exists. Use --overwrite to replace it." - ) + raise SystemExit(f"{aimodel_path} already exists. Use --overwrite to replace it.") bundle_path.mkdir(parents=True, exist_ok=True) return bundle_path, aimodel_path @@ -300,12 +292,8 @@ def _export_programs( decomp_table = _ios_decomp_table() block_size = model.extend.model.kv_block_size - def trace( - module, kwargs: dict, dynamic_shapes=None - ) -> torch.export.ExportedProgram: - return torch.export.export( - module, args=(), kwargs=kwargs, dynamic_shapes=dynamic_shapes - ) + def trace(module, kwargs: dict, dynamic_shapes=None) -> torch.export.ExportedProgram: + return torch.export.export(module, args=(), kwargs=kwargs, dynamic_shapes=dynamic_shapes) programs: list[tuple[str, str, int, torch.export.ExportedProgram]] = [] with torch.no_grad(): @@ -351,8 +339,7 @@ def trace( for q in qlens: name = f"{entrypoint}_{ctx}_{q}" logger.info( - f"Exporting {name} " - f"(flash chunks={(ctx + block_size - 1) // block_size})..." + f"Exporting {name} (flash chunks={(ctx + block_size - 1) // block_size})..." ) ref, _ = _reference_inputs( model, @@ -389,9 +376,7 @@ async def _convert_to_coreai( states = model.export_state_names() outputs = model.export_output_names() - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode) register_custom_torch_lowering(converter) for name, graph, _, program in programs: @@ -434,14 +419,10 @@ async def _export_blocked_ladder( logger.info(f"iOS context ladder: {buckets}") programs = _export_programs(model, config, buckets) gather_qlens = set(SHIPPING_EXTEND_QLENS) | set(SHIPPING_PROMPT_QLENS) - return await _convert_to_coreai( - model, programs, config, gather_qlens, include_debug_info - ) + return await _convert_to_coreai(model, programs, config, gather_qlens, include_debug_info) -def _palettization_inputs( - model: Gemma4ForCausalLMForiOS, config, max_context_length: int -) -> tuple: +def _palettization_inputs(model: Gemma4ForCausalLMForiOS, config, max_context_length: int) -> tuple: """The palettizer's calibration inputs: the smallest rung's reference inputs, as ``model.forward``'s positional arguments. @@ -479,9 +460,7 @@ async def _export_ios(args: argparse.Namespace) -> str: # unset. Resolved before any weights are loaded so a bad recipe fails fast. palettization_config = None if args.compression_config is not None: - palettization_config = _load_compression_config_object( - args.compression_config, "iOS" - ) + palettization_config = _load_compression_config_object(args.compression_config, "iOS") compression = args.compression_config.stem else: compression = "none" @@ -494,13 +473,9 @@ async def _export_ios(args: argparse.Namespace) -> str: compression_config_object=palettization_config, ) ) - bundle_path, aimodel_path = _resolve_bundle_paths( - args.output_dir, output_name, args.overwrite - ) + bundle_path, aimodel_path = _resolve_bundle_paths(args.output_dir, output_name, args.overwrite) - logger.info( - f"Loading {hf_model_id} (iOS, dtype={target_dtype}, max_ctx={max_ctx})..." - ) + logger.info(f"Loading {hf_model_id} (iOS, dtype={target_dtype}, max_ctx={max_ctx})...") # Move loaded weights to disk-backed mmap tensors so the OS can evict weight # pages during palettization and the long blocked-ladder conversion. The temp @@ -617,8 +592,7 @@ def build_parser() -> argparse.ArgumentParser: "--max-context-length", type=int, default=None, - help="Maximum context length, a power of two " - f"(default and max: {IOS_MAX_CONTEXT_LENGTH})", + help=f"Maximum context length, a power of two (default and max: {IOS_MAX_CONTEXT_LENGTH})", ) parser.add_argument( "--output-dir", @@ -630,9 +604,7 @@ def build_parser() -> argparse.ArgumentParser: default=None, help="Custom bundle name (without extension)", ) - parser.add_argument( - "--overwrite", action="store_true", help="Overwrite an existing bundle" - ) + parser.add_argument("--overwrite", action="store_true", help="Overwrite an existing bundle") parser.add_argument( "--include-debug-info", action="store_true", @@ -640,9 +612,7 @@ def build_parser() -> argparse.ArgumentParser: "Default: off, which embeds minimum debug information and makes the exported " "asset smaller.", ) - parser.add_argument( - "--verbose", "-v", action="store_true", help="Enable DEBUG logging" - ) + parser.add_argument("--verbose", "-v", action="store_true", help="Enable DEBUG logging") return parser @@ -655,10 +625,7 @@ def _resolve_defaults(args: argparse.Namespace) -> None: recipe. ``--compression`` only accepts ``none``, which skips compression entirely. """ - if ( - args.max_context_length is not None - and args.max_context_length > IOS_MAX_CONTEXT_LENGTH - ): + if args.max_context_length is not None and args.max_context_length > IOS_MAX_CONTEXT_LENGTH: raise SystemExit( f"--max-context-length supports at most {IOS_MAX_CONTEXT_LENGTH} tokens " f"(got {args.max_context_length}); the static-shape ladder tops out " @@ -667,8 +634,7 @@ def _resolve_defaults(args: argparse.Namespace) -> None: # The top context bucket is --max-context-length itself, so it must be one. if args.max_context_length is not None and ( - args.max_context_length <= 0 - or args.max_context_length & (args.max_context_length - 1) + args.max_context_length <= 0 or args.max_context_length & (args.max_context_length - 1) ): raise SystemExit( f"--max-context-length must be a power of two (got {args.max_context_length})." diff --git a/models/parakeet/export.py b/models/parakeet/export.py index e2d2bd90..70c3a658 100644 --- a/models/parakeet/export.py +++ b/models/parakeet/export.py @@ -32,7 +32,6 @@ from coreai.runtime import AIModelAssetMetadata from coreai_torch import TorchConverter, get_decomp_table - # Parakeet TDT exports as three separate graphs because the autoregressive # transducer loop (encoder frame pointer + (token, duration) sampling) lives # in ParakeetTDTGenerationMixin, not in `forward`, and torch.export cannot @@ -63,9 +62,7 @@ def __init__(self, model: "transformers.ParakeetForTDT"): self._encoder = model.encoder self._encoder_projector = model.encoder_projector - def forward( - self, input_features: torch.Tensor, attention_mask: torch.Tensor - ) -> torch.Tensor: + def forward(self, input_features: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: # attention_mask is a (B, T_audio) bool mask marking real-audio frames. # It makes the encoder exclude padding from self-attention *and* the # subsampling / conformer conv modules (matching HF). Without it a @@ -99,9 +96,7 @@ def forward( cell_state: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: embeddings = self._embedding(input_ids) - lstm_output, (new_hidden, new_cell) = self._lstm( - embeddings, (hidden_state, cell_state) - ) + lstm_output, (new_hidden, new_cell) = self._lstm(embeddings, (hidden_state, cell_state)) decoder_output = self._projector(lstm_output) return decoder_output, new_hidden, new_cell @@ -155,9 +150,7 @@ def _encoder_frame_count(mel_frames: int, subsampling_factor: int) -> int: if mel_frames <= 0 or subsampling_factor <= 1: return max(0, mel_frames) if subsampling_factor & (subsampling_factor - 1) != 0: - raise ValueError( - f"subsampling_factor must be a power of two, got {subsampling_factor}" - ) + raise ValueError(f"subsampling_factor must be a power of two, got {subsampling_factor}") length, factor = mel_frames, subsampling_factor while factor > 1: length = (length - 1) // 2 + 1 @@ -314,9 +307,7 @@ def _convert( dynamic_shapes=dynamic_shapes, ) exported = exported.run_decompositions(get_decomp_table()) - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=input_names, @@ -387,19 +378,13 @@ def _save_program(program, model_path: Path, graph: str) -> None: def _prepare_bundle_dir(bundle_dir: Path, overwrite: bool) -> None: if bundle_dir.exists(): if not overwrite: - raise FileExistsError( - f"{bundle_dir} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{bundle_dir} already exists. Pass --overwrite to replace it.") shutil.rmtree(bundle_dir) bundle_dir.mkdir(parents=True, exist_ok=True) -def _write_processor( - dest: Path, processor: "transformers.ProcessorMixin", model_name: str -) -> None: - print( - f"[INFO] Saving processor (feature extractor + tokenizer) from {model_name} to {dest}..." - ) +def _write_processor(dest: Path, processor: "transformers.ProcessorMixin", model_name: str) -> None: + print(f"[INFO] Saving processor (feature extractor + tokenizer) from {model_name} to {dest}...") processor.save_pretrained(str(dest)) @@ -461,9 +446,7 @@ def _streaming_encoder_inputs( one-frame slice error is 80 ms of audio and would drop or duplicate words at every chunk boundary. Costs one forward pass. """ - features = _audio_features_samples( - processor, dtype, geometry["window_sample_count"] - ) + features = _audio_features_samples(processor, dtype, geometry["window_sample_count"]) if features.shape[1] != geometry["window_mel_frames"]: raise ValueError( f"streaming geometry mismatch: {geometry['window_sample_count']} samples " @@ -537,9 +520,7 @@ def create_parakeet( if geometry is not None: encoder_inputs = _streaming_encoder_inputs(processor, model, dtype, geometry) else: - encoder_inputs = _encoder_inputs( - _audio_features(processor, dtype, audio_seconds) - ) + encoder_inputs = _encoder_inputs(_audio_features(processor, dtype, audio_seconds)) encoder_program = _convert( ParakeetEncoderModule(model), @@ -593,9 +574,7 @@ def create_parakeet( _AUDIO_SECONDS_FLAG = ("--audio-seconds", "audio_seconds") -def _warn_ignored_shape_args( - parser: argparse.ArgumentParser, args: argparse.Namespace -) -> None: +def _warn_ignored_shape_args(parser: argparse.ArgumentParser, args: argparse.Namespace) -> None: """Warn about window flags the chosen shape mode never reads. Each mode sizes the encoder trace a different way, and a flag belonging to @@ -612,11 +591,7 @@ def _warn_ignored_shape_args( candidates = list(_WINDOW_FLAGS) reason = "the frame counts only apply with --streaming" - ignored = [ - flag - for flag, dest in candidates - if getattr(args, dest) != parser.get_default(dest) - ] + ignored = [flag for flag, dest in candidates if getattr(args, dest) != parser.get_default(dest)] if ignored: print(f"[WARN] Ignoring {', '.join(ignored)} — {reason}.") @@ -709,7 +684,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() _warn_ignored_shape_args(parser, args) diff --git a/models/pvt/export.py b/models/pvt/export.py index 5eead8de..09d8bfca 100644 --- a/models/pvt/export.py +++ b/models/pvt/export.py @@ -56,18 +56,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -81,7 +77,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "W. Wang et al." metadata.license = "Apache-2.0" - metadata.model_description = "PVT v2 (Pyramid Vision Transformer v2) uses a pyramid structure as an effective backbone for dense prediction tasks. Source: https://huggingface.co/docs/transformers/model_doc/pvt" + metadata.model_description = ( + "PVT v2 (Pyramid Vision Transformer v2) uses a pyramid structure as an effective backbone " + "for dense prediction tasks. Source: " + "https://huggingface.co/docs/transformers/model_doc/pvt" + ) metadata.creation_date = int(time.time()) return metadata @@ -104,15 +104,11 @@ def create_pvt( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["x"], @@ -127,9 +123,7 @@ def create_pvt( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for PVT." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for PVT.") parser.add_argument( "--model", choices=["pvt_v2_b0"], @@ -161,7 +155,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/roberta/export.py b/models/roberta/export.py index 31a51d89..0b4b6990 100644 --- a/models/roberta/export.py +++ b/models/roberta/export.py @@ -59,18 +59,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -84,7 +80,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "Y. Liu et al." metadata.license = "MIT" - metadata.model_description = "RoBERTa is a transformer encoder model that improves on BERT by training with larger batches, more data, longer sequences, and without the next-sentence prediction objective. Source: https://huggingface.co/roberta-base" + metadata.model_description = ( + "RoBERTa is a transformer encoder model that improves on BERT by training with larger " + "batches, more data, longer sequences, and without the next-sentence prediction objective. " + "Source: https://huggingface.co/roberta-base" + ) metadata.creation_date = int(time.time()) return metadata @@ -98,9 +98,7 @@ def create_roberta( include_debug_info: bool, ): print("[INFO] Sourcing model...") - model = transformers.RobertaModel.from_pretrained( - model_name, add_pooling_layer=False - ) + model = transformers.RobertaModel.from_pretrained(model_name, add_pooling_layer=False) model.eval() model.to(dtype) print("[INFO] Model sourced. Running torch export with decompositions...") @@ -109,15 +107,11 @@ def create_roberta( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["input_ids"], @@ -132,9 +126,7 @@ def create_roberta( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for RoBERTa." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for RoBERTa.") parser.add_argument( "--model", choices=["roberta-base"], @@ -166,7 +158,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/ruff.toml b/models/ruff.toml new file mode 100644 index 00000000..dd0bf749 --- /dev/null +++ b/models/ruff.toml @@ -0,0 +1,2 @@ +# models/ is held to the same ruff rules as python/ (line length, rule set, isort). +extend = "../python/pyproject.toml" diff --git a/models/t5/export.py b/models/t5/export.py index c193f347..02cbdf7b 100644 --- a/models/t5/export.py +++ b/models/t5/export.py @@ -44,9 +44,7 @@ def reference_inputs( ) -> dict[str, torch.Tensor]: tokenizer = transformers.AutoTokenizer.from_pretrained(model_name) input_ids = tokenizer("The weather is nice today", return_tensors="pt")["input_ids"] - decoder_input_ids = tokenizer("Il fait beau aujourd'hui", return_tensors="pt")[ - "input_ids" - ] + decoder_input_ids = tokenizer("Il fait beau aujourd'hui", return_tensors="pt")["input_ids"] inputs = { "input_ids": input_ids.to(torch.int32), "decoder_input_ids": decoder_input_ids.to(torch.int32), @@ -99,18 +97,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -124,7 +118,12 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "C. Raffel et al." metadata.license = "Apache-2.0" - metadata.model_description = "T5 (Text-to-Text Transfer Transformer) is an encoder-decoder model pre-trained on a mixture of unsupervised and supervised tasks. Works well on many tasks via input prefixes (e.g. 'translate English to German: ...', 'summarize: ...'). Source: https://huggingface.co/docs/transformers/model_doc/t5" + metadata.model_description = ( + "T5 (Text-to-Text Transfer Transformer) is an encoder-decoder model pre-trained on a " + "mixture of unsupervised and supervised tasks. Works well on many tasks via input prefixes " + "(e.g. 'translate English to German: ...', 'summarize: ...'). Source: " + "https://huggingface.co/docs/transformers/model_doc/t5" + ) metadata.creation_date = int(time.time()) return metadata @@ -147,15 +146,11 @@ def create_t5( ds = dynamic_shapes(dtype) if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["input_ids", "decoder_input_ids"], @@ -170,9 +165,7 @@ def create_t5( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for T5." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for T5.") parser.add_argument( "--model", choices=["google-t5/t5-small", "google-t5/t5-base", "google-t5/t5-large"], @@ -204,7 +197,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/wav2vec2/export.py b/models/wav2vec2/export.py index b906ffd5..cfb46795 100644 --- a/models/wav2vec2/export.py +++ b/models/wav2vec2/export.py @@ -37,9 +37,7 @@ def forward(self, waveform): return emission -def reference_inputs( - dtype: torch.dtype, dynamic: bool = False -) -> dict[str, torch.Tensor]: +def reference_inputs(dtype: torch.dtype, dynamic: bool = False) -> dict[str, torch.Tensor]: B = 2 if dynamic else 1 # 16kHz mono audio, ~5 seconds. return {"waveform": torch.randn(B, 80000).to(dtype)} @@ -66,18 +64,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -91,7 +85,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "A. Baevski et al." metadata.license = "MIT" - metadata.model_description = "Wav2Vec 2.0 is a self-supervised speech representation model that learns directly from raw audio and, after fine-tuning, transcribes speech into character-level token emissions. Source: https://pytorch.org/audio/stable/pipelines.html" + metadata.model_description = ( + "Wav2Vec 2.0 is a self-supervised speech representation model that learns directly from " + "raw audio and, after fine-tuning, transcribes speech into character-level token " + "emissions. Source: https://pytorch.org/audio/stable/pipelines.html" + ) metadata.creation_date = int(time.time()) return metadata @@ -114,15 +112,11 @@ def create_wav2vec2( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["waveform"], @@ -171,7 +165,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() diff --git a/models/whisper/export.py b/models/whisper/export.py index 8423daaf..c233e6b4 100644 --- a/models/whisper/export.py +++ b/models/whisper/export.py @@ -38,9 +38,7 @@ def __init__(self, model_name: str, dtype: torch.dtype): ) def forward(self, input_features, decoder_input_ids): - outputs = self._model( - input_features=input_features, decoder_input_ids=decoder_input_ids - ) + outputs = self._model(input_features=input_features, decoder_input_ids=decoder_input_ids) return outputs.logits @@ -73,9 +71,7 @@ def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype) -> Path: def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -89,7 +85,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "A. Radford et al." metadata.license = "Apache-2.0" - metadata.model_description = "Whisper is an automatic speech recognition (ASR) encoder-decoder model from OpenAI, trained on a large multilingual and multitask supervised dataset. Source: https://huggingface.co/openai/whisper-large-v3" + metadata.model_description = ( + "Whisper is an automatic speech recognition (ASR) encoder-decoder model from OpenAI, " + "trained on a large multilingual and multitask supervised dataset. Source: " + "https://huggingface.co/openai/whisper-large-v3" + ) metadata.creation_date = int(time.time()) return metadata @@ -122,9 +122,7 @@ def create_whisper( exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["input_features", "decoder_input_ids"], @@ -139,9 +137,7 @@ def create_whisper( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for Whisper." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for Whisper.") parser.add_argument( "--model", choices=["openai/whisper-large-v3-turbo", "openai/whisper-large-v3"], @@ -168,7 +164,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args() @@ -179,9 +176,7 @@ def main(): }[args.dtype] output_dir = args.output_dir or _default_output_dir() - create_whisper( - output_dir, args.model, dtype, args.overwrite, args.include_debug_info - ) + create_whisper(output_dir, args.model, dtype, args.overwrite, args.include_debug_info) if __name__ == "__main__": diff --git a/models/yolo/export.py b/models/yolo/export.py index 9f8d8d5d..dd5697b4 100644 --- a/models/yolo/export.py +++ b/models/yolo/export.py @@ -30,9 +30,7 @@ class YolosModule(torch.nn.Module): def __init__(self, model_name: str): super().__init__() - self._model = transformers.AutoModelForObjectDetection.from_pretrained( - model_name - ) + self._model = transformers.AutoModelForObjectDetection.from_pretrained(model_name) def forward(self, pixel_values): outputs = self._model(pixel_values=pixel_values) @@ -74,18 +72,14 @@ def _variant_name(model_name: str, dtype: torch.dtype, dynamic: bool) -> str: return f"{safe_name}_{dtype_name}_{static_or_dynamic}" -def _asset_path( - output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool -) -> Path: +def _asset_path(output_dir: str, model_name: str, dtype: torch.dtype, dynamic: bool) -> Path: return Path(output_dir) / f"{_variant_name(model_name, dtype, dynamic)}.aimodel" def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: if model_path.exists(): if not overwrite: - raise FileExistsError( - f"{model_path} already exists. Pass --overwrite to replace it." - ) + raise FileExistsError(f"{model_path} already exists. Pass --overwrite to replace it.") if model_path.is_dir(): shutil.rmtree(model_path) else: @@ -99,7 +93,11 @@ def _build_aimodel_metadata() -> AIModelAssetMetadata: metadata = AIModelAssetMetadata() metadata.author = "Y. Fang et al." metadata.license = "Apache-2.0" - metadata.model_description = "YOLOS (You Only Look at One Sequence) applies a plain Vision Transformer directly to image patches and predicts object queries as bounding boxes and class logits. Source: https://huggingface.co/hustvl/yolos-tiny" + metadata.model_description = ( + "YOLOS (You Only Look at One Sequence) applies a plain Vision Transformer directly to " + "image patches and predicts object queries as bounding boxes and class logits. Source: " + "https://huggingface.co/hustvl/yolos-tiny" + ) metadata.creation_date = int(time.time()) return metadata @@ -122,15 +120,11 @@ def create_yolos( ds = dynamic_shapes() if dynamic else None with torch.autocast(device_type="cpu", dtype=dtype): - exported = torch.export.export( - model, args=(), kwargs=example_inputs, dynamic_shapes=ds - ) + exported = torch.export.export(model, args=(), kwargs=example_inputs, dynamic_shapes=ds) exported = exported.run_decompositions(get_decomp_table()) print("[INFO] Model exported. Converting to Core AI...") - mode = ( - TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE - ) + mode = TorchConverter.Mode.DEBUG if include_debug_info else TorchConverter.Mode.RELEASE converter = TorchConverter(mode=mode).add_exported_program( exported_program=exported, input_names=["pixel_values"], @@ -145,9 +139,7 @@ def create_yolos( def main(): - parser = argparse.ArgumentParser( - description="Create and save a Core AI AIProgram for YOLOS." - ) + parser = argparse.ArgumentParser(description="Create and save a Core AI AIProgram for YOLOS.") parser.add_argument( "--model", choices=["hustvl/yolos-tiny", "hustvl/yolos-base"], @@ -179,7 +171,8 @@ def main(): "--include-debug-info", action="store_true", help="Embed debug information in the exported .aimodel for debugging a conversion. " - "Default: off, which embeds minimum debug information and makes the exported asset smaller.", + "Default: off, which embeds minimum debug information and makes the exported asset " + "smaller.", ) args = parser.parse_args()