sample_model.py

# SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
# SPDX-License-Identifier: Apache-2.0

from pathlib import Path

import torch
from executorch.backends.arm.quantizer.arm_quantizer import QuantizationSpec
from torchao.quantization.pt2e.observer import FixedQParamsObserver

from export_executorch import export_model
from export_scenario import build_scenario_from_edge_program
from utils import load_tensor, reset_generated_artifact_dirs

BUNDLE_ROOT = Path(__file__).resolve().parent
ARTIFACTS_ROOT = BUNDLE_ROOT / "artifacts"
SCENARIO_DIR = ARTIFACTS_ROOT / "scenario"
SAMPLE_INPUT_IMAGE = BUNDLE_ROOT / "sample_inputs" / "input.png"
SAMPLE_SAMPLING_GRID = BUNDLE_ROOT / "sample_inputs" / "grid.npy"


class SampleModel(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        blur_kernel = (
            torch.tensor(
                [
                    [1.0, 2.0, 1.0],
                    [2.0, 4.0, 2.0],
                    [1.0, 2.0, 1.0],
                ],
                dtype=torch.float32,
            )
            / 16.0
        )
        sharpen_kernel = torch.tensor(
            [
                [0.0, -1.0, 0.0],
                [-1.0, 5.0, -1.0],
                [0.0, -1.0, 0.0],
            ],
            dtype=torch.float32,
        )
        self.register_buffer(
            "blur_weight",
            blur_kernel.view(1, 1, 3, 3).repeat(4, 1, 1, 1),
        )
        self.register_buffer(
            "sharpen_weight",
            sharpen_kernel.view(1, 1, 3, 3).repeat(4, 1, 1, 1),
        )

    def forward(
        self,
        input_image: torch.Tensor,
        sampling_grid: torch.Tensor,
    ) -> torch.Tensor:
        blurred = torch.nn.functional.conv2d(
            input_image,
            self.blur_weight,
            padding=1,
            groups=input_image.shape[1],
        )
        rotated = torch.nn.functional.grid_sample(
            blurred,
            sampling_grid,
            mode="bilinear",
            padding_mode="zeros",
            align_corners=False,
        )
        return torch.nn.functional.conv2d(
            rotated,
            self.sharpen_weight,
            padding=1,
            groups=rotated.shape[1],
        )


def int8_symmetric_qspec() -> QuantizationSpec:
    return QuantizationSpec(
        dtype=torch.int8,
        observer_or_fake_quant_ctr=FixedQParamsObserver.with_args(
            scale=1.0 / 127.0,
            zero_point=0,
            dtype=torch.qint8,
            qscheme=torch.per_tensor_symmetric,
            quant_min=-127,
            quant_max=127,
        ),
        quant_min=-127,
        quant_max=127,
        qscheme=torch.per_tensor_symmetric,
        is_dynamic=False,
    )


def int16_symmetric_qspec() -> QuantizationSpec:
    return QuantizationSpec(
        dtype=torch.int16,
        observer_or_fake_quant_ctr=FixedQParamsObserver.with_args(
            scale=1.0 / 32767.0,
            zero_point=0,
            dtype=torch.int16,
            qscheme=torch.per_tensor_symmetric,
            quant_min=-32767,
            quant_max=32767,
        ),
        quant_min=-32767,
        quant_max=32767,
        qscheme=torch.per_tensor_symmetric,
        is_dynamic=False,
    )


def main() -> None:
    model = SampleModel().eval()
    ARTIFACTS_ROOT.mkdir(parents=True, exist_ok=True)
    reset_generated_artifact_dirs(ARTIFACTS_ROOT, SCENARIO_DIR)

    input_image = load_tensor(SAMPLE_INPUT_IMAGE)
    sampling_grid = load_tensor(SAMPLE_SAMPLING_GRID)
    model_inputs = (input_image, sampling_grid)

    edge_program, example_inputs, io_quant_params = export_model(
        model,
        model_inputs,
        artifacts_root=SCENARIO_DIR,
        intermediate_artifacts_root=ARTIFACTS_ROOT,
        quantized_input_qspecs=[int8_symmetric_qspec(), int16_symmetric_qspec()],
        quantized_output_qspecs=[int8_symmetric_qspec()],
        enable_pt2e_quantization=True,
    )
    scenario_path = build_scenario_from_edge_program(
        edge_program,
        SCENARIO_DIR / "scenario.json",
        example_inputs,
        io_quant_params=io_quant_params,
    )
    print()
    print()
    print(f"Scenario: {scenario_path}")
    print("Run the scenario when ready with:")
    print(
        f'scenario-runner --scenario "{scenario_path}" '
        f'--output "{SCENARIO_DIR}" --log-level debug'
    )


if __name__ == "__main__":
    main()