Core AI Model Authoring and Optimization

import CoreML
import Foundation
import SwiftUI

// MARK: - Core AI Model Runner
// The .aimodel artifact is produced offline by the coreai-torch Python pipeline:
//   exported = torch.export.export(model, (example_input,))
//   core_ai_program = TorchConverter().convert(exported, inputs=[...], outputs=[...])
//   asset = core_ai_program.specialize(SpecializationOptions())
//   asset.save("SAM3Compressed.aimodel")
// Then add SAM3Compressed.mlpackage to your Xcode project.

struct CoreAIInferenceView: View {
    @State private var result: String = "Tap Run to execute model"
    @State private var isRunning = false

    var body: some View {
        VStack(spacing: 20) {
            Text("Core AI Model Demo")
                .font(.title2.bold())

            Text(result)
                .multilineTextAlignment(.center)
                .padding()
                .background(.secondarySystemBackground)
                .clipShape(RoundedRectangle(cornerRadius: 12))

            Button(isRunning ? "Running…" : "Run Inference") {
                Task { await runModel() }
            }
            .buttonStyle(.borderedProminent)
            .disabled(isRunning)
        }
        .padding()
    }

    @MainActor
    func runModel() async {
        isRunning = true
        defer { isRunning = false }

        do {
            // Load the compiled Core AI / Core ML model asset
            guard let modelURL = Bundle.main.url(
                forResource: "SAM3Compressed",
                withExtension: "mlmodelc"
            ) else {
                result = "Model asset not found in bundle"
                return
            }

            // Configure compute units — Neural Engine preferred for Apple Silicon
            let config = MLModelConfiguration()
            config.computeUnits = .cpuAndNeuralEngine

            let model = try await MLModel.load(contentsOf: modelURL, configuration: config)

            // Build feature provider with input tensors
            // Inputs match the names registered during coreai-torch conversion
            let imageData = MLMultiArray(
                shape: [1, 3, 1024, 1024],
                dataType: .float16
            )
            // (populate imageData with pixel values in a real app)

            let inputProvider = try MLDictionaryFeatureProvider(dictionary: [
                "pixel_values": MLFeatureValue(multiArray: imageData)
            ])

            // Run inference asynchronously
            let output = try await model.prediction(from: inputProvider)

            // Read output feature by name registered at conversion time
            if let maskArray = output.featureValue(for: "pred_masks")?.multiArrayValue {
                result = "Mask shape: \(maskArray.shape) — inference succeeded ✓"
            } else {
                result = "Output feature not found"
            }
        } catch {
            result = "Error: \(error.localizedDescription)"
        }
    }
}

#Preview {
    CoreAIInferenceView()
}