From 04633a37d9d44d2495c014edb43d08056120a544 Mon Sep 17 00:00:00 2001 From: Andy KoKo <90698189+czkoko@users.noreply.github.com> Date: Sat, 24 Feb 2024 02:05:41 +0800 Subject: [PATCH] The final version - more Scheduler - New painting interface - New any model inpaiting support - Variable Resolution - etc. --- Mochi Diffusion.xcodeproj/project.pbxproj | 42 +-- .../xcshareddata/swiftpm/Package.resolved | 10 +- Mochi Diffusion/Model/SDModel.swift | 9 +- Mochi Diffusion/Model/Scheduler.swift | 25 +- Mochi Diffusion/Resources/de-coremldata.bin | Bin 0 -> 903 bytes Mochi Diffusion/Resources/en-coremldata.bin | Bin 0 -> 908 bytes Mochi Diffusion/Support/Extensions.swift | 1 - Mochi Diffusion/Support/ImageController.swift | 13 +- Mochi Diffusion/Support/ImageGenerator.swift | 279 +++++++++++++++--- Mochi Diffusion/Views/GalleryItemView.swift | 4 +- .../Views/GalleryPreviewView.swift | 3 +- Mochi Diffusion/Views/GalleryView.swift | 2 +- .../SidebarControls/ControlNetView.swift | 2 +- .../Views/SidebarControls/SizeView.swift | 10 +- .../SidebarControls/StartingImageView.swift | 142 ++++++++- Mochi Diffusion/Views/SidebarView.swift | 4 + 16 files changed, 461 insertions(+), 85 deletions(-) create mode 100644 Mochi Diffusion/Resources/de-coremldata.bin create mode 100644 Mochi Diffusion/Resources/en-coremldata.bin diff --git a/Mochi Diffusion.xcodeproj/project.pbxproj b/Mochi Diffusion.xcodeproj/project.pbxproj index 61a6d5da..28f0ffcc 100755 --- a/Mochi Diffusion.xcodeproj/project.pbxproj +++ b/Mochi Diffusion.xcodeproj/project.pbxproj @@ -16,7 +16,6 @@ 03173C132999E2B500B03456 /* SDModel.swift in Sources */ = {isa = PBXBuildFile; fileRef = 03173C122999E2B500B03456 /* SDModel.swift */; }; 03173C152999F5C700B03456 /* ImageController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 03173C142999F5C700B03456 /* ImageController.swift */; }; 03173C18299B3C1300B03456 /* ImageStore.swift in Sources */ = {isa = PBXBuildFile; fileRef = 03173C17299B3C1300B03456 /* ImageStore.swift */; }; - 0352E2A5294E2591003FBF25 /* Sparkle in Frameworks */ = {isa = PBXBuildFile; productRef = 0352E2A4294E2591003FBF25 /* Sparkle */; }; 0352E2A7294E3148003FBF25 /* HelpCommands.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0352E2A6294E3148003FBF25 /* HelpCommands.swift */; }; 0352E2AA294EA0E1003FBF25 /* AppView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0352E2A9294EA0E1003FBF25 /* AppView.swift */; }; 0352E2AE294EA2B4003FBF25 /* MessageBanner.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0352E2AD294EA2B4003FBF25 /* MessageBanner.swift */; }; @@ -61,6 +60,9 @@ C1ADEC2C2B16957800E142CA /* FolderMonitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = C1ADEC2B2B16957800E142CA /* FolderMonitor.swift */; }; D7B03F2029D42F9900DF89DD /* SDModelAttentionType.swift in Sources */ = {isa = PBXBuildFile; fileRef = D7B03F1F29D42F9900DF89DD /* SDModelAttentionType.swift */; }; E45FA5822B7F7E4B009E90F0 /* GuernikaKit in Frameworks */ = {isa = PBXBuildFile; productRef = E45FA5812B7F7E4B009E90F0 /* GuernikaKit */; }; + E4EC46322B890B6D00351E8C /* Sparkle in Frameworks */ = {isa = PBXBuildFile; productRef = E4EC46312B890B6D00351E8C /* Sparkle */; }; + E4F15B1D2B86720F00E1EC3C /* de-coremldata.bin in Resources */ = {isa = PBXBuildFile; fileRef = E4F15B1C2B86720F00E1EC3C /* de-coremldata.bin */; }; + E4F15B1F2B8673DB00E1EC3C /* en-coremldata.bin in Resources */ = {isa = PBXBuildFile; fileRef = E4F15B1E2B8673DB00E1EC3C /* en-coremldata.bin */; }; /* End PBXBuildFile section */ /* Begin PBXFileReference section */ @@ -132,6 +134,8 @@ C1220FBF2AFB122F007E5055 /* ImageWellView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ImageWellView.swift; sourceTree = ""; }; C1ADEC2B2B16957800E142CA /* FolderMonitor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FolderMonitor.swift; sourceTree = ""; }; D7B03F1F29D42F9900DF89DD /* SDModelAttentionType.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SDModelAttentionType.swift; sourceTree = ""; }; + E4F15B1C2B86720F00E1EC3C /* de-coremldata.bin */ = {isa = PBXFileReference; lastKnownFileType = archive.macbinary; path = "de-coremldata.bin"; sourceTree = ""; }; + E4F15B1E2B8673DB00E1EC3C /* en-coremldata.bin */ = {isa = PBXFileReference; lastKnownFileType = archive.macbinary; path = "en-coremldata.bin"; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFrameworksBuildPhase section */ @@ -140,8 +144,8 @@ buildActionMask = 2147483647; files = ( 036BFC4E294B9FDB00D8AD04 /* Path in Frameworks */, + E4EC46322B890B6D00351E8C /* Sparkle in Frameworks */, E45FA5822B7F7E4B009E90F0 /* GuernikaKit in Frameworks */, - 0352E2A5294E2591003FBF25 /* Sparkle in Frameworks */, 0388DEB0297B00FC008B1C1C /* CompactSlider in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; @@ -275,6 +279,8 @@ isa = PBXGroup; children = ( 036BFC25294B9F7600D8AD04 /* Assets.xcassets */, + E4F15B1E2B8673DB00E1EC3C /* en-coremldata.bin */, + E4F15B1C2B86720F00E1EC3C /* de-coremldata.bin */, 03FD2318295F74B6006EEEE2 /* RealESRGAN.mlmodel */, ); path = Resources; @@ -298,9 +304,9 @@ name = "Mochi Diffusion"; packageProductDependencies = ( 036BFC4D294B9FDB00D8AD04 /* Path */, - 0352E2A4294E2591003FBF25 /* Sparkle */, 0388DEAF297B00FC008B1C1C /* CompactSlider */, E45FA5812B7F7E4B009E90F0 /* GuernikaKit */, + E4EC46312B890B6D00351E8C /* Sparkle */, ); productName = "Mochi Diffusion"; productReference = 036BFC1E294B9F7500D8AD04 /* Mochi Diffusion.app */; @@ -367,9 +373,9 @@ mainGroup = 036BFC15294B9F7500D8AD04; packageReferences = ( 036BFC4C294B9FDB00D8AD04 /* XCRemoteSwiftPackageReference "Path.swift" */, - 0352E2A3294E2591003FBF25 /* XCRemoteSwiftPackageReference "Sparkle" */, 0388DEAE297B00FC008B1C1C /* XCRemoteSwiftPackageReference "CompactSlider" */, E45FA5802B7F7E4B009E90F0 /* XCRemoteSwiftPackageReference "GuernikaKit" */, + E4EC46302B890B6D00351E8C /* XCRemoteSwiftPackageReference "Sparkle" */, ); productRefGroup = 036BFC1F294B9F7500D8AD04 /* Products */; projectDirPath = ""; @@ -389,6 +395,8 @@ 03F578642967CE9A003A815F /* Localizable.strings in Resources */, 0311C6BD2989E3EF0074BCAE /* Localizable.stringsdict in Resources */, 036BFC26294B9F7600D8AD04 /* Assets.xcassets in Resources */, + E4F15B1F2B8673DB00E1EC3C /* en-coremldata.bin in Resources */, + E4F15B1D2B86720F00E1EC3C /* de-coremldata.bin in Resources */, ); runOnlyForDeploymentPostprocessing = 0; }; @@ -700,14 +708,6 @@ /* End XCConfigurationList section */ /* Begin XCRemoteSwiftPackageReference section */ - 0352E2A3294E2591003FBF25 /* XCRemoteSwiftPackageReference "Sparkle" */ = { - isa = XCRemoteSwiftPackageReference; - repositoryURL = "https://github.com/sparkle-project/Sparkle"; - requirement = { - kind = upToNextMajorVersion; - minimumVersion = 2.0.0; - }; - }; 036BFC4C294B9FDB00D8AD04 /* XCRemoteSwiftPackageReference "Path.swift" */ = { isa = XCRemoteSwiftPackageReference; repositoryURL = "https://github.com/mxcl/Path.swift.git"; @@ -732,14 +732,17 @@ minimumVersion = 1.5.0; }; }; + E4EC46302B890B6D00351E8C /* XCRemoteSwiftPackageReference "Sparkle" */ = { + isa = XCRemoteSwiftPackageReference; + repositoryURL = "https://github.com/sparkle-project/Sparkle.git"; + requirement = { + kind = upToNextMajorVersion; + minimumVersion = 2.5.2; + }; + }; /* End XCRemoteSwiftPackageReference section */ /* Begin XCSwiftPackageProductDependency section */ - 0352E2A4294E2591003FBF25 /* Sparkle */ = { - isa = XCSwiftPackageProductDependency; - package = 0352E2A3294E2591003FBF25 /* XCRemoteSwiftPackageReference "Sparkle" */; - productName = Sparkle; - }; 036BFC4D294B9FDB00D8AD04 /* Path */ = { isa = XCSwiftPackageProductDependency; package = 036BFC4C294B9FDB00D8AD04 /* XCRemoteSwiftPackageReference "Path.swift" */; @@ -755,6 +758,11 @@ package = E45FA5802B7F7E4B009E90F0 /* XCRemoteSwiftPackageReference "GuernikaKit" */; productName = GuernikaKit; }; + E4EC46312B890B6D00351E8C /* Sparkle */ = { + isa = XCSwiftPackageProductDependency; + package = E4EC46302B890B6D00351E8C /* XCRemoteSwiftPackageReference "Sparkle" */; + productName = Sparkle; + }; /* End XCSwiftPackageProductDependency section */ }; rootObject = 036BFC16294B9F7500D8AD04 /* Project object */; diff --git a/Mochi Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved b/Mochi Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved index 48939775..11c853a8 100755 --- a/Mochi Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved +++ b/Mochi Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved @@ -14,8 +14,8 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/GuernikaCore/GuernikaKit.git", "state" : { - "revision" : "2c1ce26535279b309149acba47ae65e6a9394cb5", - "version" : "1.5.0" + "revision" : "dae44bf061ffa84df7f882cb563070d484488e44", + "version" : "1.6.0" } }, { @@ -41,14 +41,14 @@ "kind" : "remoteSourceControl", "location" : "https://github.com/GuernikaCore/Schedulers.git", "state" : { - "revision" : "5e6ab9fa91aa56fc34edc6bae9de1caa7f13deb8", - "version" : "1.2.0" + "revision" : "1f517514d679e38bb9915c3a74bf04f75d5b5875", + "version" : "1.4.0" } }, { "identity" : "sparkle", "kind" : "remoteSourceControl", - "location" : "https://github.com/sparkle-project/Sparkle", + "location" : "https://github.com/sparkle-project/Sparkle.git", "state" : { "revision" : "47d3d90aee3c52b6f61d04ceae426e607df62347", "version" : "2.5.2" diff --git a/Mochi Diffusion/Model/SDModel.swift b/Mochi Diffusion/Model/SDModel.swift index ab535638..28185db2 100755 --- a/Mochi Diffusion/Model/SDModel.swift +++ b/Mochi Diffusion/Model/SDModel.swift @@ -17,7 +17,7 @@ struct SDModel: Identifiable { let attention: SDModelAttentionType let controlNet: [String] let isXL: Bool - let inputSize: CGSize? + var inputSize: CGSize? var id: URL { url } @@ -103,8 +103,7 @@ private func identifyIfXL(_ url: URL) -> Bool { private func unetMetadataURL(from url: URL) -> URL? { let potentialMetadataURLs = [ url.appending(components: "Unet.mlmodelc", "metadata.json"), - url.appending(components: "UnetChunk1.mlmodelc", "metadata.json"), - url.appending(components: "ControlledUnet.mlmodelc", "metadata.json") + url.appending(components: "UnetChunk1.mlmodelc", "metadata.json") ] return potentialMetadataURLs.first { @@ -113,11 +112,11 @@ private func unetMetadataURL(from url: URL) -> URL? { } private func identifyInputSize(_ url: URL) -> CGSize? { - let encoderMetadataURL = url.appending(path: "VAEEncoder.mlmodelc").appending(path: "metadata.json") + let encoderMetadataURL = url.appending(path: "VAEDecoder.mlmodelc").appending(path: "metadata.json") if let jsonData = try? Data(contentsOf: encoderMetadataURL), let jsonArray = try? JSONSerialization.jsonObject(with: jsonData) as? [[String: Any]], let jsonItem = jsonArray.first, - let inputSchema = jsonItem["inputSchema"] as? [[String: Any]], + let inputSchema = jsonItem["outputSchema"] as? [[String: Any]], let controlnetCond = inputSchema.first, let shapeString = controlnetCond["shape"] as? String { let shapeIntArray = shapeString.trimmingCharacters(in: CharacterSet(charactersIn: "[]")) diff --git a/Mochi Diffusion/Model/Scheduler.swift b/Mochi Diffusion/Model/Scheduler.swift index 46bbf920..b432d5be 100755 --- a/Mochi Diffusion/Model/Scheduler.swift +++ b/Mochi Diffusion/Model/Scheduler.swift @@ -9,29 +9,44 @@ import Schedulers /// Schedulers compatible with StableDiffusionPipeline enum Scheduler: String, CaseIterable { - /// Scheduler that uses a pseudo-linear multi-step (PLMS) method - case lcm = "LCM" - /// Scheduler that uses a second order DPM-Solver++ algorithm + case ddim = "DDIM" + case dpmSolverMultistep = "DPM++ 2M" case dpmSolverMultistepKarras = "DPM++ 2M Karras" + case dpmSolverSinglestep = "DPM++ SDE" + case dpmSolverSinglestepKarras = "DPM++ SDE Karras" + + case eulerDiscrete = "Euler" + + case eulerDiscreteKarras = "Euler Karras" case eulerAncenstralDiscrete = "Euler Ancenstral" + + case lcm = "LCM" } func convertScheduler(_ scheduler: Scheduler) -> Schedulers { switch scheduler { - case .lcm: - return Schedulers.lcm + case .ddim: + return Schedulers.ddim case .dpmSolverMultistep: return Schedulers.dpmSolverMultistep case .dpmSolverMultistepKarras: return Schedulers.dpmSolverMultistepKarras + case .dpmSolverSinglestep: + return Schedulers.dpmSolverSinglestep case .dpmSolverSinglestepKarras: return Schedulers.dpmSolverSinglestepKarras + case .eulerDiscrete: + return Schedulers.eulerDiscrete + case .eulerDiscreteKarras: + return Schedulers.eulerDiscreteKarras case .eulerAncenstralDiscrete: return Schedulers.eulerAncenstralDiscrete + case .lcm: + return .lcm } } diff --git a/Mochi Diffusion/Resources/de-coremldata.bin b/Mochi Diffusion/Resources/de-coremldata.bin new file mode 100644 index 0000000000000000000000000000000000000000..c6028cb30f18f63aba9bb88ec8b175dccfeffd5e GIT binary patch literal 903 zcmaJ=(MsGv6y3DhD5I2Zp4#oprS+j|H`!gWLa9(uDhpOcYhOxPW|Ny_Aejj>6A_>G z2MYa!vOmy&>Bsocr`|+YtVri!!ra_@=A3iyTsf^)yY=u|xCy5yQHyAPd^ow)T6ui^ zhkbXLbLoHmoFT+0l`t4X&Mis{SWGaEg-i`3S`}a^r;=!)P%-Af6N|bD=G<5?7E4eG zXgo)42u&nOWrUL5^|e*zIquY%I-MU=_m9(M?xNFiXm>jkWWlF91MoRpue>_goJL2(Q=WSz+c+a$cxjAm#DiR?_<@r ze`*}xgj!Kj#JEH@5J!f+-+muvh!d+FAUxW!w z<67@lT`w0AN`tF2Z-qVMww?)cY%Gd0$AS*LKKhGCC!_t*;m*O#+hDIERRpQ9Sv3iG zNqNLdW06~>NT^Yjj?iDWN=MoD_AuxNgPHe&EmIYljN?iQ>o4sdpN)48_s09jXEX0P zdtxGLyG+NVk}CZLs@mueH)md#J*7s{xX~bv>#3{xwBH7sL4W4;*vfy&+~#iu{mHPt UhzWWDRwM%TKP|bUXLon~1;2L?s{jB1 literal 0 HcmV?d00001 diff --git a/Mochi Diffusion/Resources/en-coremldata.bin b/Mochi Diffusion/Resources/en-coremldata.bin new file mode 100644 index 0000000000000000000000000000000000000000..be5ebd2c4b45085a5fd2ec10084e2a3068aed6d2 GIT binary patch literal 908 zcmaJ=O>Wab6poeDjiwa}UX{RN1YOWV#!0GB38_L=E2^NMoZ5s)vyT{uh^tax@vDOJTRJbwS*_xbtSX*8A^_n+lI?ieI!5iHF6lN*hV<~%~a z66eZ$_b~w~WL{V*G9Cl9icYz~yY>q!cuG@lL0UtVq*CdOr{W62DpPZbC7*kPmvr!@ z+uM9WT&Fo}&)S`Dv(_)?32C*Rj)Tu;r%BxQtks!y7We4~;?5g;Zz>S;cM91Q!cfTA z&`7JS#!*PHG*cm@e|@(wv#z_@>U^KIO7~}r{BX`JpQb?fMHCfAAm~~dqQ>Nidq|+XjQ~XjV8S5n`;$-F^rz1P-&(i!dlcNu{Jla zeV^+~alw?1eLgk*a4;CML4SKRq!CUE!E%ur=B?~&;MrNibI{$lcY9SI;(g71kL!@~ z3oNSnaZ8gTj!~+J2fza5-|F*pe4K+k-5np&=LcjTBU ([SDImage], URL) { var finalImageDirURL: URL @@ -152,25 +153,37 @@ struct GenerationConfig: Sendable, Identifiable { await updateState(.error("Couldn't load \(model.name) because it doesn't exist.")) throw GeneratorError.requestedModelNotFound } - var hasher = Hasher() hasher.combine(model) hasher.combine(controlNet) hasher.combine(computeUnit) hasher.combine(reduceMemory) let hash = hasher.finalize() - guard hash != self.currentPipelineHash else { return } - + let cSize = await CGSize(width: ImageController.shared.width, height: ImageController.shared.height) + + if hash == self.currentPipelineHash { + if await ImageController.shared.startingImage != nil && lastSize == cSize { + return + } else if await ImageController.shared.startingImage == nil{ + return + } + } + + lastSize = await CGSize(width: ImageController.shared.width, height: ImageController.shared.height) await updateState(.loading) let config = MLModelConfiguration() config.computeUnits = computeUnit - + await modifyInputSize(model.url, height: ImageController.shared.height, width: ImageController.shared.width) let modelresource = try GuernikaKit.load(at: model.url) - if model.isXL { + + if FileManager.default.fileExists(atPath: model.url.path() + "WuerstchenPrior.mlmodelc"){ + //pipeline = modelresource as? WuerstchenPipeline + }else if model.isXL { self.pipeline = modelresource as! StableDiffusionXLPipeline - } else { + } else if modelresource.sampleSize.width <= CGFloat(768) { self.pipeline = modelresource as! StableDiffusionMainPipeline } + self.pipeline?.reduceMemory = reduceMemory self.currentPipelineHash = hash self.tokenizer = Tokenizer(modelDir: model.url) @@ -184,6 +197,7 @@ struct GenerationConfig: Sendable, Identifiable { } await updateState(.loading) generationStopped = false + var config = inputConfig config.pipelineConfig.seed = config.pipelineConfig.seed == 0 ? UInt32.random(in: 0 ..< UInt32.max) : config.pipelineConfig.seed @@ -196,10 +210,12 @@ struct GenerationConfig: Sendable, Identifiable { sdi.steps = config.pipelineConfig.stepCount sdi.guidanceScale = Double(config.pipelineConfig.guidanceScale) + try await hackVAE(model: config.model) for index in 0 ..< config.numberOfImages { await updateQueueProgress(QueueProgress(index: index, total: inputConfig.numberOfImages)) generationStartTime = DispatchTime.now() + let images = try pipeline.generateImages(input: config.pipelineConfig) { progress in Task { @MainActor in @@ -219,39 +235,234 @@ struct GenerationConfig: Sendable, Identifiable { if generationStopped { break } - -// for image in images { - guard let image = images else { continue } - if config.upscaleGeneratedImages { - guard let upscaledImg = await Upscaler.shared.upscale(cgImage: image) else { continue } - sdi.image = upscaledImg - sdi.aspectRatio = CGFloat(Double(upscaledImg.width) / Double(upscaledImg.height)) - sdi.upscaler = "RealESRGAN" - } else { - sdi.image = image - sdi.aspectRatio = CGFloat(Double(image.width) / Double(image.height)) - } - sdi.id = UUID() - sdi.seed = config.pipelineConfig.seed - sdi.generatedDate = Date.now - sdi.path = "" - - if config.autosaveImages && !config.imageDir.isEmpty { - var pathURL = URL(fileURLWithPath: config.imageDir, isDirectory: true) - let count = ImageStore.shared.images.endIndex + 1 - pathURL.append(path: sdi.filenameWithoutExtension(count: count)) - - let type = UTType.fromString(config.imageType) - guard let path = await sdi.save(pathURL, type: type) else { continue } - sdi.path = path.path(percentEncoded: false) - } - ImageStore.shared.add(sdi) -// } + + guard let image = images else { continue } + if config.upscaleGeneratedImages { + guard let upscaledImg = await Upscaler.shared.upscale(cgImage: image) else { continue } + sdi.image = upscaledImg + sdi.aspectRatio = CGFloat(Double(upscaledImg.width) / Double(upscaledImg.height)) + sdi.upscaler = "RealESRGAN" + } else { + sdi.image = image + sdi.aspectRatio = CGFloat(Double(image.width) / Double(image.height)) + } + sdi.id = UUID() + sdi.seed = config.pipelineConfig.seed + sdi.generatedDate = Date.now + sdi.path = "" + + if config.autosaveImages && !config.imageDir.isEmpty { + var pathURL = URL(fileURLWithPath: config.imageDir, isDirectory: true) + let count = ImageStore.shared.images.endIndex + 1 + pathURL.append(path: sdi.filenameWithoutExtension(count: count)) + + let type = UTType.fromString(config.imageType) + guard let path = await sdi.save(pathURL, type: type) else { continue } + sdi.path = path.path(percentEncoded: false) + } + ImageStore.shared.add(sdi) config.pipelineConfig.seed += 1 } await updateState(.ready(nil)) } + + func hackVAE(model: SDModel) async throws { + let resourcePath = model.url.path() + if model.isXL{ + if FileManager.default.fileExists(atPath: resourcePath + "VAEDecoder.mlmodelc/model.mil.bak"){ + let vaedecoderMIL = resourcePath + "VAEDecoder.mlmodelc/model.mil" + try FileManager.default.removeItem(atPath: vaedecoderMIL) + try FileManager.default.copyItem(atPath: resourcePath + "VAEDecoder.mlmodelc/model.mil.bak", toPath: vaedecoderMIL) + await vaeDeSDXL(vaeMIL: vaedecoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + }else{ + let vaedecoderMIL = resourcePath + "VAEDecoder.mlmodelc/model.mil" + try FileManager.default.copyItem(atPath: vaedecoderMIL, toPath: resourcePath + "VAEDecoder.mlmodelc/model.mil.bak") + try FileManager.default.removeItem(atPath: resourcePath + "VAEDecoder.mlmodelc/coremldata.bin") + try FileManager.default.copyItem(at: Bundle.main.url(forResource: "de-coremldata", withExtension: "bin")!, to: URL(fileURLWithPath: resourcePath + "VAEDecoder.mlmodelc/coremldata.bin")) + await vaeDeSDXL(vaeMIL: vaedecoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + } + + if FileManager.default.fileExists(atPath: resourcePath + "VAEEncoder.mlmodelc/model.mil.bak"){ + let vaeencoderMIL = resourcePath + "VAEEncoder.mlmodelc/model.mil" + try FileManager.default.removeItem(atPath: vaeencoderMIL) + try FileManager.default.copyItem(atPath: resourcePath + "VAEEncoder.mlmodelc/model.mil.bak", toPath: vaeencoderMIL) + await vaeEnSDXL(vaeMIL: vaeencoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + }else{ + let vaeencoderMIL = resourcePath + "VAEEncoder.mlmodelc/model.mil" + try FileManager.default.copyItem(atPath: vaeencoderMIL, toPath: resourcePath + "VAEEncoder.mlmodelc/model.mil.bak") + try FileManager.default.removeItem(atPath: resourcePath + "VAEEncoder.mlmodelc/coremldata.bin") + try FileManager.default.copyItem(at: Bundle.main.url(forResource: "en-coremldata", withExtension: "bin")!, to: URL(fileURLWithPath: resourcePath + "VAEEncoder.mlmodelc/coremldata.bin")) + await vaeEnSDXL(vaeMIL: vaeencoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + } + }else{ + if FileManager.default.fileExists(atPath: resourcePath + "VAEDecoder.mlmodelc/model.mil.bak"){ + let vaedecoderMIL = resourcePath + "VAEDecoder.mlmodelc/model.mil" + try FileManager.default.removeItem(atPath: vaedecoderMIL) + try FileManager.default.copyItem(atPath: resourcePath + "VAEDecoder.mlmodelc/model.mil.bak", toPath: vaedecoderMIL) + await vaeDeSD(vaeMIL: vaedecoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + }else{ + let vaedecoderMIL = resourcePath + "VAEDecoder.mlmodelc/model.mil" + try FileManager.default.copyItem(atPath: vaedecoderMIL, toPath: resourcePath + "VAEDecoder.mlmodelc/model.mil.bak") + try FileManager.default.removeItem(atPath: resourcePath + "VAEDecoder.mlmodelc/coremldata.bin") + try FileManager.default.copyItem(at: Bundle.main.url(forResource: "de-coremldata", withExtension: "bin")!, to: URL(fileURLWithPath: resourcePath + "VAEDecoder.mlmodelc/coremldata.bin")) + await vaeDeSD(vaeMIL: vaedecoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + } + + if FileManager.default.fileExists(atPath: resourcePath + "VAEEncoder.mlmodelc/model.mil.bak"){ + let vaeencoderMIL = resourcePath + "VAEEncoder.mlmodelc/model.mil" + try FileManager.default.removeItem(atPath: vaeencoderMIL) + try FileManager.default.copyItem(atPath: resourcePath + "VAEEncoder.mlmodelc/model.mil.bak", toPath: vaeencoderMIL) + await vaeEnSD(vaeMIL: vaeencoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + }else{ + let vaeencoderMIL = resourcePath + "VAEEncoder.mlmodelc/model.mil" + try FileManager.default.copyItem(atPath: vaeencoderMIL, toPath: resourcePath + "VAEEncoder.mlmodelc/model.mil.bak") + try FileManager.default.removeItem(atPath: resourcePath + "VAEEncoder.mlmodelc/coremldata.bin") + try FileManager.default.copyItem(at: Bundle.main.url(forResource: "en-coremldata", withExtension: "bin")!, to: URL(fileURLWithPath: resourcePath + "VAEEncoder.mlmodelc/coremldata.bin")) + await vaeEnSD(vaeMIL: vaeencoderMIL, height: ImageController.shared.height, width: ImageController.shared.width) + } + } + } + func modifyMILFile(path: String, oldDimensions: String, newDimensions: String) { + do { + let fileContent = try String(contentsOf: URL(fileURLWithPath: path), encoding: .utf8) + let modifiedContent = fileContent.replacingOccurrences(of: oldDimensions, with: newDimensions) + try modifiedContent.write(to: URL(fileURLWithPath: path),atomically: false, encoding: .utf8) + } catch { + print("Error modifying MIL file: \(error)") + } + } + + func modifyInputSize(_ url: URL, height: Int, width: Int) { + let encoderMetadataURL = url.appendingPathComponent("VAEEncoder.mlmodelc").appendingPathComponent("metadata.json") + guard let jsonData = try? Data(contentsOf: encoderMetadataURL), + var jsonArray = try? JSONSerialization.jsonObject(with: jsonData) as? [[String: Any]], + var jsonItem = jsonArray.first, + var inputSchema = jsonItem["inputSchema"] as? [[String: Any]], + var controlnetCond = inputSchema.first, + var shapeString = controlnetCond["shape"] as? String else { + return + } + + var shapeIntArray = shapeString.trimmingCharacters(in: CharacterSet(charactersIn: "[]")) + .components(separatedBy: ", ") + .compactMap { Int($0.trimmingCharacters(in: .whitespaces)) } + + shapeIntArray[3] = width + shapeIntArray[2] = height + shapeString = "[\(shapeIntArray.map { String($0) }.joined(separator: ", "))]" + + controlnetCond["shape"] = shapeString + inputSchema[0] = controlnetCond + jsonItem["inputSchema"] = inputSchema + jsonArray[0] = jsonItem + + if let updatedJsonData = try? JSONSerialization.data(withJSONObject: jsonArray, options: .prettyPrinted) { + try? updatedJsonData.write(to: encoderMetadataURL) + print("update metadata.") + } else { + print("Failed to update metadata.") + } + } + + func vaeDeSDXL(vaeMIL: String, height: Int, width: Int) { + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 4, 128, 128]", newDimensions: "[1, 4, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 128, 128]", newDimensions: "[1, 512, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 128, 128]", newDimensions: "[1, 32, 16, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 16384]", newDimensions: "[1, 32, 16, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 16384]", newDimensions: "[1, 512, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 16384, 512]", newDimensions: "[1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 16384, 1, 512]", newDimensions: "[1, \(height / 8 * width / 8), 1, 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 16384, 512]", newDimensions: "[1, 1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 16384, 16384]", newDimensions: "[1, 1, \(height / 8 * width / 8), \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 256, 256]", newDimensions: "[1, 512, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 256, 256]", newDimensions: "[1, 32, 16, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 512, 512]", newDimensions: "[1, 512, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 512, 512]", newDimensions: "[1, 32, 16, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 512, 512]", newDimensions: "[1, 256, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 512, 512]", newDimensions: "[1, 32, 8, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 1024, 1024]", newDimensions: "[1, 256, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 1024, 1024]", newDimensions: "[1, 32, 8, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 1024, 1024]", newDimensions: "[1, 128, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 4, 1024, 1024]", newDimensions: "[1, 32, 4, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 3, 1024, 1024]", newDimensions: "[1, 3, \(height), \(width)]") + } + + func vaeEnSDXL(vaeMIL: String, height: Int, width: Int) { + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 8, 128, 128]", newDimensions: "[1, 8, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 16384, 512]", newDimensions: "[1, 1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 16384, 16384]", newDimensions: "[1, 1, \(height / 8 * width / 8), \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 16384, 1, 512]", newDimensions: "[1, \(height / 8 * width / 8), 1, 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 16384, 512]", newDimensions: "[1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 16384]", newDimensions: "[1, 512, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 16384]", newDimensions: "[1, 32, 16, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 128, 128]", newDimensions: "[1, 32, 16, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 128, 128]", newDimensions: "[1, 512, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 257, 257]", newDimensions: "[1, 512, \(height / 4 + 1), \(width / 4 + 1)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 256, 256]", newDimensions: "[1, 32, 16, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 256, 256]", newDimensions: "[1, 512, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 256, 256]", newDimensions: "[1, 32, 8, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 256, 256]", newDimensions: "[1, 256, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 513, 513]", newDimensions: "[1, 256, \(height / 2 + 1), \(width / 2 + 1)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 512, 512]", newDimensions: "[1, 32, 8, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 512, 512]", newDimensions: "[1, 256, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 4, 512, 512]", newDimensions: "[1, 32, 4, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 512, 512]", newDimensions: "[1, 128, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 1025, 1025]", newDimensions: "[1, 128, \(height + 1), \(width + 1)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 1024, 1024]", newDimensions: "[1, 128, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 4, 1024, 1024]", newDimensions: "[1, 32, 4, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 3, 1024, 1024]", newDimensions: "[1, 3, \(height), \(width)]") + } + + func vaeDeSD(vaeMIL: String, height: Int, width: Int) { + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 4, 64, 64]", newDimensions: "[1, 4, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 64, 64]", newDimensions: "[1, 512, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 64, 64]", newDimensions: "[1, 32, 16, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 4096]", newDimensions: "[1, 32, 16, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 4096]", newDimensions: "[1, 512, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 4096, 512]", newDimensions: "[1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 4096, 1, 512]", newDimensions: "[1, \(height / 8 * width / 8), 1, 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 4096, 512]", newDimensions: "[1, 1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 4096, 4096]", newDimensions: "[1, 1, \(height / 8 * width / 8), \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 128, 128]", newDimensions: "[1, 512, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 128, 128]", newDimensions: "[1, 32, 16, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 256, 256]", newDimensions: "[1, 512, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 256, 256]", newDimensions: "[1, 32, 16, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 256, 256]", newDimensions: "[1, 256, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 256, 256]", newDimensions: "[1, 32, 8, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 512, 512]", newDimensions: "[1, 256, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 512, 512]", newDimensions: "[1, 32, 8, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 512, 512]", newDimensions: "[1, 128, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 4, 512, 512]", newDimensions: "[1, 32, 4, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 3, 512, 512]", newDimensions: "[1, 3, \(height), \(width)]") + } + + func vaeEnSD(vaeMIL: String, height: Int, width: Int) { + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 8, 64, 64]", newDimensions: "[1, 8, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 4096, 512]", newDimensions: "[1, 1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 1, 4096, 4096]", newDimensions: "[1, 1, \(height / 8 * width / 8), \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 4096, 1, 512]", newDimensions: "[1, \(height / 8 * width / 8), 1, 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 4096, 512]", newDimensions: "[1, \(height / 8 * width / 8), 512]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 4096]", newDimensions: "[1, 512, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 4096]", newDimensions: "[1, 32, 16, \(height / 8 * width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 64, 64]", newDimensions: "[1, 32, 16, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 64, 64]", newDimensions: "[1, 512, \(height / 8), \(width / 8)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 129, 129]", newDimensions: "[1, 512, \(height / 4 + 1), \(width / 4 + 1)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 16, 128, 128]", newDimensions: "[1, 32, 16, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 512, 128, 128]", newDimensions: "[1, 512, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 128, 128]", newDimensions: "[1, 32, 8, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 128, 128]", newDimensions: "[1, 256, \(height / 4), \(width / 4)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 257, 257]", newDimensions: "[1, 256, \(height / 2 + 1), \(width / 2 + 1)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 8, 256, 256]", newDimensions: "[1, 32, 8, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 256, 256, 256]", newDimensions: "[1, 256, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 4, 256, 256]", newDimensions: "[1, 32, 4, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 256, 256]", newDimensions: "[1, 128, \(height / 2), \(width / 2)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 513, 513]", newDimensions: "[1, 128, \(height + 1), \(width + 1)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 128, 512, 512]", newDimensions: "[1, 128, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 32, 4, 512, 512]", newDimensions: "[1, 32, 4, \(height), \(width)]") + modifyMILFile(path: vaeMIL, oldDimensions: "[1, 3, 512, 512]", newDimensions: "[1, 3, \(height), \(width)]") + } + func stopGenerate() async { generationStopped = true } diff --git a/Mochi Diffusion/Views/GalleryItemView.swift b/Mochi Diffusion/Views/GalleryItemView.swift index 596b6b4f..a853f928 100755 --- a/Mochi Diffusion/Views/GalleryItemView.swift +++ b/Mochi Diffusion/Views/GalleryItemView.swift @@ -78,6 +78,8 @@ struct GalleryItemView: View { UpscalingAnimationView() } } + .frame(width: 210, height: 210) + .background(image.averageColor?.blur(radius: 0.5).opacity(0.3)) } else { Color.clear } @@ -88,7 +90,7 @@ struct GalleryItemView: View { VStack { UpscalingAnimationView() .frame(width: 300, height: 300) - .border(.selection, width: 5) + .border(.selection, width: 3) } .frame(width: 350, height: 350) } diff --git a/Mochi Diffusion/Views/GalleryPreviewView.swift b/Mochi Diffusion/Views/GalleryPreviewView.swift index 6386e357..f252f56d 100755 --- a/Mochi Diffusion/Views/GalleryPreviewView.swift +++ b/Mochi Diffusion/Views/GalleryPreviewView.swift @@ -16,6 +16,7 @@ struct GalleryPreviewView: View { Image(image, scale: 1, label: Text("")) .resizable() .aspectRatio(contentMode: .fit) + .frame(width: 200, height: 200) if case let .running(progress) = generator.state, let progress = progress, progress.stepCount > 0 { let step = progress.step + 1 let stepValue = Double(step) / Double(progress.stepCount + 1) @@ -37,7 +38,7 @@ struct GalleryPreviewView: View { .padding(8) .background(.ultraThinMaterial, in: RoundedRectangle(cornerRadius: 8)) } - .aspectRatio(CGFloat(image.width / image.height), contentMode: .fit) +// .aspectRatio(CGFloat(image.width / image.height), contentMode: .fit) .padding(8) } } diff --git a/Mochi Diffusion/Views/GalleryView.swift b/Mochi Diffusion/Views/GalleryView.swift index e746fc1f..55cd2319 100755 --- a/Mochi Diffusion/Views/GalleryView.swift +++ b/Mochi Diffusion/Views/GalleryView.swift @@ -78,7 +78,7 @@ struct GalleryView: View { RoundedRectangle(cornerRadius: 2) .stroke( store.selectedId == sdi.id ? Color.accentColor : Color(nsColor: .controlBackgroundColor), - lineWidth: 4 + lineWidth: 2 ) ) .gesture(TapGesture(count: 2).onEnded { diff --git a/Mochi Diffusion/Views/SidebarControls/ControlNetView.swift b/Mochi Diffusion/Views/SidebarControls/ControlNetView.swift index cb7164b1..acc8c2bf 100755 --- a/Mochi Diffusion/Views/SidebarControls/ControlNetView.swift +++ b/Mochi Diffusion/Views/SidebarControls/ControlNetView.swift @@ -16,7 +16,7 @@ struct ControlNetView: View { .sidebarLabelFormat() HStack(alignment: .top) { - ImageWellView(image: controller.currentControlNets.first?.image, size: controller.currentModel?.inputSize) { image in + ImageWellView(image: controller.currentControlNets.first?.image, size: CGSize(width: controller.width, height: controller.height)) { image in if let image { await ImageController.shared.setControlNet(image: image) } else { diff --git a/Mochi Diffusion/Views/SidebarControls/SizeView.swift b/Mochi Diffusion/Views/SidebarControls/SizeView.swift index 73ec14ed..9af96425 100755 --- a/Mochi Diffusion/Views/SidebarControls/SizeView.swift +++ b/Mochi Diffusion/Views/SidebarControls/SizeView.swift @@ -9,12 +9,16 @@ import SwiftUI struct SizeView: View { @EnvironmentObject private var controller: ImageController - private let imageSizes = [ - 256, 320, 384, 448, 512, 576, 640, 704, 768 + private let sdimageSizes = [ + 512, 576, 640, 768, 832, 896 ] - + private let sdxlimageSizes = [ + 512, 576, 640, 768, 832, 896, 1024, 1152, 1216, 1280, 1344, 1536 + ] + var body: some View { HStack { + let imageSizes: [Int] = ImageController.shared.currentModel?.isXL ?? false ? sdxlimageSizes : sdimageSizes VStack(alignment: .leading) { Text( "Width:", diff --git a/Mochi Diffusion/Views/SidebarControls/StartingImageView.swift b/Mochi Diffusion/Views/SidebarControls/StartingImageView.swift index e4efa15a..7eed75ef 100755 --- a/Mochi Diffusion/Views/SidebarControls/StartingImageView.swift +++ b/Mochi Diffusion/Views/SidebarControls/StartingImageView.swift @@ -10,6 +10,8 @@ import SwiftUI struct StartingImageView: View { @EnvironmentObject private var controller: ImageController @State private var isInfoPopoverShown = false + @State private var isMaskPopoverShown = false + @State private var maskImage: NSImage? = nil var body: some View { Text( @@ -19,7 +21,7 @@ struct StartingImageView: View { .sidebarLabelFormat() HStack(alignment: .top) { - ImageWellView(image: controller.startingImage, size: controller.currentModel?.inputSize) { image in + ImageWellView(image: controller.startingImage, size: CGSize(width: controller.width, height: controller.height)) { image in if let image { ImageController.shared.setStartingImage(image: image) } else { @@ -27,18 +29,33 @@ struct StartingImageView: View { } } .frame(width: 90, height: 90) - + .overlay( + Image(nsImage: controller.maskImage.map { NSImage(cgImage: $0, size: NSSize(width: $0.width, height: $0.height)) } ?? NSImage()) + .resizable() + .aspectRatio(contentMode: .fit) + ) + Spacer() VStack(alignment: .trailing) { HStack { Button { + maskImage = nil Task { await ImageController.shared.selectStartingImage() } } label: { Image(systemName: "photo") } Button { + self.isMaskPopoverShown.toggle() + } label: { + Image(systemName: "paintbrush") + } + .disabled(controller.startingImage == nil) + + Button { + maskImage = nil + controller.maskImage = nil Task { await ImageController.shared.unsetStartingImage() } } label: { Image(systemName: "xmark") @@ -65,21 +82,130 @@ struct StartingImageView: View { .buttonStyle(PlainButtonStyle()) .popover(isPresented: self.$isInfoPopoverShown, arrowEdge: .top) { Text( - """ - Strength controls how closely the generated image resembles the starting image. - Use lower values to generate images that look similar to the starting image. - Use higher values to allow more creative freedom. + """ + Strength controls how closely the generated image resembles the starting image. + Use lower values to generate images that look similar to the starting image. + Use higher values to allow more creative freedom. - The size of the starting image must match the output image size of the current model. - """ + The size of the starting image must match the output image size of the current model. + """ ) .padding() } } MochiSlider(value: $controller.strength, bounds: 0.0...1.0, step: 0.05) + .popover(isPresented: self.$isMaskPopoverShown, arrowEdge: .top) { + MaskEditorView(startingImage: controller.startingImage, maskImage: $maskImage) + } } } +struct PathWrapper: Identifiable { + let id = UUID() + let path: Path +} + +struct MaskEditorView: View { + @EnvironmentObject private var controller: ImageController + let startingImage: CGImage? + @Binding var maskImage: NSImage? + @State private var startPoint: CGPoint? + @State private var endPoint: CGPoint? + @State private var paths: [PathWrapper] = [] + @State private var radius: CGFloat = 50 + + var body: some View { + VStack(alignment: .trailing){ + ZStack { + if let startingImage = startingImage { + Image(nsImage: NSImage(cgImage: startingImage, size: NSSize(width: startingImage.width, height: startingImage.height))) + .resizable() + .aspectRatio(contentMode: .fit) + } + if let maskImage = maskImage { + Image(nsImage: maskImage) + .resizable() + .aspectRatio(contentMode: .fit) + } + } + HStack{ + Spacer() + Button { + paths.removeAll() + maskImage = nil + controller.maskImage = nil + } label: { + Image(systemName: "arrow.clockwise.circle") + .font(.system(size: 30)) + } + .buttonBorderShape(.circle) + .padding(20) + + Slider(value: $radius, in: 30...150, step: 5) + Spacer() + } + Spacer() + } + + .gesture( + DragGesture(minimumDistance: 0) + .onChanged { value in + if startPoint == nil { + startPoint = value.location + } else { + endPoint = value.location + updateMaskImage() + } + } + .onEnded { _ in + startPoint = nil + endPoint = nil + controller.maskImage = maskImage?.cgImage + } + ) + } + + func updateMaskImage() { + guard let endPoint = endPoint else { return } + + let center = CGPoint(x: (endPoint.x) , y: (endPoint.y)) + let adjustedY = CGFloat(startingImage?.height ?? 0) - center.y + + let path = Path { path in + path.addEllipse(in: CGRect(x: center.x - radius, y: adjustedY - radius, width: radius * 2, height: radius * 2)) + } + paths.append(PathWrapper(path: path)) + drawPaths() + } + + func drawPaths() { + guard let startingImage = startingImage else { return } + let imageSize = CGSize(width: startingImage.width, height: startingImage.height) + let imageRect = CGRect(origin: .zero, size: imageSize) + + let width = Int(imageSize.width) + let height = Int(imageSize.height) + + guard let maskContext = CGContext(data: nil, width: width, height: height, bitsPerComponent: 8, bytesPerRow: 0, space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) else { return } + + maskContext.clear(imageRect) + + maskContext.setFillColor(CGColor.black) + maskContext.setBlendMode(.normal) + + for pathWrapper in paths { + maskContext.addPath(pathWrapper.path.cgPath) + maskContext.fillPath() + } + + if let maskCGImage = maskContext.makeImage() { + let maskImage = NSImage(cgImage: maskCGImage, size: imageSize) + self.maskImage = maskImage + } + } +} + + #Preview { StartingImageView() .environmentObject(ImageController.shared) diff --git a/Mochi Diffusion/Views/SidebarView.swift b/Mochi Diffusion/Views/SidebarView.swift index e0fefa62..3bcf02f2 100755 --- a/Mochi Diffusion/Views/SidebarView.swift +++ b/Mochi Diffusion/Views/SidebarView.swift @@ -23,6 +23,10 @@ struct SidebarView: View { ModelView() Spacer().frame(height: 6) } + Group { + SizeView() + Spacer().frame(height: 6) + } Group { NumberOfImagesView() Spacer().frame(height: 6)