├── .gitattributes ├── .gitignore ├── .spi.yml ├── .swiftpm └── xcode │ └── package.xcworkspace │ ├── contents.xcworkspacedata │ └── xcshareddata │ └── IDEWorkspaceChecks.plist ├── LICENSE ├── Package.resolved ├── Package.swift ├── README.md ├── Sources └── coreml-stable-diffusion-swift │ ├── enum │ ├── Errors.swift │ └── ProgressState.swift │ ├── ext │ ├── NSImage+.swift │ └── URL+.swift │ ├── helper │ ├── HelperFile.swift │ ├── HelperImage.swift │ └── HelperPipeline.swift │ ├── manager │ └── GenerativeManager.swift │ ├── model │ ├── GenerativeModel.swift │ └── Scheduler.swift │ └── protocol │ └── IGenerativeManager.swift └── Tests └── coreml-stable-diffusion-swiftTests └── coreml_stable_diffusion_swiftTests.swift /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Xcode 2 | # 3 | # gitignore contributors: remember to update Global/Xcode.gitignore, Objective-C.gitignore & Swift.gitignore 4 | 5 | ## User settings 6 | xcuserdata/ 7 | 8 | ## compatibility with Xcode 8 and earlier (ignoring not required starting Xcode 9) 9 | *.xcscmblueprint 10 | *.xccheckout 11 | 12 | ## compatibility with Xcode 3 and earlier (ignoring not required starting Xcode 4) 13 | build/ 14 | DerivedData/ 15 | *.moved-aside 16 | *.pbxuser 17 | !default.pbxuser 18 | *.mode1v3 19 | !default.mode1v3 20 | *.mode2v3 21 | !default.mode2v3 22 | *.perspectivev3 23 | !default.perspectivev3 24 | 25 | ## Obj-C/Swift specific 26 | *.hmap 27 | 28 | ## App packaging 29 | *.ipa 30 | *.dSYM.zip 31 | *.dSYM 32 | 33 | ## Playgrounds 34 | timeline.xctimeline 35 | playground.xcworkspace 36 | 37 | # Swift Package Manager 38 | # 39 | # Add this line if you want to avoid checking in source code from Swift Package Manager dependencies. 40 | # Packages/ 41 | # Package.pins 42 | # Package.resolved 43 | # *.xcodeproj 44 | # 45 | # Xcode automatically generates this directory with a .xcworkspacedata file and xcuserdata 46 | # hence it is not needed unless you have added a package configuration file to your project 47 | # .swiftpm 48 | 49 | .build/ 50 | 51 | # CocoaPods 52 | # 53 | # We recommend against adding the Pods directory to your .gitignore. However 54 | # you should judge for yourself, the pros and cons are mentioned at: 55 | # https://guides.cocoapods.org/using/using-cocoapods.html#should-i-check-the-pods-directory-into-source-control 56 | # 57 | # Pods/ 58 | # 59 | # Add this line if you want to avoid checking in source code from the Xcode workspace 60 | # *.xcworkspace 61 | 62 | # Carthage 63 | # 64 | # Add this line if you want to avoid checking in source code from Carthage dependencies. 65 | # Carthage/Checkouts 66 | 67 | Carthage/Build/ 68 | 69 | # Accio dependency management 70 | Dependencies/ 71 | .accio/ 72 | 73 | # fastlane 74 | # 75 | # It is recommended to not store the screenshots in the git repo. 76 | # Instead, use fastlane to re-generate the screenshots whenever they are needed. 77 | # For more information about the recommended setup visit: 78 | # https://docs.fastlane.tools/best-practices/source-control/#source-control 79 | 80 | fastlane/report.xml 81 | fastlane/Preview.html 82 | fastlane/screenshots/**/*.png 83 | fastlane/test_output 84 | 85 | # Code Injection 86 | # 87 | # After new code Injection tools there's a generated folder /iOSInjectionProject 88 | # https://github.com/johnno1962/injectionforxcode 89 | 90 | iOSInjectionProject/ 91 | -------------------------------------------------------------------------------- /.spi.yml: -------------------------------------------------------------------------------- 1 | version: 1 2 | builder: 3 | configs: 4 | - documentation_targets: [coreml-stable-diffusion-swift] 5 | -------------------------------------------------------------------------------- /.swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata: -------------------------------------------------------------------------------- 1 | 2 | 4 | 6 | 7 | 8 | -------------------------------------------------------------------------------- /.swiftpm/xcode/package.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | IDEDidComputeMac32BitWarning 6 | 7 | 8 | 9 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2023 Igor Shelopaev 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Package.resolved: -------------------------------------------------------------------------------- 1 | { 2 | "pins" : [ 3 | { 4 | "identity" : "files", 5 | "kind" : "remoteSourceControl", 6 | "location" : "https://github.com/JohnSundell/Files.git", 7 | "state" : { 8 | "revision" : "d273b5b7025d386feef79ef6bad7de762e106eaf", 9 | "version" : "4.2.0" 10 | } 11 | }, 12 | { 13 | "identity" : "ml-stable-diffusion", 14 | "kind" : "remoteSourceControl", 15 | "location" : "https://github.com/apple/ml-stable-diffusion.git", 16 | "state" : { 17 | "revision" : "5a170d29cf38e674b80541d7ce22929c6a11cdde", 18 | "version" : "1.1.1" 19 | } 20 | }, 21 | { 22 | "identity" : "swift-argument-parser", 23 | "kind" : "remoteSourceControl", 24 | "location" : "https://github.com/apple/swift-argument-parser.git", 25 | "state" : { 26 | "revision" : "41982a3656a71c768319979febd796c6fd111d5c", 27 | "version" : "1.5.0" 28 | } 29 | } 30 | ], 31 | "version" : 2 32 | } 33 | -------------------------------------------------------------------------------- /Package.swift: -------------------------------------------------------------------------------- 1 | // swift-tools-version: 5.7 2 | // The swift-tools-version declares the minimum version of Swift required to build this package. 3 | 4 | import PackageDescription 5 | 6 | let package = Package( 7 | name: "coreml-stable-diffusion-swift", 8 | platforms: [ 9 | .macOS(.v13), 10 | .iOS("16.2"), // Updated to support iOS 16.2 11 | ], 12 | products: [ 13 | // Products define the executables and libraries a package produces, and make them visible to other packages. 14 | .library( 15 | name: "coreml-stable-diffusion-swift", 16 | targets: ["coreml-stable-diffusion-swift"]), 17 | ], 18 | dependencies: [ 19 | // Dependencies declare other packages that this package depends on. 20 | .package(url: "https://github.com/apple/ml-stable-diffusion.git", .exactItem("1.1.1")), 21 | .package(url: "https://github.com/JohnSundell/Files.git", 22 | from: "4.2.0") 23 | ], 24 | targets: [ 25 | // Targets are the basic building blocks of a package. A target can define a module or a test suite. 26 | // Targets can depend on other targets in this package, and on products in packages this package depends on. 27 | .target( 28 | name: "coreml-stable-diffusion-swift", 29 | dependencies: [ 30 | .product(name: "StableDiffusion", package: "ml-stable-diffusion"), 31 | .product(name: "Files", package: "Files"), 32 | ]), 33 | .testTarget( 34 | name: "coreml-stable-diffusion-swiftTests", 35 | dependencies: ["coreml-stable-diffusion-swift"]), 36 | ] 37 | ) 38 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # CoreML stable diffusion image generation 2 | 3 | [![](https://img.shields.io/endpoint?url=https%3A%2F%2Fswiftpackageindex.com%2Fapi%2Fpackages%2Fswiftuiux%2Fcoreml-stable-diffusion-swift%2Fbadge%3Ftype%3Dplatforms)](https://swiftpackageindex.com/swiftuiux/coreml-stable-diffusion-swift) 4 | 5 | The package is a mediator between Apple's Core ML Stable Diffusion implementation and your app that let you run text-to-image or image-to-image models 6 | 7 | ## [SwiftUI example](https://github.com/swiftuiux/coreml-stable-diffusion-swift-example) 8 | 9 | ## [Documentation(API)](https://swiftpackageindex.com/swiftuiux/coreml-stable-diffusion-swift/main/documentation/coreml_stable_diffusion_swift) 10 | 11 | ## How to use the package 12 | ### 1. Create GenerativeManager 13 | 14 | ```swift 15 | let manager = GenerativeManager() 16 | ``` 17 | ### 2. Run async method **generate** 18 | 19 | ```swift 20 | let images: [CGImage?] = try await manager.generate( 21 | with: config, 22 | by: pipeline 23 | ) 24 | ``` 25 | ### Performance 26 | 27 | The speed can be unpredictable. Sometimes a model will suddenly run a lot slower than before. It appears as if Core ML is trying to be smart in how it schedules things, but doesn’t always optimal. 28 | 29 | 30 | 31 | 32 | ![The concept](https://github.com/swiftuiux/coreml-stable-diffusion-swift-example/blob/main/img/img_01.png) 33 | 34 | ![The concept](https://github.com/swiftuiux/coreml-stable-diffusion-swift-example/blob/main/img/img_03.png) 35 | 36 | ### Typical set of files for a model und the purpose of each file 37 | 38 | | File Name | Description | 39 | |--------------------------------------|------------------------------------------------------------------| 40 | | `TextEncoder.mlmodelc` | Encodes input text into a vector space for further processing. | 41 | | `Unet.mlmodelc` | Core model handling the transformation of encoded vectors into intermediate image representations. | 42 | | `UnetChunk1.mlmodelc` | First segment of a segmented U-Net model for optimized processing in environments with memory constraints. | 43 | | `UnetChunk2.mlmodelc` | Second segment of the segmented U-Net model, completing the tasks started by the first chunk. | 44 | | `VAEDecoder.mlmodelc` | Decodes the latent representations into final image outputs. | 45 | | `VAEEncoder.mlmodelc` | Compresses input image data into a latent space for reconstruction or further processing. | 46 | | `SafetyChecker.mlmodelc` | Ensures generated content adheres to safety guidelines by checking against predefined criteria. | 47 | | `vocab.json` | Contains the vocabulary used by the text encoder for tokenization and encoding processes. | 48 | | `merges.txt` | Stores the merging rules for byte-pair encoding used in the text encoder. | 49 | 50 | ## Documentation(API) 51 | - You need to have Xcode 13 installed in order to have access to Documentation Compiler (DocC) 52 | - Go to Product > Build Documentation or **⌃⇧⌘ D** 53 | 54 | ## Used packages 55 | - [Files](https://github.com/JohnSundell/Files) 56 | - [Stable Diffusion](https://github.com/apple/ml-stable-diffusion) 57 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/enum/Errors.swift: -------------------------------------------------------------------------------- 1 | // 2 | // Errors.swift 3 | // 4 | // 5 | // Created by Igor on 24.03.2023. 6 | // 7 | 8 | import Foundation 9 | 10 | /// Set of errors for generating images process for the current package 11 | @available(iOS 16.2, macOS 13.1, *) 12 | public enum Errors : Error{ 13 | /// Pipeline is not defined 14 | case pipelineIsNotDefined 15 | } 16 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/enum/ProgressState.swift: -------------------------------------------------------------------------------- 1 | // 2 | // ProgressState.swift 3 | // 4 | // 5 | // Created by Igor on 22.03.2023. 6 | // 7 | 8 | import Foundation 9 | 10 | /// Indicates of the current stage of generating images 11 | @available(iOS 16.2, macOS 13.1, *) 12 | public enum ProgressState{ 13 | /// At the begging after app was launched 14 | case idle 15 | /// Currently running 16 | case running 17 | /// Finished successfully 18 | case finished 19 | /// Canceling 20 | case canceling 21 | /// Failed with an error 22 | case failed 23 | } 24 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/ext/NSImage+.swift: -------------------------------------------------------------------------------- 1 | // 2 | // NSImage+.swift 3 | // 4 | // 5 | // Created by Igor on 23.03.2023. 6 | // 7 | 8 | import Foundation 9 | #if canImport(AppKit) 10 | import AppKit 11 | #endif 12 | 13 | #if os(macOS) 14 | 15 | /// https://stackoverflow.com/questions/24595908/swift-nsimage-to-cgimage 16 | @available(macOS 13.1, *) 17 | extension NSImage { 18 | 19 | /// Get CGImage from NSImage 20 | public var CGImage: CGImage? { 21 | get { 22 | guard let imageData = self.tiffRepresentation else{ 23 | return nil 24 | } 25 | let source = CGImageSourceCreateWithData(imageData as CFData, nil).unsafelyUnwrapped 26 | let maskRef = CGImageSourceCreateImageAtIndex(source, Int(0), nil) 27 | return maskRef.unsafelyUnwrapped 28 | } 29 | } 30 | } 31 | 32 | @available(macOS 13.1, *) 33 | extension NSImage { 34 | 35 | /// Returns the height of the current image. 36 | var height: CGFloat { 37 | return self.size.height 38 | } 39 | 40 | /// Returns the width of the current image. 41 | var width: CGFloat { 42 | return self.size.width 43 | } 44 | 45 | /// Returns a png representation of the current image. 46 | var PNGRepresentation: Data? { 47 | if let tiff = self.tiffRepresentation, let tiffData = NSBitmapImageRep(data: tiff) { 48 | return tiffData.representation(using: .png, properties: [:]) 49 | } 50 | 51 | return nil 52 | } 53 | 54 | /// Resizes and crops the image to the specified size. 55 | /// - Parameter innerSize: The target size to resize and crop the image to. 56 | /// - Returns: A new `NSImage` object that is resized and cropped to the specified size, or `nil` if the operation fails. 57 | func resizeAndCrop(to innerSize: NSSize) -> NSImage? { 58 | let aspectWidth = innerSize.width / self.size.width 59 | let aspectHeight = innerSize.height / self.size.height 60 | let aspectRatio = max(aspectWidth, aspectHeight) 61 | 62 | let aspectFillSize = NSSize(width: self.size.width * aspectRatio, height: self.size.height * aspectRatio) 63 | 64 | guard let resizedImage = self.resize(to: aspectFillSize) else { 65 | return nil 66 | } 67 | 68 | let croppedImage = resizedImage.crop(to: innerSize) 69 | 70 | return croppedImage 71 | } 72 | 73 | /// Resizes the image to the specified size while maintaining the aspect ratio. 74 | /// - Parameter size: The target size to resize the image to. 75 | /// - Returns: A new `NSImage` object that is resized to the specified size, or `nil` if the operation fails. 76 | func resize(to size: NSSize) -> NSImage? { 77 | let newImage = NSImage(size: size) 78 | newImage.lockFocus() 79 | self.draw(in: NSRect(origin: .zero, size: size), 80 | from: NSRect(origin: .zero, size: self.size), 81 | operation: .copy, 82 | fraction: 1.0) 83 | newImage.unlockFocus() 84 | return newImage 85 | } 86 | 87 | /// Crops the image to the specified size. 88 | /// - Parameter size: The target size to crop the image to. 89 | /// - Returns: A new `NSImage` object that is cropped to the specified size, or `nil` if the operation fails. 90 | func crop(to size: NSSize) -> NSImage? { 91 | guard let cgImage = self.cgImage(forProposedRect: nil, context: nil, hints: nil) else { 92 | return nil 93 | } 94 | 95 | let croppingRect = CGRect(x: (cgImage.width - Int(size.width)) / 2, 96 | y: (cgImage.height - Int(size.height)) / 2, 97 | width: Int(size.width), 98 | height: Int(size.height)) 99 | 100 | guard let croppedCGImage = cgImage.cropping(to: croppingRect) else { 101 | return nil 102 | } 103 | 104 | let croppedImage = NSImage(cgImage: croppedCGImage, size: size) 105 | return croppedImage 106 | } 107 | 108 | /// Saves the PNG representation of the current image to the HD. 109 | /// 110 | /// - parameter url: The location url to which to write the png file. 111 | func savePNGRepresentationToURL(url: URL) throws { 112 | if let png = self.PNGRepresentation { 113 | try png.write(to: url, options: .atomicWrite) 114 | } 115 | } 116 | } 117 | 118 | #endif 119 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/ext/URL+.swift: -------------------------------------------------------------------------------- 1 | // 2 | // URL+.swift 3 | // 4 | // 5 | // Created by Igor on 22.03.2023. 6 | // 7 | 8 | import Foundation 9 | 10 | internal extension URL { 11 | 12 | /// Check is directory 13 | @available(iOS 16.2, macOS 13.1, *) 14 | var isDirectory: Bool { 15 | return (try? resourceValues(forKeys: [.isDirectoryKey]))?.isDirectory == true 16 | } 17 | } 18 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/helper/HelperFile.swift: -------------------------------------------------------------------------------- 1 | // 2 | // HelperFile.swift 3 | // 4 | // 5 | // Created by Igor on 22.03.2023. 6 | // 7 | 8 | 9 | import Files 10 | #if canImport(AppKit) 11 | import AppKit 12 | #endif 13 | #if canImport(UIKit) 14 | import UIKit 15 | #endif 16 | 17 | #if os(macOS) 18 | /// Get list of models from the models directory 19 | /// - Returns: List of model 20 | @available(macOS 13.1, *) 21 | internal func listOfModels() async -> [GenerativeModel] { 22 | 23 | guard let docs = Folder.documents,let folder = try? docs.subfolder(at: "models") else{ 24 | return [] 25 | } 26 | 27 | return folder.subfolders.map{ 28 | GenerativeModel(url: $0.url, name: $0.name) 29 | } 30 | } 31 | 32 | /// Show the models directory in Finder 33 | @available(macOS 13.1, *) 34 | public func showInFinder() { 35 | guard let url = Folder.documents?.url else { return } 36 | 37 | if url.isDirectory { 38 | NSWorkspace.shared.selectFile(nil, inFileViewerRootedAtPath: url.path) 39 | } else { 40 | NSWorkspace.shared.activateFileViewerSelecting([url]) 41 | } 42 | } 43 | 44 | /// Create a folder 45 | /// - Parameter name: Name of the folder to create 46 | /// - Throws: Could not create a folder 47 | @available(macOS 13.1, *) 48 | public func initFolder(name : String) throws{ 49 | if let docs = Folder.documents{ 50 | if docs.containsSubfolder(named: name) == false{ 51 | try docs.createSubfolder(named: name) 52 | } 53 | } 54 | } 55 | 56 | #elseif os(iOS) 57 | /// Get list of models from the models directory 58 | /// - Returns: List of model 59 | @available(iOS 16.2, *) 60 | internal func listOfModels() async -> [GenerativeModel] { 61 | guard let docs = Folder.documents, let folder = try? docs.subfolder(at: "models") else { 62 | return [] 63 | } 64 | 65 | return folder.subfolders.map { 66 | GenerativeModel(url: $0.url, name: $0.name) 67 | } 68 | } 69 | 70 | /// Show the models directory in the Files app 71 | @available(iOS 16.2, *) 72 | public func showInFilesApp() { 73 | guard let url = Folder.documents?.url else { return } 74 | 75 | // Open the Files app at the specified URL 76 | UIApplication.shared.open(url) 77 | } 78 | 79 | /// Create a folder 80 | /// - Parameter name: Name of the folder to create 81 | /// - Throws: Could not create a folder 82 | @available(iOS 16.2, *) 83 | public func initFolder(name: String) throws { 84 | if let docs = Folder.documents { 85 | if docs.containsSubfolder(named: name) == false { 86 | try docs.createSubfolder(named: name) 87 | } 88 | } 89 | } 90 | 91 | #endif 92 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/helper/HelperImage.swift: -------------------------------------------------------------------------------- 1 | // 2 | // HelperImage.swift 3 | // 4 | // 5 | // Created by Igor on 23.03.2023. 6 | // 7 | 8 | import Foundation 9 | import SwiftUI 10 | #if canImport(AppKit) 11 | import AppKit 12 | #endif 13 | 14 | #if os(macOS) 15 | /// Get NSImage from data 16 | /// - Parameters: 17 | /// - data: Data 18 | /// - toSize: Fit to size 19 | /// - Returns: NSImage 20 | @available(macOS 13.1, *) 21 | public func getNSImage(from data : Data?, cropped toSize : NSSize? = nil) async -> NSImage? { 22 | 23 | guard let value = data else{ 24 | return nil 25 | } 26 | 27 | guard let nsImage = NSImage(data: value) else { 28 | return nil 29 | } 30 | 31 | guard let size = toSize ?? adjustSize(for: nsImage) else{ 32 | return nsImage 33 | } 34 | 35 | guard let cropped = nsImage.resizeAndCrop(to: size)else{ 36 | return nil 37 | } 38 | 39 | return cropped 40 | } 41 | 42 | /// Get image from CGImage 43 | /// - Parameter cgImage: CGImage 44 | /// - Returns: Image 45 | @available(macOS 13.1, *) 46 | public func getImage(cgImage : CGImage?) -> Image?{ 47 | 48 | guard let value = cgImage else{ 49 | return nil 50 | } 51 | 52 | let nsImage = NSImage(cgImage: value, size: .init(width: value.width, height: value.height)) 53 | 54 | return Image(nsImage: nsImage) 55 | } 56 | 57 | /// Determines the largest square size from a predefined set of sizes that can fit within the given image's dimensions. 58 | /// - Parameter image: The input `NSImage` to be adjusted. 59 | /// - Returns: An `NSSize` representing the largest possible square size that fits within the image's dimensions, or `nil` if the image is too small. 60 | @available(macOS 13.1, *) 61 | func adjustSize(for image: NSImage) -> NSSize? { 62 | let predefinedSizes: [CGFloat] = [512] 63 | 64 | let width = image.size.width 65 | let height = image.size.height 66 | 67 | guard width >= 256, height >= 256 else { 68 | return nil 69 | } 70 | 71 | let minDimension = min(width, height) 72 | 73 | // Find the largest predefined size that can fit within the image dimensions 74 | for size in predefinedSizes.reversed() { 75 | if minDimension >= size { 76 | return NSSize(width: size, height: size) 77 | } 78 | } 79 | 80 | return nil 81 | } 82 | 83 | #endif 84 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/helper/HelperPipeline.swift: -------------------------------------------------------------------------------- 1 | // 2 | // HelperPipeline.swift 3 | // 4 | // 5 | // Created by Igor on 23.03.2023. 6 | // 7 | 8 | import Foundation 9 | import StableDiffusion 10 | import CoreML 11 | 12 | /// Retrieves a StableDiffusionPipeline configured with the specified parameters. 13 | /// 14 | /// - Parameters: 15 | /// - url: The URL where the pipeline resources are located. 16 | /// - disableSafety: A boolean flag to disable safety checks. Default is `false`. 17 | /// - reduceMemory: A boolean flag to reduce memory usage. Default is `false`. 18 | /// - computeUnits: The compute units to be used by the pipeline. Default is `.cpuAndGPU`. 19 | /// 20 | /// - Throws: An error if the pipeline initialization fails. 21 | /// 22 | /// - Returns: A configured instance of `StableDiffusionPipeline`. 23 | 24 | @available(iOS 16.2, macOS 13.1, *) 25 | public func getDiffusionPipeline( 26 | for url : URL, 27 | _ controlNet: [String] = [], 28 | _ disableSafety : Bool = false, 29 | _ reduceMemory : Bool = false, 30 | _ computeUnits : MLComputeUnits = .cpuAndGPU 31 | ) throws -> StableDiffusionPipeline { 32 | 33 | // Initialize the MLModelConfiguration with the specified compute units. 34 | let config: MLModelConfiguration = .init() 35 | config.computeUnits = computeUnits 36 | 37 | // Return a new StableDiffusionPipeline instance, using the provided URL for resources. 38 | // The pipeline is configured with optional parameters to disable safety checks and reduce memory usage. 39 | return try .init( 40 | resourcesAt: url, // The URL where the pipeline resources are located. 41 | controlNet: controlNet, // An empty array for controlNet as no control network is being passed. 42 | configuration: config, // The MLModelConfiguration object with the compute units setting. 43 | disableSafety: disableSafety, // A boolean flag to disable safety checks if set to true. 44 | reduceMemory: reduceMemory // A boolean flag to reduce memory usage if set to true. 45 | ) 46 | } 47 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/manager/GenerativeManager.swift: -------------------------------------------------------------------------------- 1 | // 2 | // ModelsManager.swift 3 | // 4 | // 5 | // Created by Igor on 24.03.2023. 6 | // 7 | 8 | import Foundation 9 | import StableDiffusion 10 | import CoreML 11 | import OSLog 12 | 13 | private let logger = Logger(subsystem: "CoreML Stable diffusion", category: "GenerativeManager") 14 | 15 | /// The manager for generating images 16 | @available(iOS 16.2, macOS 13.1, *) 17 | public actor GenerativeManager: IGenerativeManager{ 18 | 19 | // MARK: - Life circle 20 | 21 | public init(){ } 22 | 23 | // MARK: - API 24 | 25 | /// Load list of models 26 | /// - Returns: List of models 27 | public func loadModels() async -> [GenerativeModel] { 28 | await listOfModels() 29 | } 30 | 31 | /// Generate images 32 | /// - Parameters: 33 | /// - config: mage generation configuration 34 | /// - pipeline: A pipeline used to generate image samples from text input using stable diffusion 35 | /// - Returns: An array of `imageCount` optional images. 36 | /// The images will be nil if safety checks were performed and found the result to be un-safe 37 | public func generate( 38 | with config: StableDiffusionPipeline.Configuration, 39 | by pipeline: StableDiffusionPipeline 40 | ) async throws -> [CGImage?] 41 | { 42 | try pipeline.generateImages(configuration: config) { progress in 43 | logger.info("Progress step: \(progress.step)") 44 | return !Task.isCancelled 45 | } 46 | } 47 | 48 | /// Get image generation configuration 49 | /// - Parameters: 50 | /// - prompt: Text prompt to guide sampling 51 | /// - negativePrompt: Negative text prompt to guide sampling 52 | /// - inputImage: Starting image for image2image or in-painting 53 | /// - seed: Random seed which to start generation 54 | /// - strength: Strength 55 | /// - guidanceScale: Controls the influence of the text prompt on sampling process (0=random images) 56 | /// - stepCount: Number of inference steps to perform 57 | /// - disableSafety: Safety checks 58 | /// - schedulerType: The type of Scheduler to use 59 | /// - Returns: Image generation configuration 60 | nonisolated 61 | public func getConfig ( 62 | _ prompt : String, 63 | _ negativePrompt : String, 64 | _ inputImage : CGImage?, 65 | _ seed : UInt32, 66 | _ strength : Float, 67 | _ guidanceScale : Float, 68 | _ stepCount : Float, 69 | _ disableSafety : Bool, 70 | _ schedulerType : StableDiffusionScheduler 71 | ) -> StableDiffusionPipeline.Configuration{ 72 | var cfg = StableDiffusionPipeline.Configuration(prompt: prompt) 73 | 74 | cfg.startingImage = inputImage 75 | cfg.negativePrompt = negativePrompt 76 | cfg.strength = strength 77 | cfg.stepCount = Int(stepCount) 78 | cfg.seed = seed 79 | cfg.guidanceScale = guidanceScale 80 | cfg.disableSafety = disableSafety 81 | cfg.schedulerType = schedulerType 82 | 83 | return cfg 84 | } 85 | } 86 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/model/GenerativeModel.swift: -------------------------------------------------------------------------------- 1 | // 2 | // GenerativeModel.swift 3 | // 4 | // 5 | // Created by Igor on 22.03.2023. 6 | // 7 | 8 | import Foundation 9 | import StableDiffusion 10 | 11 | /// Model profile description 12 | @available(iOS 16.2, macOS 13.1, *) 13 | public struct GenerativeModel: Hashable, Identifiable{ 14 | 15 | /// Identifier 16 | public let id = UUID() 17 | 18 | /// The path to the models files 19 | public let url: URL? 20 | 21 | /// Name of the model 22 | public let name: String 23 | 24 | // MARK: - Life circle 25 | public init(url: URL? = nil, name: String) { 26 | self.url = url 27 | self.name = name 28 | } 29 | } 30 | 31 | 32 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/model/Scheduler.swift: -------------------------------------------------------------------------------- 1 | // 2 | // SwiftUIView.swift 3 | // 4 | // 5 | // Created by Igor on 24.03.2023. 6 | // 7 | 8 | import SwiftUI 9 | import StableDiffusion 10 | 11 | @available(iOS 16.2, macOS 13.1, *) 12 | public extension GenerativeManager{ 13 | 14 | /// Schedulers compatible with StableDiffusionPipeline 15 | struct Scheduler: Hashable{ 16 | /// Schedulers compatible with StableDiffusionPipeline 17 | public let type : StableDiffusionScheduler 18 | ///Name 19 | public let name : String 20 | /// Description 21 | public let description : String 22 | 23 | public func hash(into hasher: inout Hasher) { 24 | hasher.combine(name) 25 | } 26 | } 27 | 28 | /// Default scheduler 29 | static let defaultScheduler : Scheduler = .init(type: .dpmSolverMultistepScheduler, name: "dpmSolverMultistep", description: "Scheduler that uses a second order DPM-Solver++ algorithm") 30 | 31 | /// Set of schedulers 32 | static let schedulers: [Scheduler] = [ 33 | defaultScheduler, 34 | .init(type: .pndmScheduler, name: "pndm", description: "Scheduler that uses a pseudo-linear multi-step (PLMS) method") 35 | ] 36 | 37 | } 38 | -------------------------------------------------------------------------------- /Sources/coreml-stable-diffusion-swift/protocol/IGenerativeManager.swift: -------------------------------------------------------------------------------- 1 | // 2 | // IGenerativeManager.swift 3 | // 4 | // 5 | // Created by Igor on 24.03.2023. 6 | // 7 | 8 | import Foundation 9 | import StableDiffusion 10 | import CoreML 11 | 12 | /// The interface of the manager for generating images 13 | @available(iOS 16.2, macOS 13.1, *) 14 | public protocol IGenerativeManager{ 15 | 16 | /// Load list of models 17 | func loadModels() async -> [GenerativeModel] 18 | 19 | /// Generate images 20 | func generate( 21 | with config: StableDiffusionPipeline.Configuration, 22 | by pipeline: StableDiffusionPipeline 23 | ) async throws -> [CGImage?] 24 | 25 | nonisolated 26 | func getConfig ( 27 | _ prompt : String, 28 | _ negativePrompt : String, 29 | _ inputImage : CGImage?, 30 | _ seed : UInt32, 31 | _ strength : Float, 32 | _ guidanceScale : Float, 33 | _ stepCount : Float, 34 | _ disableSafety : Bool, 35 | _ schedulerType : StableDiffusionScheduler 36 | ) -> StableDiffusionPipeline.Configuration 37 | 38 | } 39 | -------------------------------------------------------------------------------- /Tests/coreml-stable-diffusion-swiftTests/coreml_stable_diffusion_swiftTests.swift: -------------------------------------------------------------------------------- 1 | import XCTest 2 | @testable import coreml_stable_diffusion_swift 3 | 4 | final class coreml_stable_diffusion_swiftTests: XCTestCase { 5 | func testExample() throws { 6 | // This is an example of a functional test case. 7 | // Use XCTAssert and related functions to verify your tests produce the correct 8 | // results. 9 | // XCTAssertEqual(coreml_stable_diffusion_swift().text, "Hello, World!") 10 | } 11 | } 12 | --------------------------------------------------------------------------------