Repository: mortenjust/maple-diffusion Branch: main Commit: ad580b7a69cb Files: 53 Total size: 149.2 KB Directory structure: gitextract_i0tx6_zz/ ├── .gitignore ├── .swiftpm/ │ └── xcode/ │ └── package.xcworkspace/ │ └── contents.xcworkspacedata ├── Converter Script/ │ └── native-convert.py ├── Examples/ │ ├── Simple/ │ │ └── SimpleDiffusion/ │ │ ├── SimpleDiffusion/ │ │ │ ├── Assets.xcassets/ │ │ │ │ ├── AccentColor.colorset/ │ │ │ │ │ └── Contents.json │ │ │ │ ├── AppIcon.appiconset/ │ │ │ │ │ └── Contents.json │ │ │ │ └── Contents.json │ │ │ ├── ContentView.swift │ │ │ ├── Preview Content/ │ │ │ │ └── Preview Assets.xcassets/ │ │ │ │ └── Contents.json │ │ │ ├── SimpleDiffusion.entitlements │ │ │ └── SimpleDiffusionApp.swift │ │ └── SimpleDiffusion.xcodeproj/ │ │ ├── project.pbxproj │ │ └── project.xcworkspace/ │ │ ├── contents.xcworkspacedata │ │ └── xcshareddata/ │ │ ├── IDEWorkspaceChecks.plist │ │ └── swiftpm/ │ │ └── Package.resolved │ └── Single Line Diffusion/ │ ├── Single Line Diffusion/ │ │ ├── Assets.xcassets/ │ │ │ ├── AccentColor.colorset/ │ │ │ │ └── Contents.json │ │ │ ├── AppIcon.appiconset/ │ │ │ │ └── Contents.json │ │ │ └── Contents.json │ │ ├── ContentView.swift │ │ ├── Preview Content/ │ │ │ └── Preview Assets.xcassets/ │ │ │ └── Contents.json │ │ ├── Single_Line_Diffusion.entitlements │ │ └── Single_Line_DiffusionApp.swift │ └── Single Line Diffusion.xcodeproj/ │ ├── project.pbxproj │ └── project.xcworkspace/ │ ├── contents.xcworkspacedata │ └── xcshareddata/ │ ├── IDEWorkspaceChecks.plist │ └── swiftpm/ │ └── Package.resolved ├── LICENSE ├── Package.swift ├── README.md ├── Sources/ │ └── MapleDiffusion/ │ ├── Diffusion.swift │ ├── MPS/ │ │ ├── BPETokenizer.swift │ │ ├── Coder/ │ │ │ ├── Decoder.swift │ │ │ ├── Encoder.swift │ │ │ └── MPSGraph+Coder.swift │ │ ├── Diffuser.swift │ │ ├── MPSGraph+Transformer.swift │ │ ├── MPSGraph.swift │ │ ├── MapleDiffusion.swift │ │ ├── Scheduler.swift │ │ ├── TextGuidance.swift │ │ └── UNet.swift │ ├── Model/ │ │ ├── GenResult.swift │ │ ├── GeneratorState.swift │ │ ├── SampleInput.swift │ │ └── SampleOutput.swift │ ├── Support/ │ │ ├── CGImage + ItemProvider.swift │ │ ├── Diffusion + Generate.swift │ │ ├── Drag + Drop Helpers.swift │ │ ├── MPSGraphTensorData.swift │ │ ├── ModelFetcher.swift │ │ ├── NSImage + resizing.swift │ │ └── URLSession + async download.swift │ └── Views/ │ └── DiffusionImage.swift └── Tests/ └── MapleDiffusionTests/ └── ModelFetcherTests.swift ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ Examples/Official/maple-diffusion/bins/*.txt Examples/Official/maple-diffusion/bins/*.bin xcuserdata/ .DS_Store ================================================ FILE: .swiftpm/xcode/package.xcworkspace/contents.xcworkspacedata ================================================ ================================================ FILE: Converter Script/native-convert.py ================================================ #!/usr/bin/env python3 import sys if len(sys.argv) < 2: raise ValueError(f"Usage: {sys.argv[0]} path_to_ckpt") from pathlib import Path import torch as th import numpy as np ckpt = th.load(sys.argv[1], map_location="cpu") outpath = Path("./bins") outpath.mkdir(exist_ok=True) # vocab for clip vocab_url = "https://openaipublic.blob.core.windows.net/clip/bpe_simple_vocab_16e6.txt" vocab_dest = outpath / vocab_url.split("/")[-1] if not vocab_dest.exists(): print("downloading clip vocab") import requests with requests.get(vocab_url, stream=True) as r: assert r.status_code == 200, f"{vocab_url} failed to download. please copy it to {vocab_dest} manually." with vocab_dest.open('wb') as vf: for c in r.iter_content(chunk_size=8192): vf.write(c) print("downloaded clip vocab") # model weights for k, v in ckpt["state_dict"].items(): if not hasattr(v, "numpy"): continue v.numpy().astype('float16').tofile(outpath / (k + ".bin")) print("exporting state_dict", k, end="\r") print("\nexporting other stuff...") # other stuff th.exp(-th.log(th.tensor([10000])) * th.arange(0, 160) / 160).numpy().tofile(outpath / "temb_coefficients_fp32.bin") np.triu(np.ones((1,1,77,77), dtype=np.float16) * -65500.0, k=1).astype(np.float16).tofile(outpath / "causal_mask.bin") np.array([0.14013671875, 0.0711669921875, -0.03271484375, -0.11407470703125, 0.126220703125, 0.10101318359375, 0.034515380859375, -0.1383056640625, 0.126220703125, 0.07733154296875, 0.042633056640625, -0.177978515625]).astype(np.float16).tofile(outpath / "aux_output_conv.weight.bin") np.array([0.423828125, 0.471923828125, 0.473876953125]).astype(np.float16).tofile(outpath / "aux_output_conv.bias.bin") print(f"Done!") ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/Assets.xcassets/AccentColor.colorset/Contents.json ================================================ { "colors" : [ { "idiom" : "universal" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/Assets.xcassets/AppIcon.appiconset/Contents.json ================================================ { "images" : [ { "idiom" : "mac", "scale" : "1x", "size" : "16x16" }, { "idiom" : "mac", "scale" : "2x", "size" : "16x16" }, { "idiom" : "mac", "scale" : "1x", "size" : "32x32" }, { "idiom" : "mac", "scale" : "2x", "size" : "32x32" }, { "idiom" : "mac", "scale" : "1x", "size" : "128x128" }, { "idiom" : "mac", "scale" : "2x", "size" : "128x128" }, { "idiom" : "mac", "scale" : "1x", "size" : "256x256" }, { "idiom" : "mac", "scale" : "2x", "size" : "256x256" }, { "idiom" : "mac", "scale" : "1x", "size" : "512x512" }, { "idiom" : "mac", "scale" : "2x", "size" : "512x512" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/Assets.xcassets/Contents.json ================================================ { "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/ContentView.swift ================================================ // // ContentView.swift // SimpleDiffusion // // Created by Morten Just on 10/21/22. // import SwiftUI import MapleDiffusion import Combine struct ContentView: View { // 1 @StateObject var sd = Diffusion() @State var prompt = "" @State var steps = 20 @State var guidance : Double = 7.5 @State var image : CGImage? @State var inputImage: CGImage? @State var imagePublisher = Diffusion.placeholderPublisher @State var progress : Double = 0 var anyProgress : Double { sd.loadingProgress < 1 ? sd.loadingProgress : progress } var body: some View { VStack { DiffusionImage(image: $image, inputImage: $inputImage, progress: $progress) Spacer() TextField("Prompt", text: $prompt) // 3 .onSubmit { let input = SampleInput(prompt: prompt, initImage: inputImage, steps: steps) self.imagePublisher = sd.generate(input: input) } .disabled(!sd.isModelReady) HStack { TextField("Steps", value: $steps, formatter: NumberFormatter()) TextField("Guidance", value: $guidance, formatter: NumberFormatter()) } ProgressView(value: anyProgress) .opacity(anyProgress == 1 || anyProgress == 0 ? 0 : 1) } .task { // 2 let path = URL(string: "http://localhost:8080/Diffusion.zip")! do { try await sd.prepModels(remoteURL: path) } catch { assertionFailure("Hi, developer. You most likely don't have a local webserver running that serves the zip file with the transformed model files. See README.md") } } // 4 .onReceive(imagePublisher) { r in self.image = r.image self.progress = r.progress } .frame(minWidth: 200, minHeight: 200) } } struct ContentView_Previews: PreviewProvider { static var previews: some View { ContentView() } } extension URL { static var modelFolder = FileManager.default .urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] .appendingPathComponent("Photato/bins") } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/Preview Content/Preview Assets.xcassets/Contents.json ================================================ { "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/SimpleDiffusion.entitlements ================================================ ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion/SimpleDiffusionApp.swift ================================================ // // SimpleDiffusionApp.swift // SimpleDiffusion // // Created by Morten Just on 10/21/22. // import SwiftUI @main struct SimpleDiffusionApp: App { var body: some Scene { WindowGroup { ContentView() .ignoresSafeArea() }.windowStyle(.hiddenTitleBar) } } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion.xcodeproj/project.pbxproj ================================================ // !$*UTF8*$! { archiveVersion = 1; classes = { }; objectVersion = 56; objects = { /* Begin PBXBuildFile section */ C858BB302902A59000DF30AE /* SimpleDiffusionApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = C858BB2F2902A59000DF30AE /* SimpleDiffusionApp.swift */; }; C858BB322902A59000DF30AE /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C858BB312902A59000DF30AE /* ContentView.swift */; }; C858BB342902A59100DF30AE /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C858BB332902A59100DF30AE /* Assets.xcassets */; }; C858BB372902A59100DF30AE /* Preview Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C858BB362902A59100DF30AE /* Preview Assets.xcassets */; }; C858BB412902A68200DF30AE /* MapleDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = C858BB402902A68200DF30AE /* MapleDiffusion */; }; /* End PBXBuildFile section */ /* Begin PBXFileReference section */ C858BB2C2902A59000DF30AE /* SimpleDiffusion.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = SimpleDiffusion.app; sourceTree = BUILT_PRODUCTS_DIR; }; C858BB2F2902A59000DF30AE /* SimpleDiffusionApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SimpleDiffusionApp.swift; sourceTree = ""; }; C858BB312902A59000DF30AE /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = ""; }; C858BB332902A59100DF30AE /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; C858BB362902A59100DF30AE /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = ""; }; C858BB382902A59100DF30AE /* SimpleDiffusion.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = SimpleDiffusion.entitlements; sourceTree = ""; }; C858BB3E2902A63C00DF30AE /* maple-diffusion-package */ = {isa = PBXFileReference; lastKnownFileType = wrapper; name = "maple-diffusion-package"; path = ../../..; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFrameworksBuildPhase section */ C858BB292902A59000DF30AE /* Frameworks */ = { isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( C858BB412902A68200DF30AE /* MapleDiffusion in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXFrameworksBuildPhase section */ /* Begin PBXGroup section */ C858BB232902A59000DF30AE = { isa = PBXGroup; children = ( C858BB3E2902A63C00DF30AE /* maple-diffusion-package */, C858BB2E2902A59000DF30AE /* SimpleDiffusion */, C858BB2D2902A59000DF30AE /* Products */, C858BB3F2902A68200DF30AE /* Frameworks */, ); sourceTree = ""; }; C858BB2D2902A59000DF30AE /* Products */ = { isa = PBXGroup; children = ( C858BB2C2902A59000DF30AE /* SimpleDiffusion.app */, ); name = Products; sourceTree = ""; }; C858BB2E2902A59000DF30AE /* SimpleDiffusion */ = { isa = PBXGroup; children = ( C858BB2F2902A59000DF30AE /* SimpleDiffusionApp.swift */, C858BB312902A59000DF30AE /* ContentView.swift */, C858BB332902A59100DF30AE /* Assets.xcassets */, C858BB382902A59100DF30AE /* SimpleDiffusion.entitlements */, C858BB352902A59100DF30AE /* Preview Content */, ); path = SimpleDiffusion; sourceTree = ""; }; C858BB352902A59100DF30AE /* Preview Content */ = { isa = PBXGroup; children = ( C858BB362902A59100DF30AE /* Preview Assets.xcassets */, ); path = "Preview Content"; sourceTree = ""; }; C858BB3F2902A68200DF30AE /* Frameworks */ = { isa = PBXGroup; children = ( ); name = Frameworks; sourceTree = ""; }; /* End PBXGroup section */ /* Begin PBXNativeTarget section */ C858BB2B2902A59000DF30AE /* SimpleDiffusion */ = { isa = PBXNativeTarget; buildConfigurationList = C858BB3B2902A59100DF30AE /* Build configuration list for PBXNativeTarget "SimpleDiffusion" */; buildPhases = ( C858BB282902A59000DF30AE /* Sources */, C858BB292902A59000DF30AE /* Frameworks */, C858BB2A2902A59000DF30AE /* Resources */, ); buildRules = ( ); dependencies = ( ); name = SimpleDiffusion; packageProductDependencies = ( C858BB402902A68200DF30AE /* MapleDiffusion */, ); productName = SimpleDiffusion; productReference = C858BB2C2902A59000DF30AE /* SimpleDiffusion.app */; productType = "com.apple.product-type.application"; }; /* End PBXNativeTarget section */ /* Begin PBXProject section */ C858BB242902A59000DF30AE /* Project object */ = { isa = PBXProject; attributes = { BuildIndependentTargetsInParallel = 1; LastSwiftUpdateCheck = 1400; LastUpgradeCheck = 1400; TargetAttributes = { C858BB2B2902A59000DF30AE = { CreatedOnToolsVersion = 14.0.1; }; }; }; buildConfigurationList = C858BB272902A59000DF30AE /* Build configuration list for PBXProject "SimpleDiffusion" */; compatibilityVersion = "Xcode 14.0"; developmentRegion = en; hasScannedForEncodings = 0; knownRegions = ( en, Base, ); mainGroup = C858BB232902A59000DF30AE; productRefGroup = C858BB2D2902A59000DF30AE /* Products */; projectDirPath = ""; projectRoot = ""; targets = ( C858BB2B2902A59000DF30AE /* SimpleDiffusion */, ); }; /* End PBXProject section */ /* Begin PBXResourcesBuildPhase section */ C858BB2A2902A59000DF30AE /* Resources */ = { isa = PBXResourcesBuildPhase; buildActionMask = 2147483647; files = ( C858BB372902A59100DF30AE /* Preview Assets.xcassets in Resources */, C858BB342902A59100DF30AE /* Assets.xcassets in Resources */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXResourcesBuildPhase section */ /* Begin PBXSourcesBuildPhase section */ C858BB282902A59000DF30AE /* Sources */ = { isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( C858BB322902A59000DF30AE /* ContentView.swift in Sources */, C858BB302902A59000DF30AE /* SimpleDiffusionApp.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXSourcesBuildPhase section */ /* Begin XCBuildConfiguration section */ C858BB392902A59100DF30AE /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { ALWAYS_SEARCH_USER_PATHS = NO; CLANG_ANALYZER_NONNULL = YES; CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; CLANG_ENABLE_MODULES = YES; CLANG_ENABLE_OBJC_ARC = YES; CLANG_ENABLE_OBJC_WEAK = YES; CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; CLANG_WARN_BOOL_CONVERSION = YES; CLANG_WARN_COMMA = YES; CLANG_WARN_CONSTANT_CONVERSION = YES; CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; CLANG_WARN_DOCUMENTATION_COMMENTS = YES; CLANG_WARN_EMPTY_BODY = YES; CLANG_WARN_ENUM_CONVERSION = YES; CLANG_WARN_INFINITE_RECURSION = YES; CLANG_WARN_INT_CONVERSION = YES; CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; CLANG_WARN_STRICT_PROTOTYPES = YES; CLANG_WARN_SUSPICIOUS_MOVE = YES; CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; CLANG_WARN_UNREACHABLE_CODE = YES; CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; COPY_PHASE_STRIP = NO; DEBUG_INFORMATION_FORMAT = dwarf; ENABLE_STRICT_OBJC_MSGSEND = YES; ENABLE_TESTABILITY = YES; GCC_C_LANGUAGE_STANDARD = gnu11; GCC_DYNAMIC_NO_PIC = NO; GCC_NO_COMMON_BLOCKS = YES; GCC_OPTIMIZATION_LEVEL = 0; GCC_PREPROCESSOR_DEFINITIONS = ( "DEBUG=1", "$(inherited)", ); GCC_WARN_64_TO_32_BIT_CONVERSION = YES; GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; GCC_WARN_UNDECLARED_SELECTOR = YES; GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; GCC_WARN_UNUSED_FUNCTION = YES; GCC_WARN_UNUSED_VARIABLE = YES; MACOSX_DEPLOYMENT_TARGET = 12.3; MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; MTL_FAST_MATH = YES; ONLY_ACTIVE_ARCH = YES; SDKROOT = macosx; SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG; SWIFT_OPTIMIZATION_LEVEL = "-Onone"; }; name = Debug; }; C858BB3A2902A59100DF30AE /* Release */ = { isa = XCBuildConfiguration; buildSettings = { ALWAYS_SEARCH_USER_PATHS = NO; CLANG_ANALYZER_NONNULL = YES; CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; CLANG_ENABLE_MODULES = YES; CLANG_ENABLE_OBJC_ARC = YES; CLANG_ENABLE_OBJC_WEAK = YES; CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; CLANG_WARN_BOOL_CONVERSION = YES; CLANG_WARN_COMMA = YES; CLANG_WARN_CONSTANT_CONVERSION = YES; CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; CLANG_WARN_DOCUMENTATION_COMMENTS = YES; CLANG_WARN_EMPTY_BODY = YES; CLANG_WARN_ENUM_CONVERSION = YES; CLANG_WARN_INFINITE_RECURSION = YES; CLANG_WARN_INT_CONVERSION = YES; CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; CLANG_WARN_STRICT_PROTOTYPES = YES; CLANG_WARN_SUSPICIOUS_MOVE = YES; CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; CLANG_WARN_UNREACHABLE_CODE = YES; CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; COPY_PHASE_STRIP = NO; DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; ENABLE_NS_ASSERTIONS = NO; ENABLE_STRICT_OBJC_MSGSEND = YES; GCC_C_LANGUAGE_STANDARD = gnu11; GCC_NO_COMMON_BLOCKS = YES; GCC_WARN_64_TO_32_BIT_CONVERSION = YES; GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; GCC_WARN_UNDECLARED_SELECTOR = YES; GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; GCC_WARN_UNUSED_FUNCTION = YES; GCC_WARN_UNUSED_VARIABLE = YES; MACOSX_DEPLOYMENT_TARGET = 12.3; MTL_ENABLE_DEBUG_INFO = NO; MTL_FAST_MATH = YES; SDKROOT = macosx; SWIFT_COMPILATION_MODE = wholemodule; SWIFT_OPTIMIZATION_LEVEL = "-O"; }; name = Release; }; C858BB3C2902A59100DF30AE /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_ENTITLEMENTS = SimpleDiffusion/SimpleDiffusion.entitlements; CODE_SIGN_STYLE = Automatic; COMBINE_HIDPI_IMAGES = YES; CURRENT_PROJECT_VERSION = 1; DEVELOPMENT_ASSET_PATHS = "\"SimpleDiffusion/Preview Content\""; DEVELOPMENT_TEAM = UDGLB23X37; ENABLE_HARDENED_RUNTIME = YES; ENABLE_PREVIEWS = YES; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_KEY_NSHumanReadableCopyright = ""; LD_RUNPATH_SEARCH_PATHS = ( "$(inherited)", "@executable_path/../Frameworks", ); MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = app.otato.SimpleDiffusion; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; }; name = Debug; }; C858BB3D2902A59100DF30AE /* Release */ = { isa = XCBuildConfiguration; buildSettings = { ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_ENTITLEMENTS = SimpleDiffusion/SimpleDiffusion.entitlements; CODE_SIGN_STYLE = Automatic; COMBINE_HIDPI_IMAGES = YES; CURRENT_PROJECT_VERSION = 1; DEVELOPMENT_ASSET_PATHS = "\"SimpleDiffusion/Preview Content\""; DEVELOPMENT_TEAM = UDGLB23X37; ENABLE_HARDENED_RUNTIME = YES; ENABLE_PREVIEWS = YES; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_KEY_NSHumanReadableCopyright = ""; LD_RUNPATH_SEARCH_PATHS = ( "$(inherited)", "@executable_path/../Frameworks", ); MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = app.otato.SimpleDiffusion; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; }; name = Release; }; /* End XCBuildConfiguration section */ /* Begin XCConfigurationList section */ C858BB272902A59000DF30AE /* Build configuration list for PBXProject "SimpleDiffusion" */ = { isa = XCConfigurationList; buildConfigurations = ( C858BB392902A59100DF30AE /* Debug */, C858BB3A2902A59100DF30AE /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; C858BB3B2902A59100DF30AE /* Build configuration list for PBXNativeTarget "SimpleDiffusion" */ = { isa = XCConfigurationList; buildConfigurations = ( C858BB3C2902A59100DF30AE /* Debug */, C858BB3D2902A59100DF30AE /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; /* End XCConfigurationList section */ /* Begin XCSwiftPackageProductDependency section */ C858BB402902A68200DF30AE /* MapleDiffusion */ = { isa = XCSwiftPackageProductDependency; productName = MapleDiffusion; }; /* End XCSwiftPackageProductDependency section */ }; rootObject = C858BB242902A59000DF30AE /* Project object */; } ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion.xcodeproj/project.xcworkspace/contents.xcworkspacedata ================================================ ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist ================================================ IDEDidComputeMac32BitWarning ================================================ FILE: Examples/Simple/SimpleDiffusion/SimpleDiffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved ================================================ { "pins" : [ { "identity" : "zipfoundation", "kind" : "remoteSourceControl", "location" : "https://github.com/weichsel/ZIPFoundation.git", "state" : { "revision" : "f6a22e7da26314b38bf9befce34ae8e4b2543090", "version" : "0.9.15" } } ], "version" : 2 } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/Assets.xcassets/AccentColor.colorset/Contents.json ================================================ { "colors" : [ { "idiom" : "universal" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/Assets.xcassets/AppIcon.appiconset/Contents.json ================================================ { "images" : [ { "idiom" : "mac", "scale" : "1x", "size" : "16x16" }, { "idiom" : "mac", "scale" : "2x", "size" : "16x16" }, { "idiom" : "mac", "scale" : "1x", "size" : "32x32" }, { "idiom" : "mac", "scale" : "2x", "size" : "32x32" }, { "idiom" : "mac", "scale" : "1x", "size" : "128x128" }, { "idiom" : "mac", "scale" : "2x", "size" : "128x128" }, { "idiom" : "mac", "scale" : "1x", "size" : "256x256" }, { "idiom" : "mac", "scale" : "2x", "size" : "256x256" }, { "idiom" : "mac", "scale" : "1x", "size" : "512x512" }, { "idiom" : "mac", "scale" : "2x", "size" : "512x512" } ], "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/Assets.xcassets/Contents.json ================================================ { "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/ContentView.swift ================================================ // // ContentView.swift // Single Line Diffusion // // Created by Morten Just on 10/28/22. // import SwiftUI import MapleDiffusion struct ContentView: View { @State var image : CGImage? var body: some View { VStack { if let image { Image(image, scale: 1, label: Text("Generated")) } else { Text("Loading. See console.")} } .onAppear { Task.detached { for _ in 0...10 { image = try? await Diffusion.generate(localOrRemote: modelUrl, prompt: "cat astronaut") } } } .frame(minWidth: 500, minHeight: 500) } } struct ContentView_Previews: PreviewProvider { static var previews: some View { ContentView() } } let modelUrl = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0].appendingPathComponent("Photato/bins") ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/Preview Content/Preview Assets.xcassets/Contents.json ================================================ { "info" : { "author" : "xcode", "version" : 1 } } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/Single_Line_Diffusion.entitlements ================================================ ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion/Single_Line_DiffusionApp.swift ================================================ // // Single_Line_DiffusionApp.swift // Single Line Diffusion // // Created by Morten Just on 10/28/22. // import SwiftUI @main struct Single_Line_DiffusionApp: App { var body: some Scene { WindowGroup { ContentView() } } } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion.xcodeproj/project.pbxproj ================================================ // !$*UTF8*$! { archiveVersion = 1; classes = { }; objectVersion = 56; objects = { /* Begin PBXBuildFile section */ C8ED50B8290C204800153A3B /* Single_Line_DiffusionApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8ED50B7290C204800153A3B /* Single_Line_DiffusionApp.swift */; }; C8ED50BA290C204800153A3B /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C8ED50B9290C204800153A3B /* ContentView.swift */; }; C8ED50BC290C204800153A3B /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C8ED50BB290C204800153A3B /* Assets.xcassets */; }; C8ED50BF290C204800153A3B /* Preview Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = C8ED50BE290C204800153A3B /* Preview Assets.xcassets */; }; C8ED50C9290C207200153A3B /* MapleDiffusion in Frameworks */ = {isa = PBXBuildFile; productRef = C8ED50C8290C207200153A3B /* MapleDiffusion */; }; /* End PBXBuildFile section */ /* Begin PBXFileReference section */ C8ED50B4290C204800153A3B /* Single Line Diffusion.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = "Single Line Diffusion.app"; sourceTree = BUILT_PRODUCTS_DIR; }; C8ED50B7290C204800153A3B /* Single_Line_DiffusionApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Single_Line_DiffusionApp.swift; sourceTree = ""; }; C8ED50B9290C204800153A3B /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = ""; }; C8ED50BB290C204800153A3B /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = ""; }; C8ED50BE290C204800153A3B /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = ""; }; C8ED50C0290C204800153A3B /* Single_Line_Diffusion.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = Single_Line_Diffusion.entitlements; sourceTree = ""; }; C8ED50C6290C206600153A3B /* maple-diffusion-package */ = {isa = PBXFileReference; lastKnownFileType = wrapper; name = "maple-diffusion-package"; path = ../..; sourceTree = ""; }; /* End PBXFileReference section */ /* Begin PBXFrameworksBuildPhase section */ C8ED50B1290C204800153A3B /* Frameworks */ = { isa = PBXFrameworksBuildPhase; buildActionMask = 2147483647; files = ( C8ED50C9290C207200153A3B /* MapleDiffusion in Frameworks */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXFrameworksBuildPhase section */ /* Begin PBXGroup section */ C8ED50AB290C204800153A3B = { isa = PBXGroup; children = ( C8ED50C6290C206600153A3B /* maple-diffusion-package */, C8ED50B6290C204800153A3B /* Single Line Diffusion */, C8ED50B5290C204800153A3B /* Products */, C8ED50C7290C207200153A3B /* Frameworks */, ); sourceTree = ""; }; C8ED50B5290C204800153A3B /* Products */ = { isa = PBXGroup; children = ( C8ED50B4290C204800153A3B /* Single Line Diffusion.app */, ); name = Products; sourceTree = ""; }; C8ED50B6290C204800153A3B /* Single Line Diffusion */ = { isa = PBXGroup; children = ( C8ED50B7290C204800153A3B /* Single_Line_DiffusionApp.swift */, C8ED50B9290C204800153A3B /* ContentView.swift */, C8ED50BB290C204800153A3B /* Assets.xcassets */, C8ED50C0290C204800153A3B /* Single_Line_Diffusion.entitlements */, C8ED50BD290C204800153A3B /* Preview Content */, ); path = "Single Line Diffusion"; sourceTree = ""; }; C8ED50BD290C204800153A3B /* Preview Content */ = { isa = PBXGroup; children = ( C8ED50BE290C204800153A3B /* Preview Assets.xcassets */, ); path = "Preview Content"; sourceTree = ""; }; C8ED50C7290C207200153A3B /* Frameworks */ = { isa = PBXGroup; children = ( ); name = Frameworks; sourceTree = ""; }; /* End PBXGroup section */ /* Begin PBXNativeTarget section */ C8ED50B3290C204800153A3B /* Single Line Diffusion */ = { isa = PBXNativeTarget; buildConfigurationList = C8ED50C3290C204800153A3B /* Build configuration list for PBXNativeTarget "Single Line Diffusion" */; buildPhases = ( C8ED50B0290C204800153A3B /* Sources */, C8ED50B1290C204800153A3B /* Frameworks */, C8ED50B2290C204800153A3B /* Resources */, ); buildRules = ( ); dependencies = ( ); name = "Single Line Diffusion"; packageProductDependencies = ( C8ED50C8290C207200153A3B /* MapleDiffusion */, ); productName = "Single Line Diffusion"; productReference = C8ED50B4290C204800153A3B /* Single Line Diffusion.app */; productType = "com.apple.product-type.application"; }; /* End PBXNativeTarget section */ /* Begin PBXProject section */ C8ED50AC290C204800153A3B /* Project object */ = { isa = PBXProject; attributes = { BuildIndependentTargetsInParallel = 1; LastSwiftUpdateCheck = 1400; LastUpgradeCheck = 1400; TargetAttributes = { C8ED50B3290C204800153A3B = { CreatedOnToolsVersion = 14.0.1; }; }; }; buildConfigurationList = C8ED50AF290C204800153A3B /* Build configuration list for PBXProject "Single Line Diffusion" */; compatibilityVersion = "Xcode 14.0"; developmentRegion = en; hasScannedForEncodings = 0; knownRegions = ( en, Base, ); mainGroup = C8ED50AB290C204800153A3B; productRefGroup = C8ED50B5290C204800153A3B /* Products */; projectDirPath = ""; projectRoot = ""; targets = ( C8ED50B3290C204800153A3B /* Single Line Diffusion */, ); }; /* End PBXProject section */ /* Begin PBXResourcesBuildPhase section */ C8ED50B2290C204800153A3B /* Resources */ = { isa = PBXResourcesBuildPhase; buildActionMask = 2147483647; files = ( C8ED50BF290C204800153A3B /* Preview Assets.xcassets in Resources */, C8ED50BC290C204800153A3B /* Assets.xcassets in Resources */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXResourcesBuildPhase section */ /* Begin PBXSourcesBuildPhase section */ C8ED50B0290C204800153A3B /* Sources */ = { isa = PBXSourcesBuildPhase; buildActionMask = 2147483647; files = ( C8ED50BA290C204800153A3B /* ContentView.swift in Sources */, C8ED50B8290C204800153A3B /* Single_Line_DiffusionApp.swift in Sources */, ); runOnlyForDeploymentPostprocessing = 0; }; /* End PBXSourcesBuildPhase section */ /* Begin XCBuildConfiguration section */ C8ED50C1290C204800153A3B /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { ALWAYS_SEARCH_USER_PATHS = NO; CLANG_ANALYZER_NONNULL = YES; CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; CLANG_ENABLE_MODULES = YES; CLANG_ENABLE_OBJC_ARC = YES; CLANG_ENABLE_OBJC_WEAK = YES; CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; CLANG_WARN_BOOL_CONVERSION = YES; CLANG_WARN_COMMA = YES; CLANG_WARN_CONSTANT_CONVERSION = YES; CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; CLANG_WARN_DOCUMENTATION_COMMENTS = YES; CLANG_WARN_EMPTY_BODY = YES; CLANG_WARN_ENUM_CONVERSION = YES; CLANG_WARN_INFINITE_RECURSION = YES; CLANG_WARN_INT_CONVERSION = YES; CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; CLANG_WARN_STRICT_PROTOTYPES = YES; CLANG_WARN_SUSPICIOUS_MOVE = YES; CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; CLANG_WARN_UNREACHABLE_CODE = YES; CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; COPY_PHASE_STRIP = NO; DEBUG_INFORMATION_FORMAT = dwarf; ENABLE_STRICT_OBJC_MSGSEND = YES; ENABLE_TESTABILITY = YES; GCC_C_LANGUAGE_STANDARD = gnu11; GCC_DYNAMIC_NO_PIC = NO; GCC_NO_COMMON_BLOCKS = YES; GCC_OPTIMIZATION_LEVEL = 0; GCC_PREPROCESSOR_DEFINITIONS = ( "DEBUG=1", "$(inherited)", ); GCC_WARN_64_TO_32_BIT_CONVERSION = YES; GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; GCC_WARN_UNDECLARED_SELECTOR = YES; GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; GCC_WARN_UNUSED_FUNCTION = YES; GCC_WARN_UNUSED_VARIABLE = YES; MACOSX_DEPLOYMENT_TARGET = 12.3; MTL_ENABLE_DEBUG_INFO = INCLUDE_SOURCE; MTL_FAST_MATH = YES; ONLY_ACTIVE_ARCH = YES; SDKROOT = macosx; SWIFT_ACTIVE_COMPILATION_CONDITIONS = DEBUG; SWIFT_OPTIMIZATION_LEVEL = "-Onone"; }; name = Debug; }; C8ED50C2290C204800153A3B /* Release */ = { isa = XCBuildConfiguration; buildSettings = { ALWAYS_SEARCH_USER_PATHS = NO; CLANG_ANALYZER_NONNULL = YES; CLANG_ANALYZER_NUMBER_OBJECT_CONVERSION = YES_AGGRESSIVE; CLANG_CXX_LANGUAGE_STANDARD = "gnu++20"; CLANG_ENABLE_MODULES = YES; CLANG_ENABLE_OBJC_ARC = YES; CLANG_ENABLE_OBJC_WEAK = YES; CLANG_WARN_BLOCK_CAPTURE_AUTORELEASING = YES; CLANG_WARN_BOOL_CONVERSION = YES; CLANG_WARN_COMMA = YES; CLANG_WARN_CONSTANT_CONVERSION = YES; CLANG_WARN_DEPRECATED_OBJC_IMPLEMENTATIONS = YES; CLANG_WARN_DIRECT_OBJC_ISA_USAGE = YES_ERROR; CLANG_WARN_DOCUMENTATION_COMMENTS = YES; CLANG_WARN_EMPTY_BODY = YES; CLANG_WARN_ENUM_CONVERSION = YES; CLANG_WARN_INFINITE_RECURSION = YES; CLANG_WARN_INT_CONVERSION = YES; CLANG_WARN_NON_LITERAL_NULL_CONVERSION = YES; CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES; CLANG_WARN_OBJC_LITERAL_CONVERSION = YES; CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR; CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES; CLANG_WARN_RANGE_LOOP_ANALYSIS = YES; CLANG_WARN_STRICT_PROTOTYPES = YES; CLANG_WARN_SUSPICIOUS_MOVE = YES; CLANG_WARN_UNGUARDED_AVAILABILITY = YES_AGGRESSIVE; CLANG_WARN_UNREACHABLE_CODE = YES; CLANG_WARN__DUPLICATE_METHOD_MATCH = YES; COPY_PHASE_STRIP = NO; DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym"; ENABLE_NS_ASSERTIONS = NO; ENABLE_STRICT_OBJC_MSGSEND = YES; GCC_C_LANGUAGE_STANDARD = gnu11; GCC_NO_COMMON_BLOCKS = YES; GCC_WARN_64_TO_32_BIT_CONVERSION = YES; GCC_WARN_ABOUT_RETURN_TYPE = YES_ERROR; GCC_WARN_UNDECLARED_SELECTOR = YES; GCC_WARN_UNINITIALIZED_AUTOS = YES_AGGRESSIVE; GCC_WARN_UNUSED_FUNCTION = YES; GCC_WARN_UNUSED_VARIABLE = YES; MACOSX_DEPLOYMENT_TARGET = 12.3; MTL_ENABLE_DEBUG_INFO = NO; MTL_FAST_MATH = YES; SDKROOT = macosx; SWIFT_COMPILATION_MODE = wholemodule; SWIFT_OPTIMIZATION_LEVEL = "-O"; }; name = Release; }; C8ED50C4290C204800153A3B /* Debug */ = { isa = XCBuildConfiguration; buildSettings = { ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_ENTITLEMENTS = "Single Line Diffusion/Single_Line_Diffusion.entitlements"; CODE_SIGN_STYLE = Automatic; COMBINE_HIDPI_IMAGES = YES; CURRENT_PROJECT_VERSION = 1; DEVELOPMENT_ASSET_PATHS = "\"Single Line Diffusion/Preview Content\""; DEVELOPMENT_TEAM = UDGLB23X37; ENABLE_HARDENED_RUNTIME = YES; ENABLE_PREVIEWS = YES; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_KEY_NSHumanReadableCopyright = ""; LD_RUNPATH_SEARCH_PATHS = ( "$(inherited)", "@executable_path/../Frameworks", ); MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = "app.otato.Single-Line-Diffusion"; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; }; name = Debug; }; C8ED50C5290C204800153A3B /* Release */ = { isa = XCBuildConfiguration; buildSettings = { ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon; ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor; CODE_SIGN_ENTITLEMENTS = "Single Line Diffusion/Single_Line_Diffusion.entitlements"; CODE_SIGN_STYLE = Automatic; COMBINE_HIDPI_IMAGES = YES; CURRENT_PROJECT_VERSION = 1; DEVELOPMENT_ASSET_PATHS = "\"Single Line Diffusion/Preview Content\""; DEVELOPMENT_TEAM = UDGLB23X37; ENABLE_HARDENED_RUNTIME = YES; ENABLE_PREVIEWS = YES; GENERATE_INFOPLIST_FILE = YES; INFOPLIST_KEY_NSHumanReadableCopyright = ""; LD_RUNPATH_SEARCH_PATHS = ( "$(inherited)", "@executable_path/../Frameworks", ); MARKETING_VERSION = 1.0; PRODUCT_BUNDLE_IDENTIFIER = "app.otato.Single-Line-Diffusion"; PRODUCT_NAME = "$(TARGET_NAME)"; SWIFT_EMIT_LOC_STRINGS = YES; SWIFT_VERSION = 5.0; }; name = Release; }; /* End XCBuildConfiguration section */ /* Begin XCConfigurationList section */ C8ED50AF290C204800153A3B /* Build configuration list for PBXProject "Single Line Diffusion" */ = { isa = XCConfigurationList; buildConfigurations = ( C8ED50C1290C204800153A3B /* Debug */, C8ED50C2290C204800153A3B /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; C8ED50C3290C204800153A3B /* Build configuration list for PBXNativeTarget "Single Line Diffusion" */ = { isa = XCConfigurationList; buildConfigurations = ( C8ED50C4290C204800153A3B /* Debug */, C8ED50C5290C204800153A3B /* Release */, ); defaultConfigurationIsVisible = 0; defaultConfigurationName = Release; }; /* End XCConfigurationList section */ /* Begin XCSwiftPackageProductDependency section */ C8ED50C8290C207200153A3B /* MapleDiffusion */ = { isa = XCSwiftPackageProductDependency; productName = MapleDiffusion; }; /* End XCSwiftPackageProductDependency section */ }; rootObject = C8ED50AC290C204800153A3B /* Project object */; } ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion.xcodeproj/project.xcworkspace/contents.xcworkspacedata ================================================ ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion.xcodeproj/project.xcworkspace/xcshareddata/IDEWorkspaceChecks.plist ================================================ IDEDidComputeMac32BitWarning ================================================ FILE: Examples/Single Line Diffusion/Single Line Diffusion.xcodeproj/project.xcworkspace/xcshareddata/swiftpm/Package.resolved ================================================ { "pins" : [ { "identity" : "zipfoundation", "kind" : "remoteSourceControl", "location" : "https://github.com/weichsel/ZIPFoundation.git", "state" : { "revision" : "f6a22e7da26314b38bf9befce34ae8e4b2543090", "version" : "0.9.15" } } ], "version" : 2 } ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2022 Ollin Boer Bohan Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: Package.swift ================================================ // swift-tools-version: 5.7 // The swift-tools-version declares the minimum version of Swift required to build this package. import PackageDescription let package = Package( name: "MapleDiffusion", platforms: [ .macOS("12.3"), .iOS("15.4") ], products: [ // Products define the executables and libraries a package produces, and make them visible to other packages. .library( name: "MapleDiffusion", targets: ["MapleDiffusion"]), ], dependencies: [ // Dependencies declare other packages that this package depends on. // .package(url: /* package url */, from: "1.0.0"), .package(url: "https://github.com/weichsel/ZIPFoundation.git", .upToNextMajor(from: "0.9.0")) ], targets: [ // Targets are the basic building blocks of a package. A target can define a module or a test suite. // Targets can depend on other targets in this package, and on products in packages this package depends on. .target( name: "MapleDiffusion", dependencies: ["ZIPFoundation"]), .testTarget( name: "MapleDiffusionTests", dependencies: ["MapleDiffusion", "ZIPFoundation"]), ] ) ================================================ FILE: README.md ================================================ # Native Diffusion Swift Package [Join us on Discord](https://discord.gg/XNsw7x667a) Native Diffusion runs Stable Diffusion models **locally** on macOS / iOS devices, in Swift, using the MPSGraph framework (not Python). This is the Swift Package Manager wrapper of [Maple Diffusion](https://github.com/madebyollin/maple-diffusion). It adds image-to-image, Swift Package Manager package, and convenient ways to use the code, like Combine publishers and async/await versions. It also supports downloading weights from any local or remote URL, including the app bundle itself. Would not be possible without * [@madebyollin](https://github.com/madebyollin/) who wrote the Metal Performance Shader Graph pipeline * [@GuiyeC](https://github.com/GuiyeC) who wrote the image-to-image implementation # Features Get started in 10 minutes * Extremely simple API. Generate an image in one line of code. Make it do what you want * Flexible API. Pass in prompt, guidance scale, steps, seed, and an image. * One-off conversion script from .ckpt to Native Diffusion's own memory-optimized format * Supports Dreambooth models. Built to be fun to code with * Supports async/await, Combine publisher and classic callbacks. * Optimized for SwiftUI, but can be used in any kind of project, including command line, UIKit, or AppKit Built for end-user speed and great user experience * 100% native. No Python, no environments, your user don't need to install anything first. * Model download built in. Point it to a web address with the model files in a zip archive. The package will download and install the model for later use. * As fast or faster than a server in the cloud on newer Macs Commercial use allowed * MIT Licensed (code). We'd love attribution, but it's not needed legally. * Generated images are licensed under the [CreativeML Open RAIL-M](https://github.com/CompVis/stable-diffusion/blob/main/LICENSE) license, meaning you can use the images for virtually anything, including commercial use. # Usage ## One-line diffusion In its simplest form it's as simple as one line: ```swift let image = try? await Diffusion.generate(localOrRemote: modelUrl, prompt: "cat astronaut") ``` You can give it a local or remote URL or both. If remote, the downloaded weights are saved for later. The single line version is currently limited in terms of parameters. See `examples/SingleLineDiffusion` for a working example. ## As an observable object Let's add some UI. Here's an entire working image generator app in a single SwiftUI view: ![GIF demo](https://github.com/mortenjust/maple-diffusion/blob/main/Examples/Demos/simple-diffusion.gif) ```swift struct ContentView: View { // 1 @StateObject var sd = Diffusion() @State var prompt = "" @State var image : CGImage? @State var imagePublisher = Diffusion.placeholderPublisher @State var progress : Double = 0 var anyProgress : Double { sd.loadingProgress < 1 ? sd.loadingProgress : progress } var body: some View { VStack { DiffusionImage(image: $image, progress: $progress) Spacer() TextField("Prompt", text: $prompt) // 3 .onSubmit { self.imagePublisher = sd.generate(prompt: prompt) } .disabled(!sd.isModelReady) ProgressView(value: anyProgress) .opacity(anyProgress == 1 || anyProgress == 0 ? 0 : 1) } .task { // 2 let path = URL(string: "http://localhost:8080/Diffusion.zip")! try! await sd.prepModels(remoteURL: path) } // 4 .onReceive(imagePublisher) { r in self.image = r.image self.progress = r.progress } .frame(minWidth: 200, minHeight: 200) } } ``` Here's what it does 1. Instantiate a `Diffusion` object 2. Prepare the models, download if needed 3. Submit a prompt for generation 4. Receive updates during generation See `examples/SimpleDiffusion` for a working example. ## `DiffusionImage` An optional SwiftUI view that is specialized for diffusion: - Receives drag and drop of an image from e.g. Finder and sends it back to you via a binding (macOS) - Automatically resizes the image to 512x512 (macOS) - Lets users drag the image to Finder or other apps (macOS) - Blurs the internmediate image while generating (macOS and iOS) # Install Add `https://github.com/mortenjust/native-diffusion` in the ["Swift Package Manager" tab in Xcode](https://developer.apple.com/documentation/xcode/adding_package_dependencies_to_your_app) # Preparing the weights Native Diffusion splits the weights into a binary format that is different from the typical CKPT format. It uses many small files which it then (optionally) swaps in and out of memory, enabling it to run on both macOS and iOS. You can use the converter script in the package to convert your own CKPT file. ## Option 1: Pre-converted Standard Stable Diffusion v1.5 By downloading this zip file, you accept the [creative license from StabilityAI](https://github.com/CompVis/stable-diffusion/blob/main/LICENSE). [Download ZIP](https://drive.google.com/file/d/1EU_qxlF-p6XrxsoACvcI2CAmu45NphAC/view?usp=share_link). Please don't use this URL in your software. We'll get back to what to do with it in a second. ## Option 2: Preparing your own `ckpt` file
If you want to use your own CKPT file (like a Dreambooth fine-tuning), you can convert it into Maple Diffusion format 1. Download a Stable Diffusion model checkpoint to a folder, e.g. `~/Downloads/sd` ([`sd-v1-5.ckpt`](https://huggingface.co/runwayml/stable-diffusion-v1-5), or some derivation thereof) 2. Setup & install Python with PyTorch, if you haven't already. ``` # Grab the converter script cd ~/Downloads/sd curl https://raw.githubusercontent.com/mortenjust/maple-diffusion/main/Converter%20Script/maple-convert.py > maple-convert.py # may need to install conda first https://github.com/conda-forge/miniforge#homebrew conda deactivate conda remove -n native-diffusion --all conda create -n native-diffusion python=3.10 conda activate native-diffusion pip install torch typing_extensions numpy Pillow requests pytorch_lightning ./native-convert.py ~/Downloads/sd-v1-4.ckpt ``` The script will create a new folder called `bins`. We'll get back to what to do with it in a second.
# FAQ ## Can I use a Dreambooth model? Yes. Just copy the `alpha*` files from the standard conversion. This repo will include these files in the future. See [this issue](https://github.com/madebyollin/maple-diffusion/issues/22). ## Does it support image to image prompting? Yes. Simply pass in an `initImage` to your `SampleInput` when generating. ## It crashes You may need to regenerate the model files with the python script in the repo. This happens if you converted your ckpt model file before we added image2image. ## Can I contribute? What's next? Yes! A rough roadmap: - [ ] Stable Diffusion 2.0: - new larger output images, upscaling, depth-to-image - [ ] Add in-painting and out-painting - [ ] Generate other sizes and aspects than 512x512 - [ ] Upscaling - [ ] Dreambooth training on-device - [x] Tighten up code quality overall. Most is proof of concept. - [x] Add image-to-image See Issues for smaller contributions. If you're making changes to the MPSGraph part of the codebase, consider making your contributions to the single-file repo and then integrate the changes in the wrapped file in this repo. ## How fast is it? On my MacBook Pro M1 Max, I get ~0.3s/step, which is significantly faster than any Python/PyTorch/Tensorflow installation I've tried. On an iPhone it should take a minute or two. To attain usable performance without tripping over iOS's 4GB memory limit, Native Diffusion relies internally on FP16 (NHWC) tensors, operator fusion from MPSGraph, and a truly pitiable degree of swapping models to device storage. ## Does it support Stable Diffusion 2.0? Not yet. Would love some help on this. See above. ## I have a question, comment or suggestion Feel free to post an issue! ================================================ FILE: Sources/MapleDiffusion/Diffusion.swift ================================================ import Foundation import Combine import CoreGraphics import CoreImage /** This is the package's wrapper for the `MapleDiffusion` class. It adds a few convenient ways to use @madebyollin's code. */ public class Diffusion : ObservableObject { /// Current state of the models. Only supports states for loading currently, updated via `initModels` public var state = PassthroughSubject() @Published public var isModelReady = false @Published public var loadingProgress : Double = 0.0 var modelIsCold : Bool { print("coldness: ", isModelReady, loadingProgress) return !isModelReady && loadingProgress == 0 } var mapleDiffusion : MapleDiffusion! // Local, offline Stable Diffusion generation in Swift, no Python. Download + init + generate = 1 line of code. public static var shared = Diffusion() // TODO: Move to + generate // TODO: Add steps, guiadance etc as optional params public static func generate(localOrRemote modelURL: URL, prompt: String) async throws -> CGImage? { if shared.modelIsCold { if modelURL.isFileURL { try await shared.prepModels(localUrl: modelURL) } else { try await shared.prepModels(remoteURL: modelURL) } } return await shared.generate(input: SampleInput(prompt: prompt)) } private var saveMemory = false public init(saveMemoryButBeSlower: Bool = false) { self.saveMemory = saveMemoryButBeSlower state.send(.notStarted) } /// Empty publisher for convenience in SwiftUI (since you can't listen to a nil) public static var placeholderPublisher : AnyPublisher { PassthroughSubject().eraseToAnyPublisher() } // Prep models /// Tuple type for easier access to the progress while also getting the full state. public typealias LoaderUpdate = (progress: Double, state: GeneratorState) /// Init models and update publishers. Run this off the main actor. /// TODO: 3 overloads: local only, remote only, both - build into arguments what they do var combinedProgress : Double = 0 public func prepModels(localUrl: URL, progress:((Double)->Void)? = nil) async throws { let fetcher = ModelFetcher(local: localUrl) try await initModels(fetcher: fetcher, progress: progress) } public func prepModels(remoteURL: URL, progress:((Double)->Void)? = nil) async throws { let fetcher = ModelFetcher(remote: remoteURL) try await initModels(fetcher: fetcher, progress: progress) } public func prepModels(localUrl: URL, remoteUrl: URL, progress:((Double)->Void)? = nil) async throws { let fetcher = ModelFetcher(local: localUrl, remote: remoteUrl) try await initModels(fetcher: fetcher, progress: progress) } // Init Models private func initModels( fetcher: ModelFetcher, progress: ((Double)->Void)? ) async throws { let combinedSteps : Double = 2 var combinedProgress : Double = 0 await MainActor.run { self.loadingProgress = 0.05 } /// 1. Fetch the model let modelLocation: URL = try await fetcher.fetch { p in combinedProgress = (p/combinedSteps) self.updateLoadingProgress(progress: combinedProgress, message: "Fetching models") } /// 2. instantiate MD, which has light side effects self.mapleDiffusion = MapleDiffusion(modelLocation: modelLocation, saveMemoryButBeSlower: saveMemory) // let earlierProgress = combinedProgress /// 3. Initialize models on a background thread // try await initModels() { p in // combinedProgress = (p/combinedSteps) + earlierProgress // self.updateLoadingProgress(progress: combinedProgress, message: "Loading models") // progress?(combinedProgress) // } /// 4. Done. Set published main status to true await MainActor.run { print("Model is ready") self.state.send(.ready) self.loadingProgress = 1 self.isModelReady = true } } private func updateLoadingProgress(progress: Double, message:String) { Task { await MainActor.run { self.state.send(GeneratorState.modelIsLoading(progress: progress, message: message)) self.loadingProgress = progress } } } } ================================================ FILE: Sources/MapleDiffusion/MPS/BPETokenizer.swift ================================================ // // BPETokenizer.swift // // // Created by Guillermo Cique Fernández on 9/11/22. // import Foundation import MetalPerformanceShadersGraph class BPETokenizer { // why didn't they just byte-encode func whitespaceClean(s: String) -> String { return s.components(separatedBy: .whitespacesAndNewlines) .filter { !$0.isEmpty } .joined(separator: " ") .trimmingCharacters(in: .whitespacesAndNewlines) } func getPairs(s: [String]) -> Set { return Set((1.."}) // `var vocabFileURL = modelFolder.appendingPathComponent("bpe_simple_vocab_16e6").appendingPathExtension("txt")` // let vocabFile = try! String(contentsOf: Bundle.main.url(forResource: "bins/bpe_simple_vocab_16e6", withExtension: "txt")!) let vocabFile = try! String( contentsOf: modelLocation .appendingPathComponent("bpe_simple_vocab_16e6") .appendingPathExtension("txt") ) for (i, m) in vocabFile.split(separator: "\n")[1..<48_895].enumerated() { ranks[String(m)] = i vocabList.append(m.split(separator: " ").joined(separator: "")) } vocab = vocabList.enumerated().reduce(into: [:], {$0[$1.element] = $1.offset}) } func encodeToken(s: String) -> [Int] { let token = String(s.utf8.map{bytesToUnicode[Int($0)]!}) var word = token[.."] var pairs = getPairs(s: Array(word)) var mergedWordTokens = [token + ""] var count = 0 if (!pairs.isEmpty) { while (true) { count += 1 assert(count < 8192, "encodeToken is trapped in a token factory for input \(s)") let highestRankedBigram = pairs.min(by: {ranks[$0, default: Int.max] < ranks[$1, default: Int.max]})! if (ranks[highestRankedBigram] == nil) { break } let fs = highestRankedBigram.split(separator: " ") let (first, second) = (String(fs[0]), String(fs[1])) var (newWord, i) = ([String](), 0) while (i < word.count) { let j = word[i.. [Int] { let ns = NSString(string: whitespaceClean(s: s.lowercased())) var bpe: [Int] = [] for match in pat.matches(in: String(ns), range: NSRange(location: 0, length: ns.length)) { bpe.append(contentsOf: encodeToken(s: ns.substring(with: match.range))) } if (bpe.count > 75) { print("Prompt of \(bpe.count) bpe tokens will be truncated: \(s)") } return [49406] + bpe[.. MPSGraphTensorData { return graph.run( with: queue, feeds: [input: xIn], targetTensors: [output], targetOperations: nil )[output]! } } extension MPSGraph { func makeDecoder(at folder: URL, xIn: MPSGraphTensor) -> MPSGraphTensor { var x = xIn let name = "first_stage_model.decoder" x = multiplication(x, constant(1 / 0.18215, dataType: MPSDataType.float16), name: "rescale") x = makeConv(at: folder, xIn: x, name: "first_stage_model.post_quant_conv", outChannels: 4, khw: 1) x = makeConv(at: folder, xIn: x, name: name + ".conv_in", outChannels: 512, khw: 3) // middle x = makeCoderResBlock(at: folder, xIn: x, name: name + ".mid.block_1", outChannels: 512) x = makeCoderAttention(at: folder, xIn: x, name: name + ".mid.attn_1") x = makeCoderResBlock(at: folder, xIn: x, name: name + ".mid.block_2", outChannels: 512) // block 3 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.3.block.0", outChannels: 512) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.3.block.1", outChannels: 512) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.3.block.2", outChannels: 512) x = upsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + ".up.3.upsample.conv", outChannels: 512, khw: 3) // block 2 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.2.block.0", outChannels: 512) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.2.block.1", outChannels: 512) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.2.block.2", outChannels: 512) x = upsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + ".up.2.upsample.conv", outChannels: 512, khw: 3) // block 1 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.1.block.0", outChannels: 256) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.1.block.1", outChannels: 256) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.1.block.2", outChannels: 256) x = upsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + ".up.1.upsample.conv", outChannels: 256, khw: 3) // block 0 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.0.block.0", outChannels: 128) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.0.block.1", outChannels: 128) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".up.0.block.2", outChannels: 128) x = makeGroupNormSwish(at: folder, xIn: x, name: name + ".norm_out") x = makeConv(at: folder, xIn: x, name: name + ".conv_out", outChannels: 3, khw: 3) x = addition(x, constant(1.0, dataType: MPSDataType.float16), name: nil) x = multiplication(x, constant(0.5, dataType: MPSDataType.float16), name: nil) return makeByteConverter(xIn: x) } } ================================================ FILE: Sources/MapleDiffusion/MPS/Coder/Encoder.swift ================================================ // // Encoder.swift // // // Created by Guillermo Cique Fernández on 9/11/22. // import Foundation import MetalPerformanceShadersGraph class Encoder { let device: MPSGraphDevice private let graph: MPSGraph private let encoderIn: MPSGraphTensor! private let encoderOut: MPSGraphTensor! private let noise: MPSGraphTensor! private let gaussianOut: MPSGraphTensor! private let scaled: MPSGraphTensor! private let stochasticEncode: MPSGraphTensor! private let stepIn: MPSGraphTensor! private let timestepsIn: MPSGraphTensor! init(synchronize: Bool, modelLocation: URL, device: MPSGraphDevice, inputShape: [NSNumber], outputShape: [NSNumber], timestepsShape: [NSNumber], seed: Int ) { self.device = device graph = MPSGraph(synchronize: synchronize) encoderIn = graph.placeholder(shape: inputShape, dataType: MPSDataType.uInt8, name: nil) encoderOut = graph.makeEncoder(at: modelLocation, xIn: encoderIn) noise = graph.randomTensor( withShape: outputShape, descriptor: MPSGraphRandomOpDescriptor(distribution: .normal, dataType: .float16)!, seed: seed, name: nil ) gaussianOut = graph.diagonalGaussianDistribution(encoderOut, noise: noise) scaled = graph.multiplication(gaussianOut, graph.constant(0.18215, dataType: MPSDataType.float16), name: "rescale") stepIn = graph.placeholder(shape: [1], dataType: MPSDataType.int32, name: nil) timestepsIn = graph.placeholder(shape: timestepsShape, dataType: MPSDataType.int32, name: nil) stochasticEncode = graph.stochasticEncode(at: modelLocation, stepIn: stepIn, timestepsIn: timestepsIn, imageIn: scaled, noiseIn: noise) } func run(with queue: MTLCommandQueue, image: MPSGraphTensorData, step: Int, timesteps: MPSGraphTensorData) -> MPSGraphTensorData { let stepData = step.tensorData(device: device) return graph.run( with: queue, feeds: [ encoderIn: image, stepIn: stepData, timestepsIn: timesteps ], targetTensors: [ noise, encoderOut, gaussianOut, scaled, stochasticEncode ], targetOperations: nil )[stochasticEncode]! } } extension MPSGraph { func makeEncoder(at folder: URL, xIn: MPSGraphTensor) -> MPSGraphTensor { var x = xIn // Split into RBGA let xParts = split(x, numSplits: 4, axis: 2, name: nil) // Drop alpha channel x = concatTensors(xParts.dropLast(), dimension: 2, name: nil) x = cast(x, to: .float16, name: nil) x = division(x, constant(255.0, shape: [1], dataType: .float16), name: nil) x = expandDims(x, axis: 0, name: nil) x = multiplication(x, constant(2.0, shape: [1], dataType: .float16), name: nil) x = subtraction(x, constant(1.0, shape: [1], dataType: .float16), name: nil) let name = "first_stage_model.encoder" x = makeConv(at: folder, xIn: x, name: name + ".conv_in", outChannels: 128, khw: 3) // block 0 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.0.block.0", outChannels: 128) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.0.block.1", outChannels: 128) x = downsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + ".down.0.downsample.conv", outChannels: 128, khw: 3) // block 1 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.1.block.0", outChannels: 256) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.1.block.1", outChannels: 256) x = downsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + ".down.1.downsample.conv", outChannels: 256, khw: 3) // block 2 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.2.block.0", outChannels: 512) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.2.block.1", outChannels: 512) x = downsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + ".down.2.downsample.conv", outChannels: 512, khw: 3) // block 3 x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.3.block.0", outChannels: 512) x = makeCoderResBlock(at: folder, xIn: x, name: name + ".down.3.block.1", outChannels: 512) // middle x = makeCoderResBlock(at: folder, xIn: x, name: name + ".mid.block_1", outChannels: 512) x = makeCoderAttention(at: folder, xIn: x, name: name + ".mid.attn_1") x = makeCoderResBlock(at: folder, xIn: x, name: name + ".mid.block_2", outChannels: 512) x = makeGroupNormSwish(at: folder, xIn: x, name: name + ".norm_out") x = makeConv(at: folder, xIn: x, name: name + ".conv_out", outChannels: 8, khw: 3) return makeConv(at: folder, xIn: x, name: "first_stage_model.quant_conv", outChannels: 8, khw: 1) } } ================================================ FILE: Sources/MapleDiffusion/MPS/Coder/MPSGraph+Coder.swift ================================================ // // MPSGraph+Coder.swift // // // Created by Guillermo Cique Fernández on 9/11/22. // import Foundation import MetalPerformanceShadersGraph extension MPSGraph { func makeCoderResBlock(at folder: URL, xIn: MPSGraphTensor, name: String, outChannels: NSNumber) -> MPSGraphTensor { var x = xIn x = makeGroupNormSwish(at: folder, xIn: x, name: name + ".norm1") x = makeConv(at: folder, xIn: x, name: name + ".conv1", outChannels: outChannels, khw: 3) x = makeGroupNormSwish(at: folder, xIn: x, name: name + ".norm2") x = makeConv(at: folder, xIn: x, name: name + ".conv2", outChannels: outChannels, khw: 3) if (xIn.shape![3] != outChannels) { let ninShortcut = makeConv(at: folder, xIn: xIn, name: name + ".nin_shortcut", outChannels: outChannels, khw: 1) return addition(x, ninShortcut, name: "skip") } return addition(x, xIn, name: "skip") } func makeCoderAttention(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var x = makeGroupNorm(at: folder, xIn: xIn, name: name + ".norm") let c = x.shape![3] x = reshape(x, shape: [x.shape![0], NSNumber(value:x.shape![1].intValue * x.shape![2].intValue), c], name: nil) let q = makeLinear(at: folder, xIn: x, name: name + ".q", outChannels: c, bias: false) var k = makeLinear(at: folder, xIn: x, name: name + ".k", outChannels: c, bias: false) k = multiplication(k, constant(1.0 / sqrt(c.doubleValue), dataType: MPSDataType.float16), name: nil) k = transposeTensor(k, dimension: 1, withDimension: 2, name: nil) let v = makeLinear(at: folder, xIn: x, name: name + ".v", outChannels: c, bias: false) var att = matrixMultiplication(primary: q, secondary: k, name: nil) att = softMax(with: att, axis: 2, name: nil) att = matrixMultiplication(primary: att, secondary: v, name: nil) x = makeLinear(at: folder, xIn: att, name: name + ".proj_out", outChannels: c) x = reshape(x, shape: xIn.shape!, name: nil) return addition(x, xIn, name: nil) } } ================================================ FILE: Sources/MapleDiffusion/MPS/Diffuser.swift ================================================ // // Diffuser.swift // // // Created by Guillermo Cique Fernández on 14/11/22. // import Foundation import MetalPerformanceShadersGraph class Diffuser { private let device: MPSGraphDevice private let graph: MPSGraph private let xIn: MPSGraphTensor private let etaUncondIn: MPSGraphTensor private let etaCondIn: MPSGraphTensor private let timestepIn: MPSGraphTensor private let timestepSizeIn: MPSGraphTensor private let guidanceScaleIn: MPSGraphTensor private let out: MPSGraphTensor private let auxOut: MPSGraphTensor init(synchronize: Bool, modelLocation: URL, device: MPSGraphDevice, shape: [NSNumber]) { self.device = device graph = MPSGraph(synchronize: synchronize) xIn = graph.placeholder(shape: shape, dataType: MPSDataType.float16, name: nil) etaUncondIn = graph.placeholder(shape: shape, dataType: MPSDataType.float16, name: nil) etaCondIn = graph.placeholder(shape: shape, dataType: MPSDataType.float16, name: nil) timestepIn = graph.placeholder(shape: [1], dataType: MPSDataType.int32, name: nil) timestepSizeIn = graph.placeholder(shape: [1], dataType: MPSDataType.int32, name: nil) guidanceScaleIn = graph.placeholder(shape: [1], dataType: MPSDataType.float32, name: nil) out = graph.makeDiffusionStep( at: modelLocation, xIn: xIn, etaUncondIn: etaUncondIn, etaCondIn: etaCondIn, timestepIn: timestepIn, timestepSizeIn: timestepSizeIn, guidanceScaleIn: graph.cast(guidanceScaleIn, to: MPSDataType.float16, name: "this string must not be the empty string") ) auxOut = graph.makeAuxUpsampler(at: modelLocation, xIn: out) } func run( with queue: MTLCommandQueue, latent: MPSGraphTensorData, timestep: Int, timestepSize: Int, etaUncond: MPSGraphTensorData, etaCond: MPSGraphTensorData, guidanceScale: MPSGraphTensorData ) -> (MPSGraphTensorData, MPSGraphTensorData?) { let timestepData = timestep.tensorData(device: device) let timestepSizeData = timestepSize.tensorData(device: device) let outputs = graph.run( with: queue, feeds: [ xIn: latent, etaUncondIn: etaUncond, etaCondIn: etaCond, timestepIn: timestepData, timestepSizeIn: timestepSizeData, guidanceScaleIn: guidanceScale ], targetTensors: [out, auxOut], targetOperations: nil ) return (outputs[out]!, outputs[auxOut]) } } fileprivate extension MPSGraph { func makeDiffusionStep( at folder: URL, xIn: MPSGraphTensor, etaUncondIn: MPSGraphTensor, etaCondIn: MPSGraphTensor, timestepIn: MPSGraphTensor, timestepSizeIn: MPSGraphTensor, guidanceScaleIn: MPSGraphTensor ) -> MPSGraphTensor { // superconditioning var deltaCond = multiplication(subtraction(etaCondIn, etaUncondIn, name: nil), guidanceScaleIn, name: nil) deltaCond = tanh(with: deltaCond, name: nil) // NOTE: normal SD doesn't clamp here iirc let eta = addition(etaUncondIn, deltaCond, name: nil) // scheduler conditioning let alphasCumprod = loadConstant(at: folder, name: "alphas_cumprod", shape: [1000]) let alphaIn = gatherAlongAxis(0, updates: alphasCumprod, indices: timestepIn, name: nil) let prevTimestep = maximum( constant(0, dataType: MPSDataType.int32), subtraction(timestepIn, timestepSizeIn, name: nil), name: nil ) let alphaPrevIn = gatherAlongAxis(0, updates: alphasCumprod, indices: prevTimestep, name: nil) // scheduler step let deltaX0 = multiplication(squareRootOfOneMinus(alphaIn), eta, name: nil) let predX0Unscaled = subtraction(xIn, deltaX0, name: nil) let predX0 = division(predX0Unscaled, squareRoot(with: alphaIn, name: nil), name: nil) let dirX = multiplication(squareRootOfOneMinus(alphaPrevIn), eta, name: nil) let xPrevBase = multiplication(squareRoot(with: alphaPrevIn, name: nil), predX0, name:nil) return addition(xPrevBase, dirX, name: nil) } func makeAuxUpsampler(at folder: URL, xIn: MPSGraphTensor) -> MPSGraphTensor { var x = xIn x = makeConv(at: folder, xIn: xIn, name: "aux_output_conv", outChannels: 3, khw: 1) x = upsampleNearest(xIn: x, scaleFactor: 8) return makeByteConverter(xIn: x) } } ================================================ FILE: Sources/MapleDiffusion/MPS/MPSGraph+Transformer.swift ================================================ // // MPSGraph+Transformer.swift // // // Created by Guillermo Cique Fernández on 9/11/22. // import Foundation import MetalPerformanceShadersGraph extension MPSGraph { func makeSpatialTransformerBlock(at folder: URL, xIn: MPSGraphTensor, name: String, contextIn: MPSGraphTensor, saveMemory: Bool) -> MPSGraphTensor { let n, h, w, c: NSNumber (n, h, w, c) = (xIn.shape![0], xIn.shape![1], xIn.shape![2], xIn.shape![3]) var x = xIn x = makeGroupNorm(at: folder, xIn: x, name: name + ".norm") x = makeConv(at: folder, xIn: x, name: name + ".proj_in", outChannels: c, khw: 1) x = reshape(x, shape: [n, (h.intValue * w.intValue) as NSNumber, c], name: nil) x = makeBasicTransformerBlock(at: folder, xIn: x, name: name + ".transformer_blocks.0", contextIn: contextIn, saveMemory: saveMemory) x = reshape(x, shape: [n, h, w, c], name: nil) x = makeConv(at: folder, xIn: x, name: name + ".proj_out", outChannels: c, khw: 1) return addition(x, xIn, name: nil) } fileprivate func makeBasicTransformerBlock(at folder: URL, xIn: MPSGraphTensor, name: String, contextIn: MPSGraphTensor, saveMemory: Bool) -> MPSGraphTensor { var x = xIn var attn1 = makeLayerNorm(at: folder, xIn: x, name: name + ".norm1") attn1 = makeCrossAttention(at: folder, xIn: attn1, name: name + ".attn1", context: nil, saveMemory: saveMemory) x = addition(attn1, x, name: nil) var attn2 = makeLayerNorm(at: folder, xIn: x, name: name + ".norm2") attn2 = makeCrossAttention(at: folder, xIn: attn2, name: name + ".attn2", context: contextIn, saveMemory: saveMemory) x = addition(attn2, x, name: nil) var ff = makeLayerNorm(at: folder, xIn: x, name: name + ".norm3") ff = makeFeedForward(at: folder, xIn: ff, name: name + ".ff.net") return addition(ff, x, name: nil) } fileprivate func makeFeedForward(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { assert(xIn.shape!.count == 3) let dim = xIn.shape![2] let dimMult = dim.intValue * 4 let dimProj = NSNumber(value: dimMult * 2) let proj = makeLinear(at: folder, xIn: xIn, name: name + ".0.proj", outChannels: dimProj) var x = sliceTensor(proj, dimension: 2, start: 0, length: dimMult, name: nil) var gate = sliceTensor(proj, dimension: 2, start: dimMult, length: dimMult, name: nil) gate = gelu(gate) x = multiplication(x, gate, name: nil) return makeLinear(at: folder, xIn: x, name: name + ".2", outChannels: dim) } fileprivate func makeCrossAttention(at folder: URL, xIn: MPSGraphTensor, name: String, context: MPSGraphTensor?, saveMemory: Bool) -> MPSGraphTensor { let c = xIn.shape![2] let (nHeads, dHead) = (NSNumber(8), NSNumber(value: c.intValue / 8)) var q = makeLinear(at: folder, xIn: xIn, name: name + ".to_q", outChannels: c, bias: false) let context = context ?? xIn var k = makeLinear(at: folder, xIn: context, name: name + ".to_k", outChannels: c, bias: false) var v = makeLinear(at: folder, xIn: context, name: name + ".to_v", outChannels: c, bias: false) let n = xIn.shape![0] let hw = xIn.shape![1] let t = context.shape![1] q = reshape(q, shape: [n, hw, nHeads, dHead], name: nil) k = reshape(k, shape: [n, t, nHeads, dHead], name: nil) v = reshape(v, shape: [n, t, nHeads, dHead], name: nil) q = transposeTensor(q, dimension: 1, withDimension: 2, name: nil) k = transposeTensor(k, dimension: 1, withDimension: 2, name: nil) k = transposeTensor(k, dimension: 2, withDimension: 3, name: nil) k = multiplication(k, constant(1.0 / sqrt(dHead.doubleValue), dataType: MPSDataType.float16), name: nil) v = transposeTensor(v, dimension: 1, withDimension: 2, name: nil) var att: MPSGraphTensor if (saveMemory) { // MEM-HACK - silly graph seems to use less peak memory var attRes = [MPSGraphTensor]() let sliceSize = 1 for i in 0.. MPSGraphTensor { let numels = shape.reduce(into: 1, { accumulator, value in accumulator *= value.intValue }) let fileUrl: URL = folder.appendingPathComponent(name + (fp32 ? "_fp32" : "")).appendingPathExtension("bin") let data: Data = try! Data(contentsOf: fileUrl, options: Data.ReadingOptions.alwaysMapped) let expectedCount = numels * (fp32 ? 4 : 2) assert(data.count == expectedCount, "Mismatch between byte count of data \(data.count) and expected size \(expectedCount) for \(numels) els in \(fileUrl)") return constant(data, shape: shape, dataType: fp32 ? MPSDataType.float32 : MPSDataType.float16) } func makeConv(at folder: URL, xIn: MPSGraphTensor, name: String, outChannels: NSNumber, khw: NSNumber, stride: Int = 1, bias: Bool = true) -> MPSGraphTensor { let w = loadConstant(at: folder, name: name + ".weight", shape: [outChannels, xIn.shape![3], khw, khw]) let p: Int = khw.intValue / 2; let convDesc = MPSGraphConvolution2DOpDescriptor( strideInX: stride, strideInY: stride, dilationRateInX: 1, dilationRateInY: 1, groups: 1, paddingLeft: p, paddingRight: p, paddingTop: p, paddingBottom: p, paddingStyle: MPSGraphPaddingStyle.explicit, dataLayout: MPSGraphTensorNamedDataLayout.NHWC, weightsLayout: MPSGraphTensorNamedDataLayout.OIHW )! let conv = convolution2D(xIn, weights: w, descriptor: convDesc, name: nil) if (bias) { let b = loadConstant(at: folder, name: name + ".bias", shape: [1, 1, 1, outChannels]) return addition(conv, b, name: nil) } return conv } func makeLinear(at folder: URL, xIn: MPSGraphTensor, name: String, outChannels: NSNumber, bias: Bool = true) -> MPSGraphTensor { if (xIn.shape!.count == 2) { var x = reshape(xIn, shape: [xIn.shape![0], 1, 1, xIn.shape![1]], name: nil) x = makeConv(at: folder, xIn: x, name: name, outChannels: outChannels, khw: 1, bias: bias) return reshape(x, shape: [xIn.shape![0], outChannels], name: nil) } var x = reshape(xIn, shape: [xIn.shape![0], 1, xIn.shape![1], xIn.shape![2]], name: nil) x = makeConv(at: folder, xIn: x, name: name, outChannels: outChannels, khw: 1, bias: bias) return reshape(x, shape: [xIn.shape![0], xIn.shape![1], outChannels], name: nil) } func makeLayerNorm(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { assert(xIn.shape!.count == 3, "layernorm requires NTC") let gamma = loadConstant(at: folder, name: name + ".weight", shape: [1, 1, xIn.shape![2]]) let beta = loadConstant(at: folder, name: name + ".bias", shape: [1, 1, xIn.shape![2]]) let mean = mean(of: xIn, axes: [2], name: nil) let variance = variance(of: xIn, axes: [2], name: nil) let x = normalize(xIn, mean: mean, variance: variance, gamma: gamma, beta: beta, epsilon: 1e-5, name: nil) return reshape(x, shape: xIn.shape!, name: nil) } func makeGroupNorm(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var x = xIn if (xIn.shape!.count == 3) { x = expandDims(x, axes: [1], name: nil) } let shape = x.shape! let nGroups: NSNumber = 32 let nGrouped: NSNumber = shape[3].floatValue / nGroups.floatValue as NSNumber let gamma = loadConstant(at: folder, name: name + ".weight", shape: [1, 1, 1, nGroups, nGrouped]) let beta = loadConstant(at: folder, name: name + ".bias", shape: [1, 1, 1, nGroups, nGrouped]) x = reshape(x, shape: [shape[0], shape[1], shape[2], nGroups, nGrouped], name: nil) let mean = mean(of: x, axes: [1, 2, 4], name: nil) let variance = variance(of: x, axes: [1, 2, 4], name: nil) x = normalize(x, mean: mean, variance: variance, gamma: gamma, beta: beta, epsilon: 1e-5, name: nil) return reshape(x, shape: xIn.shape!, name: nil) } func makeGroupNormSwish(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { return swish(makeGroupNorm(at: folder, xIn: xIn, name: name)) } func makeByteConverter(xIn: MPSGraphTensor) -> MPSGraphTensor { var x = xIn x = clamp(x, min: constant(0, shape: [1], dataType: MPSDataType.float16), max: constant(1.0, shape: [1], dataType: MPSDataType.float16), name: nil) x = multiplication(x, constant(255, shape: [1], dataType: MPSDataType.float16), name: nil) x = round(with: x, name: nil) x = cast(x, to: MPSDataType.uInt8, name: "cast to uint8 rgba") let alpha = constant(255, shape: [1, x.shape![1], x.shape![2], 1], dataType: MPSDataType.uInt8) return concatTensors([x, alpha], dimension: 3, name: nil) } func stochasticEncode(at folder: URL, stepIn: MPSGraphTensor, timestepsIn: MPSGraphTensor, imageIn: MPSGraphTensor, noiseIn: MPSGraphTensor) -> MPSGraphTensor { let alphasCumprod = loadConstant(at: folder, name: "alphas_cumprod", shape: [1000]) let alphas = gatherAlongAxis(0, updates: alphasCumprod, indices: timestepsIn, name: nil) let sqrtAlphasCumprod = squareRoot(with: alphas, name: nil) let sqrtOneMinusAlphasCumprod = squareRootOfOneMinus(alphas) let imageAlphas = multiplication(extractIntoTensor(a: sqrtAlphasCumprod, t: stepIn, shape: imageIn.shape!), imageIn, name: nil) let noiseAlphas = multiplication(extractIntoTensor(a: sqrtOneMinusAlphasCumprod, t: stepIn, shape: imageIn.shape!), noiseIn, name: nil) return addition(imageAlphas, noiseAlphas, name: nil) } } // MARK: Operations extension MPSGraph { func swish(_ tensor: MPSGraphTensor) -> MPSGraphTensor { return multiplication(tensor, sigmoid(with: tensor, name: nil), name: nil) } func upsampleNearest(xIn: MPSGraphTensor, scaleFactor: Int = 2) -> MPSGraphTensor { return resize( xIn, size: [ NSNumber(value:xIn.shape![1].intValue * scaleFactor), NSNumber(value:xIn.shape![2].intValue * scaleFactor) ], mode: MPSGraphResizeMode.nearest, centerResult: true, alignCorners: false, layout: MPSGraphTensorNamedDataLayout.NHWC, name: nil ) } func downsampleNearest(xIn: MPSGraphTensor, scaleFactor: Int = 2) -> MPSGraphTensor { return resize( xIn, size: [ NSNumber(value:xIn.shape![1].intValue / scaleFactor), NSNumber(value:xIn.shape![2].intValue / scaleFactor) ], mode: MPSGraphResizeMode.nearest, centerResult: true, alignCorners: false, layout: MPSGraphTensorNamedDataLayout.NHWC, name: nil ) } func squareRootOfOneMinus(_ tensor: MPSGraphTensor) -> MPSGraphTensor { return squareRoot(with: subtraction(constant(1.0, dataType: MPSDataType.float16), tensor, name: nil), name: nil) } // Gaussian Error Linear Units func gelu(_ tensor: MPSGraphTensor) -> MPSGraphTensor { var x = tensor x = multiplication(x, constant(1/sqrt(2), dataType: MPSDataType.float16), name: nil) x = erf(with: x, name: nil) x = addition(x, constant(1, dataType: MPSDataType.float16), name: nil) x = multiplication(x, constant(0.5, dataType: MPSDataType.float16), name: nil) return multiplication(tensor, x, name: nil) } func diagonalGaussianDistribution(_ tensor: MPSGraphTensor, noise: MPSGraphTensor) -> MPSGraphTensor { let chunks = split(tensor, numSplits: 2, axis: 3, name: nil) let mean = chunks[0] let logvar = clamp(chunks[1], min: constant(-30, shape: [1], dataType: MPSDataType.float16), max: constant(20, shape: [1], dataType: MPSDataType.float16), name: nil) let std = exponent(with: multiplication(constant(0.5, shape: [1], dataType: MPSDataType.float16), logvar, name: nil), name: nil) return addition(mean, multiplication(std, noise, name: nil), name: nil) } func extractIntoTensor(a: MPSGraphTensor, t: MPSGraphTensor, shape: [NSNumber]) -> MPSGraphTensor { let out = gatherAlongAxis(-1, updates: a, indices: t, name: nil) return reshape(out, shape: [t.shape!.first!] + [NSNumber](repeating: 1, count: shape.count-1), name: nil) } } ================================================ FILE: Sources/MapleDiffusion/MPS/MapleDiffusion.swift ================================================ import MetalPerformanceShadersGraph import Foundation // Maple Diffusion implements stable diffusion (original v1.4 model) // inference via MPSGraph. iOS has a hard memory limit of 4GB (with // a special entitlement), so this implementation trades off latency // for memory usage in many places (tagged with MEM-HACK) in order to // stay under the limit and minimize probability of oom. /** When updating 1. Paste in the entire new file 2. Replace references to Bundle.main like this `let fileUrl: URL = Bundle.main.url(forResource: "bins/" + name + (fp32 ? "_fp32" : ""), withExtension: ".bin")!` to `let fileUrl: URL = modelFolder.appendingPathComponent(name + (fp32 ? "_fp32" : "")).appendingPathExtension("bin")` and also ``` let vocabFile = try! String( contentsOf: modelFolder .appendingPathComponent("bpe_simple_vocab_16e6") .appendingPathExtension("txt") ) ``` */ // madebyollin's code starts here: class MapleDiffusion { let device: MTLDevice let graphDevice: MPSGraphDevice let commandQueue: MTLCommandQueue let saveMemory: Bool let synchronize: Bool let modelLocation: URL private var _textGuidance: TextGuidance? var textGuidance: TextGuidance { if let _textGuidance { return _textGuidance } let textGuidance = TextGuidance( synchronize: synchronize, modelLocation: modelLocation, device: graphDevice ) _textGuidance = textGuidance return textGuidance } private var _uNet: UNet? var uNet: UNet { if let _uNet { return _uNet } let uNet = UNet( synchronize: synchronize, modelLocation: modelLocation, saveMemory: saveMemory, device: graphDevice, shape: [1, height, width, 4] ) _uNet = uNet return uNet } lazy var diffuser: Diffuser = { Diffuser( synchronize: synchronize, modelLocation: modelLocation, device: graphDevice, shape: [1, height, width, 4] ) }() private var _decoder: Decoder? var decoder: Decoder { if let _decoder { return _decoder } let decoder = Decoder( synchronize: synchronize, modelLocation: modelLocation, device: graphDevice, shape: [1, height, width, 4] ) _decoder = decoder return decoder } var width: NSNumber = 64 var height: NSNumber = 64 public init(modelLocation: URL, saveMemoryButBeSlower: Bool = true) { self.modelLocation = modelLocation saveMemory = saveMemoryButBeSlower device = MTLCreateSystemDefaultDevice()! graphDevice = MPSGraphDevice(mtlDevice: device) commandQueue = device.makeCommandQueue()! synchronize = !device.hasUnifiedMemory } private func runTextGuidance(prompt: String, negativePrompt: String) -> (MPSGraphTensorData, MPSGraphTensorData) { let guidance = textGuidance.run(with: commandQueue, prompt: prompt, negativePrompt: negativePrompt) if saveMemory { // MEM-HACK unload the text guidance to fit the unet _textGuidance = nil } return guidance } private func initLatent(input: SampleInput, scheduler: Scheduler) -> MPSGraphTensorData { if let image = input.initImage, let strength = input.strength { let imageData = MPSGraphTensorData(device: graphDevice, cgImage: image) let timestepsData = scheduler.timestepsData let startStep = Int(Float(input.steps) * strength) let encoder = Encoder( synchronize: synchronize, modelLocation: modelLocation, device: graphDevice, inputShape: imageData.shape, outputShape: [1, height, width, 4], timestepsShape: timestepsData.shape, seed: input.seed ) return encoder.run(with: commandQueue, image: imageData, step: startStep, timesteps: timestepsData) } else { let graph = MPSGraph(synchronize: synchronize) let out = graph.randomTensor( withShape: [1, height, width, 4], descriptor: MPSGraphRandomOpDescriptor(distribution: .normal, dataType: .float16)!, seed: input.seed, name: nil ) return graph.run(with: commandQueue, feeds: [:], targetTensors: [out], targetOperations: nil)[out]! } } private func sample( latent: inout MPSGraphTensorData, input: SampleInput, baseGuidance: MPSGraphTensorData, textGuidance: MPSGraphTensorData, scheduler: Scheduler, completion: @escaping (CGImage?, Float, String) -> () ) { let guidanceScaleData = input.guidanceScale.tensorData(device: graphDevice) let actualTimesteps = scheduler.timesteps(strength: input.strength) for (index, timestep) in actualTimesteps.enumerated() { let tick = CFAbsoluteTimeGetCurrent() let temb = scheduler.run(with: commandQueue, timestep: timestep) let (etaUncond, etaCond) = uNet.run( with: commandQueue, latent: latent, baseGuidance: baseGuidance, textGuidance: textGuidance, temb: temb ) let (newLatent, auxOut) = diffuser.run( with: commandQueue, latent: latent, timestep: timestep, timestepSize: scheduler.timestepSize, etaUncond: etaUncond, etaCond: etaCond, guidanceScale: guidanceScaleData ) latent = newLatent // update ui let tock = CFAbsoluteTimeGetCurrent() let stepRuntime = String(format:"%.2fs", tock - tick) let progressDesc = index == 0 ? "Decoding..." : "Step \(index) / \(actualTimesteps.count) (\(stepRuntime) / step)" let outImage = auxOut?.cgImage let progress = 0.1 + (Float(index) / Float(actualTimesteps.count)) * 0.8 completion(outImage, progress, progressDesc) } if saveMemory { // MEM-HACK: unload the unet to fit the decoder _uNet = nil } } private func runDecoder(latent: MPSGraphTensorData) -> CGImage? { let decodedLatent = decoder.run(with: commandQueue, xIn: latent) if saveMemory { // MEM-HACK unload the decoder _decoder = nil } return decodedLatent.cgImage } public func generate( input: SampleInput, completion: @escaping (CGImage?, Float, String) -> () ) { let mainTick = CFAbsoluteTimeGetCurrent() // 1. String -> Embedding completion(input.initImage, 0, "Tokenizing...") let (baseGuidance, textGuidance) = runTextGuidance(prompt: input.prompt, negativePrompt: input.negativePrompt) // 2. Noise generation completion(input.initImage, 0.05, "Generating noise...") let scheduler = Scheduler(synchronize: synchronize, modelLocation: modelLocation, device: graphDevice, steps: input.steps) var latent = initLatent(input: input, scheduler: scheduler) // 3. Diffusion let startImage: CGImage? if saveMemory { startImage = input.initImage } else { startImage = runDecoder(latent: latent) } completion(startImage, 0.1, "Starting diffusion...") sample( latent: &latent, input: input, baseGuidance: baseGuidance, textGuidance: textGuidance, scheduler: scheduler, completion: completion ) // 4. Decoder let finalImage = runDecoder(latent: latent) completion(finalImage, 1.0, "Cooling down...") let mainTock = CFAbsoluteTimeGetCurrent() let runtime = String(format:"%.2fs", mainTock - mainTick) print("Time", runtime) } } ================================================ FILE: Sources/MapleDiffusion/MPS/Scheduler.swift ================================================ // // Scheduler.swift // // // Created by Guillermo Cique Fernández on 13/11/22. // import Foundation import MetalPerformanceShadersGraph class Scheduler { let count: Int private let timesteps: [Int] let timestepSize: Int var timestepsData: MPSGraphTensorData { let data = timesteps.map { Int32($0) }.withUnsafeBufferPointer { Data(buffer: $0) } return MPSGraphTensorData( device: device, data: data, shape: [NSNumber(value: timesteps.count)], dataType: MPSDataType.int32 ) } private let device: MPSGraphDevice private let graph: MPSGraph private let timestepIn: MPSGraphTensor private let tembOut: MPSGraphTensor init(synchronize: Bool, modelLocation: URL, device: MPSGraphDevice, steps: Int) { self.device = device count = steps timestepSize = 1000 / steps timesteps = Array(stride(from: 1, to: 1000, by: timestepSize)) graph = MPSGraph(synchronize: synchronize) timestepIn = graph.placeholder(shape: [1], dataType: MPSDataType.int32, name: nil) tembOut = graph.makeTimeFeatures(at: modelLocation, tIn: timestepIn) } func timesteps(strength: Float?) -> [Int] { guard let strength else { return timesteps.reversed() } let startStep = Int(Float(count) * strength) return timesteps[0.. MPSGraphTensorData { let timestepData = [Int32(timestep)].withUnsafeBufferPointer { Data(buffer: $0) } let data = MPSGraphTensorData(device: device, data: timestepData, shape: [1], dataType: MPSDataType.int32) return graph.run( with: queue, feeds: [timestepIn: data], targetTensors: [tembOut], targetOperations: nil )[tembOut]! } } extension MPSGraph { func makeTimeFeatures(at folder: URL, tIn: MPSGraphTensor) -> MPSGraphTensor { var temb = cast(tIn, to: MPSDataType.float32, name: "temb") var coeffs = loadConstant(at: folder, name: "temb_coefficients", shape: [160], fp32: true) coeffs = cast(coeffs, to: MPSDataType.float32, name: "coeffs") temb = multiplication(temb, coeffs, name: nil) temb = concatTensors([cos(with: temb, name: nil), sin(with: temb, name: nil)], dimension: 0, name: nil) temb = reshape(temb, shape: [1, 320], name: nil) return cast(temb, to: MPSDataType.float16, name: "temb fp16") } } ================================================ FILE: Sources/MapleDiffusion/MPS/TextGuidance.swift ================================================ // // TextGuidance.swift // // // Created by Guillermo Cique Fernández on 13/11/22. // import Foundation import MetalPerformanceShadersGraph class TextGuidance { private let device: MPSGraphDevice private let tokenizer: BPETokenizer private let executable: MPSGraphExecutable init(synchronize: Bool, modelLocation: URL, device: MPSGraphDevice) { self.device = device self.tokenizer = BPETokenizer(modelLocation: modelLocation) let graph = MPSGraph(synchronize: synchronize) let textGuidanceIn = graph.placeholder(shape: [2, 77], dataType: MPSDataType.int32, name: nil) let textGuidanceOut = graph.makeTextGuidance(at: modelLocation, xIn: textGuidanceIn, name: "cond_stage_model.transformer.text_model") let textGuidanceOut0 = graph.sliceTensor(textGuidanceOut, dimension: 0, start: 0, length: 1, name: nil) let textGuidanceOut1 = graph.sliceTensor(textGuidanceOut, dimension: 0, start: 1, length: 1, name: nil) self.executable = graph.compile( with: device, feeds: [ textGuidanceIn: MPSGraphShapedType(shape: textGuidanceIn.shape, dataType: MPSDataType.int32) ], targetTensors: [textGuidanceOut0, textGuidanceOut1], targetOperations: nil, compilationDescriptor: nil ) } func run(with queue: MTLCommandQueue, prompt: String, negativePrompt: String) -> (MPSGraphTensorData, MPSGraphTensorData) { let baseTokens = tokenizer.encode(s: negativePrompt) let tokens = tokenizer.encode(s: prompt) let data = (baseTokens + tokens).map {Int32($0)} .withUnsafeBufferPointer { Data(buffer: $0) } let tensorData = MPSGraphTensorData(device: device, data: data, shape: [2, 77], dataType: MPSDataType.int32) let res = executable.run(with: queue, inputs: [tensorData], results: nil, executionDescriptor: nil) return (res[0], res[1]) } } fileprivate extension MPSGraph { func makeTextGuidance(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var x = makeTextEmbeddings(at: folder, xIn: xIn, name: name + ".embeddings") x = makeTextEncoder(at: folder, xIn: x, name: name + ".encoder") return makeLayerNorm(at: folder, xIn: x, name: name + ".final_layer_norm") } func makeTextEmbeddings(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var tokenEmbeddings = loadConstant(at: folder, name: name + ".token_embedding.weight", shape: [1, 49408, 768]) tokenEmbeddings = broadcast(tokenEmbeddings, shape: [2, 49408, 768], name: nil) let positionEmbeddings = loadConstant(at: folder, name: name + ".position_embedding.weight", shape: [1, 77, 768]) var embeddings = broadcast(expandDims(xIn, axes: [2], name: nil), shape: [2, 77, 768], name: nil) embeddings = gatherAlongAxis(1, updates: tokenEmbeddings, indices: embeddings, name: nil) return addition(embeddings, positionEmbeddings, name: nil) } func makeTextAttention(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { let nHeads: NSNumber = 12 let dHead: NSNumber = 64 let c: NSNumber = 768 var q = makeLinear(at: folder, xIn: xIn, name: name + ".q_proj", outChannels: c) var k = makeLinear(at: folder, xIn: xIn, name: name + ".k_proj", outChannels: c) var v = makeLinear(at: folder, xIn: xIn, name: name + ".v_proj", outChannels: c) let n = xIn.shape![0] let t = xIn.shape![1] q = reshape(q, shape: [n, t, nHeads, dHead], name: nil) k = reshape(k, shape: [n, t, nHeads, dHead], name: nil) v = reshape(v, shape: [n, t, nHeads, dHead], name: nil) q = transposeTensor(q, dimension: 1, withDimension: 2, name: nil) k = transposeTensor(k, dimension: 1, withDimension: 2, name: nil) v = transposeTensor(v, dimension: 1, withDimension: 2, name: nil) var att = matrixMultiplication(primary: q, secondary: transposeTensor(k, dimension: 2, withDimension: 3, name: nil), name: nil) att = multiplication(att, constant(1.0 / sqrt(dHead.doubleValue), dataType: MPSDataType.float16), name: nil) att = addition(att, loadConstant(at: folder, name: "causal_mask", shape: [1, 1, 77, 77]), name: nil) att = softMax(with: att, axis: 3, name: nil) att = matrixMultiplication(primary: att, secondary: v, name: nil) att = transposeTensor(att, dimension: 1, withDimension: 2, name: nil) att = reshape(att, shape: [n, t, c], name: nil) return makeLinear(at: folder, xIn: att, name: name + ".out_proj", outChannels: c) } func makeTextEncoderLayer(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var x = xIn x = makeLayerNorm(at: folder, xIn: x, name: name + ".layer_norm1") x = makeTextAttention(at: folder, xIn: x, name: name + ".self_attn") x = addition(x, xIn, name: nil) let skip = x x = makeLayerNorm(at: folder, xIn: x, name: name + ".layer_norm2") x = makeLinear(at: folder, xIn: x, name: name + ".mlp.fc1", outChannels: 3072) x = gelu(x) x = makeLinear(at: folder, xIn: x, name: name + ".mlp.fc2", outChannels: 768) return addition(x, skip, name: nil) } func makeTextEncoder(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var x = xIn for i in 0..<12 { x = makeTextEncoderLayer(at: folder, xIn: x, name: name + ".layers.\(i)") } return x } } ================================================ FILE: Sources/MapleDiffusion/MPS/UNet.swift ================================================ // // UNet.swift // // // Created by Guillermo Cique Fernández on 9/11/22. // import Foundation import MetalPerformanceShadersGraph class UNet { let synchronize: Bool let modelLocation: URL let saveMemory: Bool let device: MPSGraphDevice // MEM-HACK: split into subgraphs private var unetAnUnexpectedJourneyExecutable: MPSGraphExecutable? private var anUnexpectedJourneyShapes = [[NSNumber]]() private var unetTheDesolationOfSmaugExecutable: MPSGraphExecutable? private var theDesolationOfSmaugShapes = [[NSNumber]]() private var theDesolationOfSmaugIndices = [MPSGraphTensor: Int]() private var unetTheBattleOfTheFiveArmiesExecutable: MPSGraphExecutable? private var theBattleOfTheFiveArmiesIndices = [MPSGraphTensor: Int]() init(synchronize: Bool, modelLocation: URL, saveMemory: Bool, device: MPSGraphDevice, shape: [NSNumber]) { self.synchronize = synchronize self.modelLocation = modelLocation self.saveMemory = saveMemory self.device = device loadAnUnexpectedJourney(shape: shape) loadTheDesolationOfSmaug() loadTheBattleOfTheFiveArmies() } private func loadAnUnexpectedJourney(shape: [NSNumber]) { let graph = MPSGraph(synchronize: synchronize) let xIn = graph.placeholder(shape: shape, dataType: MPSDataType.float16, name: nil) let condIn = graph.placeholder(shape: [saveMemory ? 1 : 2, 77, 768], dataType: MPSDataType.float16, name: nil) let tembIn = graph.placeholder(shape: [1, 320], dataType: MPSDataType.float16, name: nil) let unetOuts = graph.makeUNetAnUnexpectedJourney( at: modelLocation, xIn: xIn, tembIn: tembIn, condIn: condIn, name: "model.diffusion_model", saveMemory: saveMemory ) let unetFeeds = [xIn, condIn, tembIn].reduce(into: [:], {$0[$1] = MPSGraphShapedType(shape: $1.shape!, dataType: $1.dataType)}) unetAnUnexpectedJourneyExecutable = graph.compile( with: device, feeds: unetFeeds, targetTensors: unetOuts, targetOperations: nil, compilationDescriptor: nil ) anUnexpectedJourneyShapes = unetOuts.map{$0.shape!} } private func loadTheDesolationOfSmaug() { let graph = MPSGraph(synchronize: synchronize) let condIn = graph.placeholder(shape: [saveMemory ? 1 : 2, 77, 768], dataType: MPSDataType.float16, name: nil) let placeholders = anUnexpectedJourneyShapes.map{graph.placeholder(shape: $0, dataType: MPSDataType.float16, name: nil)} + [condIn] theDesolationOfSmaugIndices.removeAll() for i in 0.. [MPSGraphTensorData] { var out = [MPSGraphTensorData]() for r in unetAnUnexpectedJourneyExecutable!.feedTensors! { for i in x { if (i.shape == r.shape) { out.append(i) } } } return out } private func reorderTheDesolationOfSmaug(x: [MPSGraphTensorData]) -> [MPSGraphTensorData] { var out = [MPSGraphTensorData]() for r in unetTheDesolationOfSmaugExecutable!.feedTensors! { out.append(x[theDesolationOfSmaugIndices[r]!]) } return out } private func reorderTheBattleOfTheFiveArmies(x: [MPSGraphTensorData]) -> [MPSGraphTensorData] { var out = [MPSGraphTensorData]() for r in unetTheBattleOfTheFiveArmiesExecutable!.feedTensors! { out.append(x[theBattleOfTheFiveArmiesIndices[r]!]) } return out } private func runUNet( with queue: MTLCommandQueue, latent: MPSGraphTensorData, guidance: MPSGraphTensorData, temb: MPSGraphTensorData ) -> MPSGraphTensorData { var x = unetAnUnexpectedJourneyExecutable!.run( with: queue, inputs: reorderAnUnexpectedJourney(x: [latent, guidance, temb]), results: nil, executionDescriptor: nil ) x = unetTheDesolationOfSmaugExecutable!.run( with: queue, inputs: reorderTheDesolationOfSmaug(x: x + [guidance]), results: nil, executionDescriptor: nil ) return unetTheBattleOfTheFiveArmiesExecutable!.run( with: queue, inputs: reorderTheBattleOfTheFiveArmies(x: x + [guidance]), results: nil, executionDescriptor: nil )[0] } private func runBatchedUNet( with queue: MTLCommandQueue, latent: MPSGraphTensorData, baseGuidance: MPSGraphTensorData, textGuidance: MPSGraphTensorData, temb: MPSGraphTensorData ) -> (MPSGraphTensorData, MPSGraphTensorData) { // concat var graph = MPSGraph(synchronize: synchronize) let bg = graph.placeholder(shape: baseGuidance.shape, dataType: MPSDataType.float16, name: nil) let tg = graph.placeholder(shape: textGuidance.shape, dataType: MPSDataType.float16, name: nil) let concatGuidance = graph.concatTensors([bg, tg], dimension: 0, name: nil) let concatGuidanceData = graph.run( with: queue, feeds: [ bg : baseGuidance, tg: textGuidance ], targetTensors: [concatGuidance], targetOperations : nil )[concatGuidance]! // run let concatEtaData = runUNet(with: queue, latent: latent, guidance: concatGuidanceData, temb: temb) // split graph = MPSGraph(synchronize: synchronize) let etas = graph.placeholder(shape: concatEtaData.shape, dataType: concatEtaData.dataType, name: nil) let eta0 = graph.sliceTensor(etas, dimension: 0, start: 0, length: 1, name: nil) let eta1 = graph.sliceTensor(etas, dimension: 0, start: 1, length: 1, name: nil) let etaRes = graph.run( with: queue, feeds: [etas: concatEtaData], targetTensors: [eta0, eta1], targetOperations: nil ) return (etaRes[eta0]!, etaRes[eta1]!) } func run( with queue: MTLCommandQueue, latent: MPSGraphTensorData, baseGuidance: MPSGraphTensorData, textGuidance: MPSGraphTensorData, temb: MPSGraphTensorData ) -> (MPSGraphTensorData, MPSGraphTensorData) { if (saveMemory) { // MEM-HACK: un/neg-conditional and text-conditional are run in two separate passes (not batched) to save memory let etaUncond = runUNet(with: queue, latent: latent, guidance: baseGuidance, temb: temb) let etaCond = runUNet(with: queue, latent: latent, guidance: textGuidance, temb: temb) return (etaUncond, etaCond) } else { return runBatchedUNet(with: queue, latent: latent, baseGuidance: baseGuidance, textGuidance: textGuidance, temb: temb) } } } extension MPSGraph { func makeTimeEmbed(at folder: URL, xIn: MPSGraphTensor, name: String) -> MPSGraphTensor { var x = xIn x = makeLinear(at: folder, xIn: x, name: name + ".0", outChannels: 1280) x = swish(x) return makeLinear(at: folder, xIn: x, name: name + ".2", outChannels: 1280) } func makeUNetResBlock(at folder: URL, xIn: MPSGraphTensor, embIn: MPSGraphTensor, name: String, inChannels: NSNumber, outChannels: NSNumber) -> MPSGraphTensor { var x = xIn x = makeGroupNormSwish(at: folder, xIn: x, name: name + ".in_layers.0") x = makeConv(at: folder, xIn: x, name: name + ".in_layers.2", outChannels: outChannels, khw: 3) var emb = embIn emb = swish(emb) emb = makeLinear(at: folder, xIn: emb, name: name + ".emb_layers.1", outChannels: outChannels) emb = expandDims(emb, axes: [1, 2], name: nil) x = addition(x, emb, name: nil) x = makeGroupNormSwish(at: folder, xIn: x, name: name + ".out_layers.0") x = makeConv(at: folder, xIn: x, name: name + ".out_layers.3", outChannels: outChannels, khw: 3) var skip = xIn if (inChannels != outChannels) { skip = makeConv(at: folder, xIn: xIn, name: name + ".skip_connection", outChannels: outChannels, khw: 1) } return addition(x, skip, name: nil) } func makeOutputBlock(at folder: URL, xIn: MPSGraphTensor, embIn: MPSGraphTensor, condIn: MPSGraphTensor, inChannels: NSNumber, outChannels: NSNumber, dHead: NSNumber, name: String, saveMemory: Bool, spatialTransformer: Bool = true, upsample: Bool = false) -> MPSGraphTensor { var x = xIn x = makeUNetResBlock(at: folder, xIn: x, embIn: embIn, name: name + ".0", inChannels: inChannels, outChannels: outChannels) if (spatialTransformer) { x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".1", contextIn: condIn, saveMemory: saveMemory) } if (upsample) { x = upsampleNearest(xIn: x) x = makeConv(at: folder, xIn: x, name: name + (spatialTransformer ? ".2" : ".1") + ".conv", outChannels: outChannels, khw: 3) } return x } func makeUNetAnUnexpectedJourney(at folder: URL, xIn: MPSGraphTensor, tembIn: MPSGraphTensor, condIn: MPSGraphTensor, name: String, saveMemory: Bool = true) -> [MPSGraphTensor] { let emb = makeTimeEmbed(at: folder, xIn: tembIn, name: name + ".time_embed") var savedInputs = [MPSGraphTensor]() var x = xIn if (!saveMemory) { // need to explicitly batch to avoid shape errors later iirc // TODO: did we actually need this x = broadcast(x, shape: [condIn.shape![0], x.shape![1], x.shape![2], x.shape![3]], name: nil) } // input blocks x = makeConv(at: folder, xIn: x, name: name + ".input_blocks.0.0", outChannels: 320, khw: 3) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.1.0", inChannels: 320, outChannels: 320) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".input_blocks.1.1", contextIn: condIn, saveMemory: saveMemory) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.2.0", inChannels: 320, outChannels: 320) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".input_blocks.2.1", contextIn: condIn, saveMemory: saveMemory) savedInputs.append(x) // downsample x = makeConv(at: folder, xIn: x, name: name + ".input_blocks.3.0.op", outChannels: 320, khw: 3, stride: 2) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.4.0", inChannels: 320, outChannels: 640) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".input_blocks.4.1", contextIn: condIn, saveMemory: saveMemory) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.5.0", inChannels: 640, outChannels: 640) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".input_blocks.5.1", contextIn: condIn, saveMemory: saveMemory) savedInputs.append(x) // downsample x = makeConv(at: folder, xIn: x, name: name + ".input_blocks.6.0.op", outChannels: 640, khw: 3, stride: 2) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.7.0", inChannels: 640, outChannels: 1280) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".input_blocks.7.1", contextIn: condIn, saveMemory: saveMemory) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.8.0", inChannels: 1280, outChannels: 1280) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".input_blocks.8.1", contextIn: condIn, saveMemory: saveMemory) savedInputs.append(x) // downsample x = makeConv(at: folder, xIn: x, name: name + ".input_blocks.9.0.op", outChannels: 1280, khw: 3, stride: 2) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.10.0", inChannels: 1280, outChannels: 1280) savedInputs.append(x) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".input_blocks.11.0", inChannels: 1280, outChannels: 1280) savedInputs.append(x) // middle blocks x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".middle_block.0", inChannels: 1280, outChannels: 1280) x = makeSpatialTransformerBlock(at: folder, xIn: x, name: name + ".middle_block.1", contextIn: condIn, saveMemory: saveMemory) x = makeUNetResBlock(at: folder, xIn: x, embIn: emb, name: name + ".middle_block.2", inChannels: 1280, outChannels: 1280) return savedInputs + [emb] + [x] } func makeUNetTheDesolationOfSmaug(at folder: URL, savedInputsIn: [MPSGraphTensor], name: String, saveMemory: Bool = true) -> [MPSGraphTensor] { var savedInputs = savedInputsIn let condIn = savedInputs.popLast()! var x = savedInputs.popLast()! let emb = savedInputs.popLast()! // output blocks x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 2560, outChannels: 1280, dHead: 160, name: name + ".output_blocks.0", saveMemory: saveMemory, spatialTransformer: false, upsample: false) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 2560, outChannels: 1280, dHead: 160, name: name + ".output_blocks.1", saveMemory: saveMemory, spatialTransformer: false, upsample: false) // upsample x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 2560, outChannels: 1280, dHead: 160, name: name + ".output_blocks.2", saveMemory: saveMemory, spatialTransformer: false, upsample: true) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 2560, outChannels: 1280, dHead: 160, name: name + ".output_blocks.3", saveMemory: saveMemory, spatialTransformer: true, upsample: false) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 2560, outChannels: 1280, dHead: 160, name: name + ".output_blocks.4", saveMemory: saveMemory, spatialTransformer: true, upsample: false) return savedInputs + [emb] + [x] } func makeUNetTheBattleOfTheFiveArmies(at folder: URL, savedInputsIn: [MPSGraphTensor], name: String, saveMemory: Bool = true) -> MPSGraphTensor { var savedInputs = savedInputsIn let condIn = savedInputs.popLast()! var x = savedInputs.popLast()! let emb = savedInputs.popLast()! // upsample x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 1920, outChannels: 1280, dHead: 160, name: name + ".output_blocks.5", saveMemory: saveMemory, spatialTransformer: true, upsample: true) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 1920, outChannels: 640, dHead: 80, name: name + ".output_blocks.6", saveMemory: saveMemory, spatialTransformer: true, upsample: false) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 1280, outChannels: 640, dHead: 80, name: name + ".output_blocks.7", saveMemory: saveMemory, spatialTransformer: true, upsample: false) // upsample x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 960, outChannels: 640, dHead: 80, name: name + ".output_blocks.8", saveMemory: saveMemory, spatialTransformer: true, upsample: true) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 960, outChannels: 320, dHead: 40, name: name + ".output_blocks.9", saveMemory: saveMemory, spatialTransformer: true, upsample: false) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 640, outChannels: 320, dHead: 40, name: name + ".output_blocks.10", saveMemory: saveMemory, spatialTransformer: true, upsample: false) x = concatTensors([x, savedInputs.popLast()!], dimension: 3, name: nil) x = makeOutputBlock(at: folder, xIn: x, embIn: emb, condIn: condIn, inChannels: 640, outChannels: 320, dHead: 40, name: name + ".output_blocks.11", saveMemory: saveMemory, spatialTransformer: true, upsample: false) // out x = makeGroupNormSwish(at: folder, xIn: x, name: "model.diffusion_model.out.0") return makeConv(at: folder, xIn: x, name: "model.diffusion_model.out.2", outChannels: 4, khw: 3) } } ================================================ FILE: Sources/MapleDiffusion/Model/GenResult.swift ================================================ import Foundation //import AppKit import CoreImage import CoreGraphics public struct GenResult { internal init(image: CGImage?, progress: Double, stage: String) { self.image = image self.progress = progress self.stage = stage } public let image : CGImage? public let progress : Double public let stage : String } ================================================ FILE: Sources/MapleDiffusion/Model/GeneratorState.swift ================================================ // Created by Morten Just on 10/20/22. // import Foundation public enum GeneratorState: Equatable { public static func == (lhs: GeneratorState, rhs: GeneratorState) -> Bool { switch(lhs, rhs) { case (.ready, .ready): return true // case (.modelIsLoading(progress: <#T##Double#>, message: <#T##String#>)) case (.notStarted, .notStarted): return true default: return false } } case notStarted case modelIsLoading(progress: Double, message: String) case ready } ================================================ FILE: Sources/MapleDiffusion/Model/SampleInput.swift ================================================ // // SampleInput.swift // // // Created by Guillermo Cique Fernández on 14/11/22. // import Foundation import CoreGraphics public struct SampleInput { var prompt: String var negativePrompt: String var initImage: CGImage? { didSet { checkSize() }} var strength: Float? var seed: Int var steps: Int var guidanceScale: Float public init( prompt: String, negativePrompt: String = "", seed: Int = Int.random(in: 0...Int.max), steps: Int = 20, guidanceScale: Float = 7.5 ) { self.prompt = prompt self.negativePrompt = negativePrompt self.initImage = nil self.strength = nil self.seed = seed self.steps = steps self.guidanceScale = guidanceScale } public init( prompt: String, negativePrompt: String = "", initImage: CGImage?, strength: Float = 0.75, seed: Int = Int.random(in: 0...Int.max), steps: Int = 20, guidanceScale: Float = 5.0 ) { self.prompt = prompt self.negativePrompt = negativePrompt self.initImage = initImage self.strength = strength self.seed = seed self.steps = steps self.guidanceScale = guidanceScale } private func checkSize() { guard let initImage else { return } if initImage.width != 512 || initImage.height != 512 { assertionFailure("Please make sure your input image is exactly 512x512. You can use your own cropping mechanism, or the extension in NSImage:crop:to (macOS). Feel free to contribute a general solution. See Github issues for ideas.") } } } #if os(iOS) import UIKit public extension SampleInput { init( prompt: String, negativePrompt: String = "", initImage: UIImage, strength: Float = 0.75, seed: Int = Int.random(in: 0...Int.max), steps: Int = 50, guidanceScale: Float = 5.0 ) { self.prompt = prompt self.negativePrompt = negativePrompt self.initImage = initImage.cgImage! self.strength = strength self.seed = seed self.steps = steps self.guidanceScale = guidanceScale } } #endif #if os(macOS) import AppKit public extension SampleInput { init( prompt: String, negativePrompt: String = "", initImage: NSImage, strength: Float = 0.75, seed: Int = Int.random(in: 0...Int.max), steps: Int = 50, guidanceScale: Float = 5.0 ) { self.prompt = prompt self.negativePrompt = negativePrompt var imageRect = CGRect(x: 0, y: 0, width: initImage.size.width, height: initImage.size.height) self.initImage = initImage.cgImage(forProposedRect: &imageRect, context: nil, hints: nil)! self.strength = strength self.seed = seed self.steps = steps self.guidanceScale = guidanceScale } } #endif ================================================ FILE: Sources/MapleDiffusion/Model/SampleOutput.swift ================================================ // // SampleOutput.swift // // // Created by Guillermo Cique Fernández on 14/11/22. // import Foundation import CoreGraphics public struct SampleOutput { let image: CGImage? let input: SampleInput } ================================================ FILE: Sources/MapleDiffusion/Support/CGImage + ItemProvider.swift ================================================ // // File.swift // // // Created by Morten Just on 10/25/22. // #if os(macOS) import Foundation import AppKit import CoreGraphics import CoreImage import UniformTypeIdentifiers extension CGImage { func itemProvider(filename: String? = nil) -> NSItemProvider? { let rep = NSBitmapImageRep(cgImage: self) guard let data = rep.representation(using: .png, properties: [:]) else { return nil } let filename = filename ?? "Generated Image \(UUID().uuidString.prefix(4))" let tempUrl = FileManager.default.temporaryDirectory.appendingPathComponent(filename).appendingPathExtension("png") do { try data.write(to: tempUrl) let item = NSItemProvider(item: tempUrl as NSSecureCoding, typeIdentifier: UTType.fileURL.identifier) item.suggestedName = "\(filename).png" return item } catch { return nil } } } #endif ================================================ FILE: Sources/MapleDiffusion/Support/Diffusion + Generate.swift ================================================ // // File.swift // // // Created by Morten Just on 10/27/22. // import Foundation import Combine import CoreGraphics /** Generator functions. Mostly convenience wrappers around generate:input[..] */ public extension Diffusion { /// Generate an image async. Optional callback with progress and intermediate image. func generate( input: SampleInput, progress: ((GenResult) -> Void)? = nil, remoteUrl: String? = nil ) async -> CGImage? { var combinedSteps : Double = 1 var combinedProgress : Double = 0 // maybe load model first if modelIsCold, let remoteUrl, let url = URL(string: remoteUrl) { combinedSteps += 1 print("async gen: cold model, loading") try! await prepModels(remoteURL: url) { p in combinedProgress = (p/combinedSteps) let res = GenResult(image: nil, progress: combinedProgress, stage: "Loading Model") progress?(res) } } // generate image return await withCheckedContinuation { continuation in print("async gen: generating", input.prompt) self.generate(input: input) { (image, progressFloat, stage) in let genResult = GenResult(image: image, progress: Double(progressFloat), stage: stage) progress?(genResult) print("async gen: ", genResult.stage) if progressFloat == 1, let finalImage = genResult.image { continuation.resume(returning: finalImage) } if progressFloat == 1, genResult.image == nil { continuation.resume(returning: nil) } } } } /// Generate an image asynchronously. Optional callback with progress and intermediate image. Run inside a Task.detached to run in background. /// ``` /// // without progress reporting /// let image = await diffusion.generate("astronaut in the ocean") /// /// // with progress /// let image = await diffusion.generate("astronaut in the ocean") { progress in /// print("progress: \(progress.progress)") } func generate(prompt: String, negativePrompt: String = "", inputImage: CGImage? = nil, seed: Int = Int.random(in: 0...Int.max), steps:Int = 20, guidanceScale:Float = 7.5, progress: ((GenResult) -> Void)? = nil, remoteUrl: String? = nil ) async -> CGImage? { /// This is a just a wrapper that adds `SampleInput`, mainly for simplicity and to not break pre-SampleInput builds. let sampleInput = SampleInput(prompt: prompt, negativePrompt: negativePrompt, seed: seed, steps: steps, guidanceScale: guidanceScale) return await generate(input: sampleInput, progress: progress) } /// Generate an image and get a publisher for progress and intermediate image. Runs detached with "initiated" priority. /// ``` /// diffusion.generate("Astronaut in the ocean") /// .sink { result in /// print("progress: \(result.progress)") // result also contains the intermediate image /// }.store(in: ...) /// func generate(input: SampleInput) -> AnyPublisher { let publisher = PassthroughSubject() Task.detached(priority: .userInitiated) { // TODO: Test other priorities and consider making it an optional argument self.generate(input: input) { (cgImage, progress, stage) in Task { print("RAW progress", progress) await MainActor.run { let result = GenResult(image: cgImage, progress: Double(progress), stage: stage) publisher.send(result) if progress >= 1 { publisher.send(completion: .finished)} } } } } return publisher.eraseToAnyPublisher() } func generate(prompt: String, negativePrompt: String = "", seed: Int = Int.random(in: 0...Int.max), steps:Int = 20, guidanceScale:Float = 7.5) -> AnyPublisher { let sampleInput = SampleInput(prompt: prompt, negativePrompt: negativePrompt, seed: seed, steps: steps, guidanceScale: guidanceScale) return generate(input: sampleInput) } /// Generate and image with a callback on current thread. /// func generate( input: SampleInput, completion: @escaping (CGImage?, Float, String)->() ) { mapleDiffusion.generate(input: input) { cgImage, progress, stage in // penultimate step is also marked progress 1 in MD currently, working around it var realProgress = progress if progress == 1 && stage.contains("Decoding") { realProgress = 0.97 } completion(cgImage, realProgress, stage) } } } ================================================ FILE: Sources/MapleDiffusion/Support/Drag + Drop Helpers.swift ================================================ // // File.swift // // // Created by Morten Just on 11/21/22. // import Foundation import CoreGraphics import SwiftUI #if os(macOS) /// Allow safe and optional item dragging struct FinderDraggable : ViewModifier { let image: CGImage let enabled: Bool func body(content: Content) -> some View { if enabled, let provider = image.itemProvider() { content.onDrag { provider } } else { content } } } struct FinderImageDroppable : ViewModifier { var isTargeted: Binding? let forceSize: CGSize? var onDrop : ((CGImage) -> Void) func loadImage(url:URL) { let nsImage = NSImage(contentsOf: url) let cgImage = nsImage?.cgImage(forProposedRect: nil, context: nil, hints: nil) // resize here if forcesize is there if let cgImage { onDrop(cgImage) } } func body(content: Content) -> some View { content .droppable(isTargeted: isTargeted) { url in loadImage(url: url) } } } struct FinderDroppable : ViewModifier { var isTargeted : Binding? var onDrop : (URL) -> Void func body(content: Content) -> some View { content .onDrop(of: [.fileURL], isTargeted: isTargeted) { providers in if let provider = (providers.first { $0.canLoadObject(ofClass: URL.self ) }) { let _ = provider.loadObject(ofClass: URL.self) { reading, error in if let reading { onDrop(reading) } if let error { print("error", error) } } return true } return false } } } extension View { public func draggable(enabled: Bool, image: CGImage) -> some View { self.modifier(FinderDraggable(image: image, enabled: enabled)) } public func droppable(isTargeted:Binding? = nil, onDrop: @escaping (URL)->Void) -> some View { self.modifier(FinderDroppable(isTargeted: isTargeted, onDrop: onDrop)) } public func imageDroppable(isTargeted: Binding? = nil, forceSize: CGSize? = nil, onDrop: @escaping (CGImage)->Void) -> some View { self.modifier(FinderImageDroppable(isTargeted: isTargeted, forceSize:forceSize, onDrop: onDrop)) } } #endif ================================================ FILE: Sources/MapleDiffusion/Support/MPSGraphTensorData.swift ================================================ // // MPSGraphTensorData.swift // // // Created by Guillermo Cique Fernández on 9/11/22. // import Foundation import MetalPerformanceShadersGraph extension MPSGraphTensorData { var cgImage: CGImage? { let shape = self.shape.map{ $0.intValue } var imageArrayCPUBytes = [UInt8](repeating: 0, count: shape.reduce(1, *)) self.mpsndarray().readBytes(&imageArrayCPUBytes, strideBytes: nil) return CGImage( width: shape[2], height: shape[1], bitsPerComponent: 8, bitsPerPixel: 32, bytesPerRow: shape[2]*shape[3], space: CGColorSpaceCreateDeviceRGB(), bitmapInfo: CGBitmapInfo(rawValue: CGBitmapInfo.byteOrder32Big.rawValue | CGImageAlphaInfo.noneSkipLast.rawValue), provider: CGDataProvider(data: NSData(bytes: &imageArrayCPUBytes, length: imageArrayCPUBytes.count))!, decode: nil, shouldInterpolate: true, intent: CGColorRenderingIntent.defaultIntent ) } public convenience init(device: MPSGraphDevice, cgImage: CGImage) { let shape: [NSNumber] = [NSNumber(value: cgImage.height), NSNumber(value: cgImage.width), 4] let data = cgImage.dataProvider!.data! as Data self.init(device: device, data: data, shape: shape, dataType: .uInt8) } } /* bitsPerPixel 32 bytesPerRow 2048 byteOrderInfo CGImageByteOrderInfo colorSpace Optional( (kCGColorSpaceICCBased; kCGColorSpaceModelRGB; sRGB IEC61966-2.1)) */ extension Int { func tensorData(device: MPSGraphDevice) -> MPSGraphTensorData { let data = [Int32(self)].withUnsafeBufferPointer { Data(buffer: $0) } return MPSGraphTensorData(device: device, data: data, shape: [1], dataType: MPSDataType.int32) } } extension Float { func tensorData(device: MPSGraphDevice) -> MPSGraphTensorData { let data = [Float32(self)].withUnsafeBufferPointer { Data(buffer: $0) } return MPSGraphTensorData(device: device, data: data, shape: [1], dataType: MPSDataType.float32) } } ================================================ FILE: Sources/MapleDiffusion/Support/ModelFetcher.swift ================================================ // // File.swift // // // Created by Morten Just on 10/22/22. // import Foundation import Combine import ZIPFoundation class ModelFetcher { let local: URL? let remote: URL? private var bin = Set() // Initializer overloads /// We saved our model locally, e.g. in the bundle. init(local: URL) { self.local = local self.remote = nil } /// (Recommended) Use the models at the default location. Download it first if not there. init(remote: URL) { self.remote = remote self.local = nil } /// Use the models at the given local location. Download to this destination if not there. init(local: URL, remote: URL) { self.local = local self.remote = remote } struct FileDownload { let url : URL? let progress: Double let stage:String } // Helpers private var bundleId : String { Bundle.main.bundleIdentifier ?? "app.otato.diffusion" } var finalLocalUrl: URL { local ?? fallbackModelFolder } /// Used if no local URL is provided var fallbackModelFolder : URL { // TODO: If on iOS use the bundle FileManager.default .urls(for: .applicationSupportDirectory, in: .userDomainMask)[0] .appendingPathComponent(bundleId) .appendingPathComponent("bins") } /// Checks final destination and returns true if it's not empty var isModelPresentLocally : Bool { do { let files = try FileManager.default .contentsOfDirectory(at: finalLocalUrl, includingPropertiesForKeys: nil) return !files.isEmpty } catch { return false } } // methods enum ModelFetcherError : Error { case emptyFolderAndNoRemoteURL } func fetch(progress: ProgressClosure?) async throws -> URL { let combinedSteps = 2.0 // Download and unzip var combinedProgress = 0.0 // is the calculated local folder, and it's not empty? go! if isModelPresentLocally { progress?(1.0) return finalLocalUrl } // create the folder if it doesn't exist try FileManager.default.createDirectory(at: finalLocalUrl, withIntermediateDirectories: true) // if the local folder is empty and we have a remote, start downloading and return a stream if !isModelPresentLocally, let remote { // 1 download let downloadedURL = try await URLSession .shared .downloadWithProgress(url: remote) { p in combinedProgress += (p/combinedSteps) progress?(combinedProgress) } // 2. unzip try await moveAndUnzip(downloadedURL) { p in combinedProgress += (p/combinedSteps) progress?(combinedSteps) } // 3 done! return finalLocalUrl } // with the overloads we shouldn't end up here assertionFailure() throw ModelFetcherError.emptyFolderAndNoRemoteURL // fatalError("The model folder is empty or not there, and no remote URL was provided") } func moveAndUnzip(_ url : URL, completion: ProgressClosure?) async throws { let progress = Progress() var bin = Set() progress.publisher(for: \.fractionCompleted) .map { Double(integerLiteral: $0) } .sink { value in completion?(value) }.store(in: &bin) try FileManager.default.unzipItem(at: url, to: finalLocalUrl, progress: progress) } } public typealias ProgressClosure = (Double) -> Void ================================================ FILE: Sources/MapleDiffusion/Support/NSImage + resizing.swift ================================================ // // Image resizing and cropping for input images // // // Created by Morten Just on 11/22/22. // #if os(macOS) import Foundation import AppKit extension NSImage { var isLandscape : Bool { size.height < size.width } func scale(factor: CGFloat) -> NSImage { let newWidth = self.size.width * factor let newHeight = self.size.height * factor let newSize = NSSize(width: newWidth, height: newHeight) // Draw self into a new image with the new size let newImage = NSImage(size: newSize, flipped: false) { rect in self.draw(in: .init(x: 0, y: 0, width: newWidth, height: newHeight)) return true } return newImage } @objc var cgImage: CGImage? { get { guard let imageData = self.tiffRepresentation else { return nil } guard let sourceData = CGImageSourceCreateWithData(imageData as CFData, nil) else { return nil } return CGImageSourceCreateImageAtIndex(sourceData, 0, nil) } } func crop(to newSize:NSSize) -> NSImage { // scale var factor : CGFloat = 1 if isLandscape { factor = newSize.height / self.size.height } else { factor = newSize.width / self.size.width } let scaledImage = scale(factor: factor) /// Find the center crop rect var fromRect = NSRect(x: 0, y: 0, width: newSize.width, height: newSize.height) if self.isLandscape { fromRect.origin.x = (0.5 * scaledImage.size.width) - (0.5 * newSize.width) } else { fromRect.origin.y = (0.5 * scaledImage.size.height) - (0.5 * newSize.height) } guard let rep = NSBitmapImageRep( bitmapDataPlanes: nil, pixelsWide: Int(newSize.width), pixelsHigh: Int(newSize.height), bitsPerSample: 8, samplesPerPixel: 4, hasAlpha: true, isPlanar: false, colorSpaceName: .deviceRGB, bytesPerRow: 0, bitsPerPixel: 0 ) else { preconditionFailure() } /// Get rid of retina pixels NSGraphicsContext.saveGraphicsState() NSGraphicsContext.current = NSGraphicsContext(bitmapImageRep: rep) scaledImage.draw(at: .zero, from: fromRect, operation: .sourceOver, fraction: 1.0) NSGraphicsContext.restoreGraphicsState() let data = rep.representation(using: .tiff, properties: [:]) let newImage = NSImage(data: data!) return newImage! } } #endif ================================================ FILE: Sources/MapleDiffusion/Support/URLSession + async download.swift ================================================ // // File.swift // // // Created by Morten Just on 10/24/22. // import Foundation import Combine extension URLSession { func downloadWithProgress(url: URL, progress:((Double)->Void)? = nil) async throws -> URL { var downloadTask: URLSessionDownloadTask? var bin = Set() return try await withCheckedThrowingContinuation({ continuation in let request = URLRequest(url: url) downloadTask = URLSession.shared.downloadTask(with: request) { url, response, error in if let url { continuation.resume(returning: url) } if let error { continuation.resume(throwing: error) } } /// Send progress to closure if we've got one' if let progress, let downloadTask { downloadTask.publisher(for: \.progress.fractionCompleted) .sink(receiveValue: { p in progress(p) }).store(in: &bin) } downloadTask?.resume() }) } } ================================================ FILE: Sources/MapleDiffusion/Views/DiffusionImage.swift ================================================ // // File.swift // // // Created by Morten Just on 10/21/22. // import Foundation import SwiftUI import UniformTypeIdentifiers /// Displays a CGImage and blurs it according to its generation progress public struct DiffusionImage : View { @Binding var image : CGImage? @Binding var progress : Double var inputImage : Binding? @State var isTargeted = false let label : String let draggable : Bool public init(image: Binding, inputImage: Binding? = nil, progress: Binding,c label: String = "Generated Image", draggable: Bool = true) { self._image = image self._progress = progress self.inputImage = inputImage self.label = label self.draggable = draggable } var enableDrag : Bool { guard draggable else { return false } return progress == 1 ? true : false } public var body: some View { ZStack { if let i = inputImage?.wrappedValue { Image(i, scale: 1, label: Text("Input image")) .resizable() .aspectRatio(contentMode: .fit) } if let image { Image(image, scale: 1, label: Text(label)) .resizable() .aspectRatio(contentMode: .fit) .animation(nil) .blur(radius: (1 - sqrt(progress)) * 100 ) .blendMode(progress < 1 ? .sourceAtop : .normal) .animation(.linear(duration: 1), value: progress) .clipShape(Rectangle()) #if os(macOS) .draggable(enabled: enableDrag, image: image) #endif } } #if os(macOS) .frame(maxWidth: .infinity, maxHeight: .infinity) .imageDroppable() { image in let ns = NSImage(cgImage: image, size: .init(width: image.width, height: image.height)) let cropped = ns.crop(to: .init(width: 512, height: 512)) let cg = cropped.cgImage self.inputImage?.wrappedValue = cg self.image = nil } #endif } } ================================================ FILE: Tests/MapleDiffusionTests/ModelFetcherTests.swift ================================================ // // ModelFetcherTests.swift // // // Created by Morten Just on 10/24/22. // import XCTest @testable import MapleDiffusion /** Warning: These tests are very much work in progress, and far from pure. They require set up and behaviors change from time to time. For example, if you run `testRemoteOnly` with no downloaded folders, it will download. On the second run, it wili not download. I guess the next steps would be to mainipulate the file system to create the same test context every time. */ final class ModelFetcherTests: XCTestCase { /// To test, start a web server in the folder you're hosting Diffusion.zip (a zip of all the converted model files) let remoteUrl = URL(string: "http://localhost:8080/Diffusion.zip")! let localUrl = FileManager.default.urls(for: .applicationSupportDirectory, in: .userDomainMask)[0].appendingPathComponent("com.example.myapp/bins") override func setUpWithError() throws { } override func tearDownWithError() throws { } /// In this case, the developer handles the downloadnig or is embedding in the bundle. func testFetchLocalOnly() async throws { /// Here, we're given only a file URL. We'll return an async stream that only outputs one element let fetcher = ModelFetcher(local: localUrl) var url : URL? for await status in fetcher.fetch() { print("status: ", status) if let u = status.url { url = u } } XCTAssertNotNil(url) print("--> got ", url!) let fileCount = try! FileManager.default.contentsOfDirectory(at: url!, includingPropertiesForKeys: nil).count XCTAssert(fileCount > 0) } func testRemoteOnly() async throws { let fetcher = ModelFetcher(remote: remoteUrl) var url: URL? for await status in fetcher.fetch() { print("status: ", status) if let u = status.url { url = u } } XCTAssertNotNil(url) print("--> got", url!) let fileCount = try! FileManager.default.contentsOfDirectory(at: url!, includingPropertiesForKeys: nil).count XCTAssert(fileCount > 0) } func testLocalAndRemoteOnCleanSlate() async throws { try FileManager.default.contentsOfDirectory(at: localUrl, includingPropertiesForKeys: nil) .forEach { url in print("deleting", url) try FileManager.default.removeItem(at: url) } try await testLocalAndRemote() } func testLocalAndRemote() async throws { let fetcher = ModelFetcher( local: localUrl, remote: remoteUrl) var url: URL? for await status in fetcher.fetch() { print("status: ", status) if let u = status.url { url = u } } XCTAssertNotNil(url) print("--> Final URL", url!) let fileCount = try! FileManager.default.contentsOfDirectory(at: url!, includingPropertiesForKeys: nil).count XCTAssert(fileCount > 0) } func testPerformanceExample() throws { // This is an example of a performance test case. self.measure { // Put the code you want to measure the time of here. } } }