InputButton.swift 4.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. import SwiftUI
  2. struct InputButton: View {
  3. @ObservedObject var llamaState: LlamaState
  4. @State private var inputLink: String = ""
  5. @State private var status: String = "download"
  6. @State private var filename: String = ""
  7. @State private var downloadTask: URLSessionDownloadTask?
  8. @State private var progress = 0.0
  9. @State private var observation: NSKeyValueObservation?
  10. private static func extractModelInfo(from link: String) -> (modelName: String, filename: String)? {
  11. guard let url = URL(string: link),
  12. let lastPathComponent = url.lastPathComponent.components(separatedBy: ".").first,
  13. let modelName = lastPathComponent.components(separatedBy: "-").dropLast().joined(separator: "-").removingPercentEncoding,
  14. let filename = lastPathComponent.removingPercentEncoding else {
  15. return nil
  16. }
  17. return (modelName, filename)
  18. }
  19. private static func getFileURL(filename: String) -> URL {
  20. FileManager.default.urls(for: .documentDirectory, in: .userDomainMask)[0].appendingPathComponent(filename)
  21. }
  22. private func download() {
  23. guard let extractedInfo = InputButton.extractModelInfo(from: inputLink) else {
  24. // Handle invalid link or extraction failure
  25. return
  26. }
  27. let (modelName, filename) = extractedInfo
  28. self.filename = filename // Set the state variable
  29. status = "downloading"
  30. print("Downloading model \(modelName) from \(inputLink)")
  31. guard let url = URL(string: inputLink) else { return }
  32. let fileURL = InputButton.getFileURL(filename: filename)
  33. downloadTask = URLSession.shared.downloadTask(with: url) { temporaryURL, response, error in
  34. if let error = error {
  35. print("Error: \(error.localizedDescription)")
  36. return
  37. }
  38. guard let response = response as? HTTPURLResponse, (200...299).contains(response.statusCode) else {
  39. print("Server error!")
  40. return
  41. }
  42. do {
  43. if let temporaryURL = temporaryURL {
  44. try FileManager.default.copyItem(at: temporaryURL, to: fileURL)
  45. print("Writing to \(filename) completed")
  46. llamaState.cacheCleared = false
  47. let model = Model(name: modelName, url: self.inputLink, filename: filename, status: "downloaded")
  48. llamaState.downloadedModels.append(model)
  49. status = "downloaded"
  50. }
  51. } catch let err {
  52. print("Error: \(err.localizedDescription)")
  53. }
  54. }
  55. observation = downloadTask?.progress.observe(\.fractionCompleted) { progress, _ in
  56. self.progress = progress.fractionCompleted
  57. }
  58. downloadTask?.resume()
  59. }
  60. var body: some View {
  61. VStack {
  62. HStack {
  63. TextField("Paste Quantized Download Link", text: $inputLink)
  64. .textFieldStyle(RoundedBorderTextFieldStyle())
  65. Button(action: {
  66. downloadTask?.cancel()
  67. status = "download"
  68. }) {
  69. Text("Cancel")
  70. }
  71. }
  72. if status == "download" {
  73. Button(action: download) {
  74. Text("Download Custom Model")
  75. }
  76. } else if status == "downloading" {
  77. Button(action: {
  78. downloadTask?.cancel()
  79. status = "download"
  80. }) {
  81. Text("Downloading \(Int(progress * 100))%")
  82. }
  83. } else if status == "downloaded" {
  84. Button(action: {
  85. let fileURL = InputButton.getFileURL(filename: self.filename)
  86. if !FileManager.default.fileExists(atPath: fileURL.path) {
  87. download()
  88. return
  89. }
  90. do {
  91. try llamaState.loadModel(modelUrl: fileURL)
  92. } catch let err {
  93. print("Error: \(err.localizedDescription)")
  94. }
  95. }) {
  96. Text("Load Custom Model")
  97. }
  98. } else {
  99. Text("Unknown status")
  100. }
  101. }
  102. .onDisappear() {
  103. downloadTask?.cancel()
  104. }
  105. .onChange(of: llamaState.cacheCleared) { newValue in
  106. if newValue {
  107. downloadTask?.cancel()
  108. let fileURL = InputButton.getFileURL(filename: self.filename)
  109. status = FileManager.default.fileExists(atPath: fileURL.path) ? "downloaded" : "download"
  110. }
  111. }
  112. }
  113. }