Misc fixes for the demo app. (#6064)

Misc fixes for the demo app. (#6051) Summary: Pull Request resolved: #6051 . Reviewed By: kirklandsign Differential Revision: D64119638 fbshipit-source-id: 486e499d899625b397ddb32d6b0400e740804bae (cherry picked from commit 9364c06) Co-authored-by: Anthony Shoumikhin <[email protected]> Co-authored-by: Anthony Shoumikhin <[email protected]>
pytorch · Oct 11, 2024 · f17c9e1 · f17c9e1
1 parent 6e788c7
commit f17c9e1
Show file tree

Hide file tree

Showing 9 changed files with 50 additions and 583 deletions.
diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/project.pbxproj b/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/project.pbxproj
diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMA.xcscheme b/examples/demo-apps/apple_ios/LLaMA/LLaMA.xcodeproj/xcshareddata/xcschemes/LLaMA.xcscheme
@@ -23,14 +23,14 @@
       </BuildActionEntries>
    </BuildAction>
    <TestAction
-      buildConfiguration = "Debug"
+      buildConfiguration = "Release"
       selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
       selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
       shouldUseLaunchSchemeArgsEnv = "YES"
       shouldAutocreateTestPlan = "YES">
    </TestAction>
    <LaunchAction
-      buildConfiguration = "Debug"
+      buildConfiguration = "Release"
       selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
       selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
       launchStyle = "0"

diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift
@@ -23,7 +23,6 @@ extension UIImage {
     return UIGraphicsImageRenderer(size: newSize, format: format).image {
       _ in draw(in: CGRect(origin: .zero, size: newSize))
     }
-    return image
   }
 
   func toRGBArray() -> [UInt8]? {
@@ -177,36 +176,38 @@ struct ContentView: View {
         .padding([.leading, .trailing, .bottom], 10)
         .sheet(isPresented: $isImagePickerPresented, onDismiss: addSelectedImageMessage) {
           ImagePicker(selectedImage: $selectedImage, sourceType: imagePickerSourceType)
+            .id(imagePickerSourceType.rawValue)
         }
       }
       .navigationBarTitle(title, displayMode: .inline)
-      .navigationBarItems(leading:
-                            Button(action: {
-                              showingSettings.toggle()
-                            }) {
-                              Image(systemName: "gearshape")
-                                .imageScale(.large)
-                            })
-      .navigationBarItems(trailing:
-                            HStack {
-                              Menu {
-                                Section(header: Text("Memory")) {
-                                  Text("Used: \(resourceMonitor.usedMemory) Mb")
-                                  Text("Available: \(resourceMonitor.availableMemory) Mb")
-                                }
-                              } label: {
-                                Text("\(resourceMonitor.usedMemory) Mb")
-                              }
-                              .onAppear {
-                                resourceMonitor.start()
-                              }
-                              .onDisappear {
-                                resourceMonitor.stop()
-                              }
-                              Button(action: { showingLogs = true }) {
-                                Image(systemName: "list.bullet.rectangle")
-                              }
-                            }
+      .navigationBarItems(
+        leading:
+          Button(action: {
+            showingSettings.toggle()
+          }) {
+            Image(systemName: "gearshape")
+              .imageScale(.large)
+          },
+        trailing:
+          HStack {
+            Menu {
+              Section(header: Text("Memory")) {
+                Text("Used: \(resourceMonitor.usedMemory) Mb")
+                Text("Available: \(resourceMonitor.usedMemory) Mb")
+              }
+            } label: {
+              Text("\(resourceMonitor.usedMemory) Mb")
+            }
+            .onAppear {
+              resourceMonitor.start()
+            }
+            .onDisappear {
+              resourceMonitor.stop()
+            }
+            Button(action: { showingLogs = true }) {
+              Image(systemName: "list.bullet.rectangle")
+            }
+          }
       )
       .sheet(isPresented: $showingLogs) {
         NavigationView {
@@ -247,20 +248,24 @@ struct ContentView: View {
     isGenerating = true
     shouldStopGenerating = false
     shouldStopShowingToken = false
-    let text = prompt
+    let text = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
     let seq_len = 768 // text: 256, vision: 768
     let modelPath = resourceManager.modelPath
     let tokenizerPath = resourceManager.tokenizerPath
-    let useLlama = modelPath.range(of: "llama", options: .caseInsensitive) != nil
+    let useLlama = modelPath.lowercased().contains("llama")
 
     prompt = ""
     hideKeyboard()
     showingSettings = false
 
+    messages.append(Message(text: text))
+    messages.append(Message(type: useLlama ? .llamagenerated : .llavagenerated))
+
     runnerQueue.async {
       defer {
         DispatchQueue.main.async {
           isGenerating = false
+          selectedImage = nil
         }
       }
 
@@ -272,10 +277,7 @@ struct ContentView: View {
 
       guard !shouldStopGenerating else { return }
       if useLlama {
-        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 {
@@ -305,10 +307,7 @@ struct ContentView: View {
           }
         }
       } else {
-        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 {
@@ -367,7 +366,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 {
@@ -383,7 +382,7 @@ struct ContentView: View {
                   }
                 }
                 if shouldStopGenerating {
-                  runnerHolder.runner?.stop()
+                  runnerHolder.llavaRunner?.stop()
                 }
               }
             }
@@ -394,7 +393,7 @@ struct ContentView: View {
           try runnerHolder.runner?.generate(llama3_prompt, sequenceLength: seq_len) { token in
 
             NSLog(">>> token={\(token)}")
-            if token != llama3_prompt && !shouldStopShowingToken {
+            if token != llama3_prompt {
               // hack to fix the issue that extension/llm/runner/text_token_generator.h
               // keeps generating after <|eot_id|>
               if token == "<|eot_id|>" {
@@ -458,6 +457,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/LLaMAPerfBenchmark/LLaMAPerfBenchmark.entitlements b/examples/demo-apps/apple_ios/LLaMA/LLaMAPerfBenchmark/LLaMAPerfBenchmark.entitlements
diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMAPerfBenchmark/LLaMAPerfBenchmarkApp.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMAPerfBenchmark/LLaMAPerfBenchmarkApp.swift
diff --git a/examples/demo-apps/apple_ios/LLaMA/LLaMAPerfBenchmarkTests/LLaMAPerfBenchmarkTests.swift b/examples/demo-apps/apple_ios/LLaMA/LLaMAPerfBenchmarkTests/LLaMAPerfBenchmarkTests.swift
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();
 }
 

diff --git a/extension/llm/runner/text_token_generator.h b/extension/llm/runner/text_token_generator.h
@@ -18,7 +18,7 @@ namespace executorch {
 namespace extension {
 namespace llm {
 
-class TextTokenGenerator {
+class ET_EXPERIMENTAL TextTokenGenerator {
  public:
   TextTokenGenerator(
       Tokenizer* tokenizer,
@@ -72,10 +72,11 @@ class TextTokenGenerator {
     // initialize tensor wrappers
     auto tokens_managed = from_blob(
         token_data.data(), token_shape, executorch::aten::ScalarType::Long);
-
     auto start_pos_managed =
         from_blob(&pos, {1}, executorch::aten::ScalarType::Long);
 
+    should_stop_ = false;
+
     // Generate our tokens
     while (pos < seq_len - 1) {
       // Run the model