Run ruff check and format on models/ with the same rules as python/

.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'

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()
 

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()
 

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()
 

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,19 +397,15 @@ 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"],
         default="bfloat16",
     )
     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,

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()