DrawViewController.swift 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345
  1. //
  2. // DrawViewController.swift
  3. // AIPaint
  4. //
  5. // Created by Fengyu He on 2022/11/29.
  6. //
  7. import UIKit
  8. import Photos
  9. import SnapKit
  10. import SwiftyJSON
  11. import SwiftUI
  12. import CoreData
  13. class DrawViewController: UIViewController {
  14. lazy var generateView: UIView = {
  15. let view = UIView()
  16. view.backgroundColor = .clear
  17. view.layer.borderWidth = 1
  18. view.layer.borderColor = UIColor.black.cgColor
  19. return view
  20. }()
  21. lazy var promptTextField: UITextField = {
  22. let textField = UITextField()
  23. textField.placeholder = "请输入你的作品的组成元素,以空格分割。"
  24. textField.layer.borderWidth = 1
  25. textField.layer.borderColor = UIColor.black.cgColor
  26. return textField
  27. }()
  28. lazy var negativePromptTextField: UITextField = {
  29. let textField = UITextField()
  30. textField.placeholder = "请输入你不想在作品中看到的元素,以空格分割。"
  31. textField.layer.borderWidth = 1
  32. textField.layer.borderColor = UIColor.black.cgColor
  33. return textField
  34. }()
  35. lazy var generateButton: UIButton = {
  36. let button = UIButton()
  37. button.layer.backgroundColor = UIColor.blue.cgColor
  38. button.addTarget(self, action: #selector(generateAction), for: .touchUpInside)
  39. button.setTitle("生成图片", for: .normal)
  40. button.isUserInteractionEnabled = true
  41. return button
  42. }()
  43. var longPress: UILongPressGestureRecognizer?
  44. lazy var generateImage: UIImageView = {
  45. let imageView = UIImageView()
  46. imageView.layer.borderWidth = 1
  47. imageView.layer.borderColor = UIColor.black.cgColor
  48. imageView.isUserInteractionEnabled = true
  49. return imageView
  50. }()
  51. lazy var remindText: UILabel = {
  52. var label = UILabel()
  53. label.text = "长按图片进行保存"
  54. label.textColor = .systemGray5
  55. label.textAlignment = .center
  56. return label
  57. }()
  58. var prompt = ""
  59. var tranPrompt = ""
  60. var negativePrompt = ""
  61. var tranNegativePrompt = ""
  62. var sessionHash = ""
  63. var delegate: AppDelegate?
  64. var context: NSManagedObjectContext?
  65. var painting: Painting?
  66. let alertController = UIAlertController(title: "保存成功!", message: nil, preferredStyle: .alert)
  67. override func viewDidLoad() {
  68. super.viewDidLoad()
  69. if UITraitCollection.current.userInterfaceStyle == .dark {
  70. view.backgroundColor = .black
  71. } else {
  72. view.backgroundColor = .white
  73. }
  74. self.title = "AI Paint"
  75. delegate = UIApplication.shared.delegate as? AppDelegate
  76. view.addSubview(generateView)
  77. generateView.addSubview(promptTextField)
  78. generateView.addSubview(negativePromptTextField)
  79. generateView.addSubview(generateButton)
  80. view.addSubview(remindText)
  81. view.addSubview(generateImage)
  82. longPress = UILongPressGestureRecognizer(target: self, action: #selector(saveImageToPhotos))
  83. generateImage.addGestureRecognizer(longPress!)
  84. generateView.snp.makeConstraints { (make) in
  85. make.top.equalTo(view.safeAreaLayoutGuide.snp.top)
  86. make.width.equalTo(view.safeAreaLayoutGuide.snp.width)
  87. make.height.equalTo(view.safeAreaLayoutGuide.snp.height).multipliedBy(0.35)
  88. make.centerX.equalTo(view.safeAreaLayoutGuide.snp.centerX)
  89. }
  90. promptTextField.snp.makeConstraints { (make) in
  91. make.top.equalTo(generateView.snp.top).offset(5)
  92. make.width.equalTo(generateView.snp.width).multipliedBy(0.95)
  93. make.height.equalTo(generateView.snp.height).multipliedBy(0.35)
  94. make.centerX.equalTo(generateView.snp.centerX)
  95. }
  96. negativePromptTextField.snp.makeConstraints { (make) in
  97. make.top.equalTo(promptTextField.snp.bottom).offset(5)
  98. make.width.equalTo(generateView.snp.width).multipliedBy(0.95)
  99. make.height.equalTo(generateView.snp.height).multipliedBy(0.35)
  100. make.centerX.equalTo(generateView.snp.centerX)
  101. }
  102. generateButton.snp.makeConstraints { (make) in
  103. make.top.equalTo(negativePromptTextField.snp.bottom).offset(5)
  104. make.width.equalTo(generateView.snp.width).multipliedBy(0.95)
  105. make.height.equalTo(generateView.snp.height).multipliedBy(0.2)
  106. make.centerX.equalTo(generateView.snp.centerX)
  107. }
  108. generateImage.snp.makeConstraints { (make) in
  109. make.top.equalTo(generateView.snp.bottom).offset(10)
  110. make.width.equalTo(view.safeAreaLayoutGuide.snp.width)
  111. make.height.equalTo(view.safeAreaLayoutGuide.snp.height).multipliedBy(0.60)
  112. make.centerX.equalTo(view.safeAreaLayoutGuide.snp.centerX)
  113. }
  114. remindText.snp.makeConstraints { (make) in
  115. make.top.equalTo(generateImage.snp.top).offset(2)
  116. make.width.equalTo(generateImage.snp.width)
  117. make.height.equalTo(generateImage.snp.height).multipliedBy(0.40)
  118. make.centerX.equalTo(generateImage.snp.centerX)
  119. }
  120. }
  121. func translatePrompt(prompt: String, negativePrompt: String, completion: @escaping (_ tranPrompt: String, _ tranNegativePrompt: String) -> Void) {
  122. let APP_KEY = "3bb96f14d2af67ec"
  123. let APP_SECRET = "SRQF1jGwAMuYciFK33KfyRGKa7B1x38u"
  124. let q = prompt + "|" + negativePrompt
  125. let salt = UUID().uuidString
  126. let curtime = Int(Date().timeIntervalSince1970)
  127. let inputQ = q.count <= 20 ? q : String(q.prefix(10)) + String(q.count) + String(q.suffix(10))
  128. let signStr = APP_KEY + inputQ + salt + String(curtime) + APP_SECRET
  129. let cryptoSign = signStr.sha256
  130. let paramDic = ["q": q, "from": "auto", "to": "en", "appKey": APP_KEY, "salt": salt, "sign": cryptoSign, "signType": "v3", "curtime": String(curtime)]
  131. let list = NSMutableArray()
  132. for subDic in paramDic {
  133. let tmpStr = "\(subDic.0)=\(subDic.1)"
  134. list.add(tmpStr)
  135. }
  136. let paramStr = list.componentsJoined(by: "&")
  137. let paramData = "https://openapi.youdao.com/api?" + paramStr
  138. print(paramData)
  139. guard let addr = paramData.addingPercentEncoding(withAllowedCharacters: .urlQueryAllowed) else {
  140. print("error")
  141. return
  142. }
  143. guard let url = URL(string: addr) else {
  144. print("error")
  145. return
  146. }
  147. var request = URLRequest(url: url)
  148. request.httpMethod = "GET"
  149. var promptValue = ""
  150. var negativePromotValue = ""
  151. let task = URLSession.shared.dataTask(with: request) { (data, response, error) in
  152. do {
  153. guard let unwrapedData = data else {
  154. print("error")
  155. return
  156. }
  157. let jsonData = try JSON(data: unwrapedData)
  158. let promptData = jsonData["translation"][0].stringValue
  159. print(promptData)
  160. if String(promptData.first ?? String.Element("")) == "|" {
  161. if #available(iOS 16.0, *) {
  162. let subStringArr = promptData.split(separator: "|")
  163. negativePromotValue = String(subStringArr[0])
  164. } else {
  165. // Fallback on earlier versions
  166. let subStringArr = promptData.components(separatedBy: "|")
  167. negativePromotValue = String(subStringArr[0])
  168. }
  169. } else if String(promptData.last ?? String.Element("")) == "|" {
  170. if #available(iOS 16.0, *) {
  171. let subStringArr = promptData.split(separator: " |")
  172. promptValue = String(subStringArr[0])
  173. } else {
  174. // Fallback on earlier versions
  175. let subStringArr = promptData.components(separatedBy: "|")
  176. promptValue = String(subStringArr[0])
  177. }
  178. } else if promptData == "|"{
  179. // do nothing
  180. } else {
  181. if #available(iOS 16.0, *) {
  182. let subStringArr = promptData.split(separator: "|")
  183. promptValue = String(subStringArr[0])
  184. negativePromotValue = String(subStringArr[1])
  185. } else {
  186. // Fallback on earlier versions
  187. let subStringArr = promptData.components(separatedBy: "|")
  188. promptValue = String(subStringArr[0])
  189. negativePromotValue = String(subStringArr[1])
  190. }
  191. }
  192. } catch {}
  193. completion(promptValue, negativePromotValue)
  194. }
  195. task.resume()
  196. }
  197. @objc func generateAction() {
  198. let promptText = promptTextField.text ?? ""
  199. let negativePromptText = negativePromptTextField.text ?? ""
  200. translatePrompt(prompt: promptTextField.text!, negativePrompt: negativePromptTextField.text!) { [self] promptValue, negativePromotValue in
  201. prompt = "masterpiece, best quality, " + promptValue
  202. negativePrompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digits, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, bad feet, " + negativePromotValue
  203. sessionHash = UUID().uuidString
  204. guard let unwrapedDelegate = delegate else {
  205. return
  206. }
  207. context = unwrapedDelegate.persistentContainer.viewContext
  208. painting = NSEntityDescription.insertNewObject(forEntityName: "Painting", into: context!) as? Painting
  209. painting?.prompt = promptText
  210. painting?.negative_prompt = negativePromptText
  211. painting?.session_hash = sessionHash
  212. painting?.create_time = Date()
  213. guard let url = URL(string: "https://aipaint.hefengyu.org/api/predict/") else {
  214. print("error")
  215. return
  216. }
  217. var request = URLRequest(url: url)
  218. request.httpMethod = "POST"
  219. let postData: [String: Any] = ["fn_index":50, "data":[prompt,negativePrompt,"None","None",20,"Euler a",false,false,1,1,7,-1,-1,0,0,0,false,512,512,false,0.7,0,0,"None",false,false,false,"","Seed","","Nothing","",true,false,false,"null","",""], "session_hash":sessionHash]
  220. let jsonData = try? JSONSerialization.data(withJSONObject: postData, options: [])
  221. request.setValue("application/json", forHTTPHeaderField: "Content-Type")
  222. request.httpBody = jsonData
  223. let dataTask = URLSession.shared.dataTask(with: request) { (data, request, error) in
  224. do {
  225. guard let data = data else {
  226. return
  227. }
  228. guard let json = try? JSON(data: data) else {
  229. return
  230. }
  231. let imageAddr = json["data"][0][0]["name"].stringValue
  232. let addrProcess = imageAddr.replacingOccurrences(of: "\\", with: "/")
  233. let imageStr = "https://aipaint.hefengyu.org/file=\(addrProcess)"
  234. guard let imageUrl = URL(string: imageStr) else {
  235. return
  236. }
  237. guard let imageData = try? Data(contentsOf: imageUrl) else {
  238. return
  239. }
  240. DispatchQueue.main.async { [self] in
  241. painting?.image = imageData
  242. self.generateImage.image = UIImage(data: imageData)
  243. self.generateButton.backgroundColor = .blue
  244. self.generateButton.isUserInteractionEnabled = true
  245. saveContext()
  246. }
  247. }
  248. }
  249. dataTask.resume()
  250. }
  251. generateButton.backgroundColor = .gray
  252. generateButton.isUserInteractionEnabled = false
  253. }
  254. @objc func saveImageToPhotos() {
  255. if (longPress!.state == .began) {
  256. guard let img = generateImage.image else {
  257. print("error")
  258. return
  259. }
  260. switch PHPhotoLibrary.authorizationStatus() {
  261. case .authorized:
  262. saveImg(image: img)
  263. case .notDetermined:
  264. PHPhotoLibrary.requestAuthorization { (status) in
  265. if status == .authorized {
  266. self.saveImg(image: img)
  267. } else {
  268. print("用户拒绝访问")
  269. }
  270. }
  271. case .restricted, .denied:
  272. if let url = URL.init(string: UIApplication.openSettingsURLString) {
  273. if UIApplication.shared.canOpenURL(url) {
  274. UIApplication.shared.open(url)
  275. }
  276. }
  277. case .limited:
  278. return
  279. @unknown default:
  280. return
  281. }
  282. }else if (longPress!.state == .ended){
  283. return
  284. }
  285. }
  286. func saveContext() {
  287. if context!.hasChanges {
  288. do {
  289. try context?.save()
  290. print("save success.")
  291. } catch {}
  292. }
  293. }
  294. func saveImg(image: UIImage) {
  295. PHPhotoLibrary.shared().performChanges({
  296. PHAssetChangeRequest.creationRequestForAsset(from: image)
  297. }, completionHandler: { (isSuccess, error) in
  298. DispatchQueue.main.async {
  299. if isSuccess{
  300. self.present(self.alertController, animated: true, completion: nil)
  301. DispatchQueue.main.asyncAfter(deadline: DispatchTime.now() + 2) {
  302. self.alertController.dismiss(animated: false)
  303. }
  304. }
  305. }
  306. })
  307. }
  308. override func touchesBegan(_ touches: Set<UITouch>, with event: UIEvent?) {
  309. promptTextField.resignFirstResponder()
  310. negativePromptTextField.resignFirstResponder()
  311. }
  312. }