Export aoti for preprocess (#5354)

Summary:
Export AOTI for preprocess.
Requires: D62651605

Pull Request resolved: #5354

Test Plan:
```
python -m unittest examples/models/flamingo/preprocess/test_preprocess.py
```

Reviewed By: larryliu0820

Differential Revision: D62662418

Pulled By: lucylq

fbshipit-source-id: a094a14870fcca820fb4b739684b5e75aefd36d3

Author: lucylq

examples/models/flamingo/preprocess/export_preprocess.py

@@ -4,16 +4,30 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from export_preprocess_lib import export_preprocess, lower_to_executorch_preprocess
+import torch
+from executorch.examples.models.flamingo.preprocess.export_preprocess_lib import (
+    export_preprocess,
+    get_example_inputs,
+    lower_to_executorch_preprocess,
+)
 
 
 def main():
+    # Export
     ep = export_preprocess()
-    et = lower_to_executorch_preprocess(ep)
 
-    with open("preprocess.pte", "wb") as file:
+    # ExecuTorch
+    et = lower_to_executorch_preprocess(ep)
+    with open("preprocess_et.pte", "wb") as file:
         et.write_to_file(file)
 
+    # AOTInductor
+    torch._inductor.aot_compile(
+        ep.module(),
+        get_example_inputs(),
+        options={"aot_inductor.output_path": "preprocess_aoti.so"},
+    )
+
 
 if __name__ == "__main__":
     main()

examples/models/flamingo/preprocess/test_preprocess.py

@@ -37,7 +37,11 @@
 )
 from torchvision.transforms.v2 import functional as F
 
-from .export_preprocess_lib import export_preprocess, lower_to_executorch_preprocess
+from .export_preprocess_lib import (
+    export_preprocess,
+    get_example_inputs,
+    lower_to_executorch_preprocess,
+)
 
 
 @dataclass
@@ -206,6 +210,11 @@ def test_preprocess(
         executorch_model = lower_to_executorch_preprocess(exported_model)
         executorch_module = _load_for_executorch_from_buffer(executorch_model.buffer)
 
+        aoti_path = torch._inductor.aot_compile(
+            exported_model.module(),
+            get_example_inputs(),
+        )
+
         # Prepare image input.
         image = (
             np.random.randint(0, 256, np.prod(image_size))
@@ -266,3 +275,9 @@ def test_preprocess(
         )
         self.assertTrue(torch.allclose(reference_image, et_image))
         self.assertEqual(reference_ar, et_ar.tolist())
+
+        # Run aoti model and check it matches reference model.
+        aoti_model = torch._export.aot_load(aoti_path, "cpu")
+        aoti_image, aoti_ar = aoti_model(image_tensor, inscribed_size, best_resolution)
+        self.assertTrue(torch.allclose(reference_image, aoti_image))
+        self.assertEqual(reference_ar, aoti_ar.tolist())