Misc fixes for the demo app.

Summary: . Differential Revision: D64119638
pytorch · Oct 9, 2024 · 83a7e5a · 83a7e5a
1 parent 566902b
commit 83a7e5a
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 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
@@ -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
@@ -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
@@ -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();
 }