From 83a7e5ad180c67c63d5aa2a0a35d4fa39209048a Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <shoumikhin@meta.com>
Date: Wed, 9 Oct 2024 09:58:56 -0700
Subject: [PATCH] Misc fixes for the demo app.

Summary: .

Differential Revision: D64119638
---
 .../LLaMA/LLaMA/Application/ContentView.swift | 21 +++++++++----------
 .../LLaMARunner/Exported/LLaMARunner.h        |  4 ++--
 .../LLaMARunner/Exported/LLaMARunner.mm       |  4 ++--
 3 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift
index f47771b8f26..8f7c752f1c5 100644
--- a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift
+++ b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift
@@ -18,12 +18,9 @@ class RunnerHolder: ObservableObject {
 
 extension UIImage {
   func resized(to newSize: CGSize) -> UIImage {
-    let format = UIGraphicsImageRendererFormat.default()
-    let renderer = UIGraphicsImageRenderer(size: newSize, format: format)
-    let image = renderer.image { _ in
+    UIGraphicsImageRenderer(size: newSize).image { _ in
       draw(in: CGRect(origin: .zero, size: newSize))
     }
-    return image
   }
 
   func toRGBArray() -> [UInt8]? {
@@ -71,10 +68,10 @@ extension UIImage {
         let g = UInt8(pixelData[pixelIndex + 1])
         let b = UInt8(pixelData[pixelIndex + 2])
 
-        let rgbIndex = (y * width + x)
+        let rgbIndex = (y * width + x) * 3
         rgbArray[rgbIndex] = r
-        rgbArray[rgbIndex + height * width] = g
-        rgbArray[rgbIndex + 2 * height * width] = b
+        rgbArray[rgbIndex + 1] = g
+        rgbArray[rgbIndex + 2] = b
       }
     }
 
@@ -288,6 +285,7 @@ struct ContentView: View {
       defer {
         DispatchQueue.main.async {
           isGenerating = false
+          selectedImage = nil
         }
       }
 
@@ -302,7 +300,7 @@ struct ContentView: View {
         messages.append(Message(text: text))
         messages.append(Message(type: .llamagenerated))
 
-        if let runner = runnerHolder.runner, !runner.isloaded() {
+        if let runner = runnerHolder.runner, !runner.isLoaded() {
           var error: Error?
           let startLoadTime = Date()
           do {
@@ -335,7 +333,7 @@ struct ContentView: View {
         messages.append(Message(text: text))
         messages.append(Message(type: .llavagenerated))
 
-        if let runner = runnerHolder.llavaRunner, !runner.isloaded() {
+        if let runner = runnerHolder.llavaRunner, !runner.isLoaded() {
           var error: Error?
           let startLoadTime = Date()
           do {
@@ -394,7 +392,7 @@ struct ContentView: View {
             if token != llava_prompt {
               if token == "</s>" {
                 shouldStopGenerating = true
-                runnerHolder.runner?.stop()
+                runnerHolder.llavaRunner?.stop()
               } else {
                 tokens.append(token)
                 if tokens.count > 2 {
@@ -410,7 +408,7 @@ struct ContentView: View {
                   }
                 }
                 if shouldStopGenerating {
-                  runnerHolder.runner?.stop()
+                  runnerHolder.llavaRunner?.stop()
                 }
               }
             }
@@ -485,6 +483,7 @@ struct ContentView: View {
       }
       runnerQueue.async {
         runnerHolder.runner = nil
+        runnerHolder.llavaRunner = nil
       }
       switch pickerType {
       case .model:
diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h
index c4edea2af93..1290a75d975 100644
--- a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h
+++ b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.h
@@ -18,7 +18,7 @@ NS_SWIFT_NAME(Runner)
 
 - (instancetype)initWithModelPath:(NSString*)filePath
                     tokenizerPath:(NSString*)tokenizerPath;
-- (BOOL)isloaded;
+- (BOOL)isLoaded;
 - (BOOL)loadWithError:(NSError**)error;
 - (BOOL)generate:(NSString*)prompt
        sequenceLength:(NSInteger)seq_len
@@ -36,7 +36,7 @@ NS_SWIFT_NAME(LLaVARunner)
 
 - (instancetype)initWithModelPath:(NSString*)filePath
                     tokenizerPath:(NSString*)tokenizerPath;
-- (BOOL)isloaded;
+- (BOOL)isLoaded;
 - (BOOL)loadWithError:(NSError**)error;
 - (BOOL)generate:(void*)imageBuffer
                 width:(CGFloat)width
diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm
index 3136d2745fd..e6a4722ddbf 100644
--- a/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm
+++ b/examples/demo-apps/apple_ios/LLaMA/LLaMARunner/LLaMARunner/Exported/LLaMARunner.mm
@@ -40,7 +40,7 @@ - (void)dealloc {
   [ExecuTorchLog.sharedLog removeSink:self];
 }
 
-- (BOOL)isloaded {
+- (BOOL)isLoaded {
   return _runner->is_loaded();
 }
 
@@ -128,7 +128,7 @@ - (void)dealloc {
   [ExecuTorchLog.sharedLog removeSink:self];
 }
 
-- (BOOL)isloaded {
+- (BOOL)isLoaded {
   return _runner->is_loaded();
 }