unet_script.py

import torch
import numpy as np
from shark.shark_inference import SharkInference
from shark.shark_importer import SharkImporter
from shark.shark_downloader import download_torch_model

mlir_model, func_name, inputs, golden_out = download_torch_model(
    "stable_diff_quant", tank_url="gs://shark_tank/prashant_nod"
)

shark_module = SharkInference(mlir_model, func_name, mlir_dialect="linalg")
shark_module.compile()
result = shark_module.forward(inputs)
np.testing.assert_allclose(golden_out, result, rtol=1e-02, atol=1e-03)