Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improved metadata, second attempt #406

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 6 additions & 2 deletions Mochi Diffusion.xcodeproj/project.pbxproj
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@
C1220FC02AFB122F007E5055 /* ImageWellView.swift in Sources */ = {isa = PBXBuildFile; fileRef = C1220FBF2AFB122F007E5055 /* ImageWellView.swift */; };
C1ADEC2C2B16957800E142CA /* FolderMonitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = C1ADEC2B2B16957800E142CA /* FolderMonitor.swift */; };
D7B03F2029D42F9900DF89DD /* SDModelAttentionType.swift in Sources */ = {isa = PBXBuildFile; fileRef = D7B03F1F29D42F9900DF89DD /* SDModelAttentionType.swift */; };
DF1521DE2B95014100D7A82E /* MetadataHelper.swift in Sources */ = {isa = PBXBuildFile; fileRef = DF1521DD2B95014100D7A82E /* MetadataHelper.swift */; };
/* End PBXBuildFile section */

/* Begin PBXFileReference section */
Expand Down Expand Up @@ -132,6 +133,7 @@
C1220FBF2AFB122F007E5055 /* ImageWellView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ImageWellView.swift; sourceTree = "<group>"; };
C1ADEC2B2B16957800E142CA /* FolderMonitor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FolderMonitor.swift; sourceTree = "<group>"; };
D7B03F1F29D42F9900DF89DD /* SDModelAttentionType.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SDModelAttentionType.swift; sourceTree = "<group>"; };
DF1521DD2B95014100D7A82E /* MetadataHelper.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MetadataHelper.swift; sourceTree = "<group>"; };
/* End PBXFileReference section */

/* Begin PBXFrameworksBuildPhase section */
Expand Down Expand Up @@ -174,6 +176,7 @@
9B7392C2298DEAC2006F03D5 /* Tokenizer.swift */,
03FD231A295F7A81006EEEE2 /* Upscaler.swift */,
71D315152B2A584E007FFDEB /* NotificationController.swift */,
DF1521DD2B95014100D7A82E /* MetadataHelper.swift */,
);
path = Support;
sourceTree = "<group>";
Expand Down Expand Up @@ -439,6 +442,7 @@
03D280F2294FD60E00C7D184 /* PromptView.swift in Sources */,
03FD2319295F74B6006EEEE2 /* RealESRGAN.mlmodel in Sources */,
035EE000295A224900080D18 /* ModelView.swift in Sources */,
DF1521DE2B95014100D7A82E /* MetadataHelper.swift in Sources */,
03E075792968A153003488F9 /* CircularProgressView.swift in Sources */,
58FF218829FAE96200A70C7C /* ControlNetView.swift in Sources */,
C11A03882B099D9600540F7F /* JobQueueView.swift in Sources */,
Expand Down Expand Up @@ -641,7 +645,7 @@
CURRENT_PROJECT_VERSION = 5.0;
DEAD_CODE_STRIPPING = YES;
DEVELOPMENT_ASSET_PATHS = "\"Mochi Diffusion/Preview Content\"";
DEVELOPMENT_TEAM = TCQ6328PP6;
DEVELOPMENT_TEAM = SY87JRZ2T6;
ENABLE_HARDENED_RUNTIME = YES;
ENABLE_PREVIEWS = YES;
GENERATE_INFOPLIST_FILE = YES;
Expand Down Expand Up @@ -674,7 +678,7 @@
CURRENT_PROJECT_VERSION = 5.0;
DEAD_CODE_STRIPPING = YES;
DEVELOPMENT_ASSET_PATHS = "\"Mochi Diffusion/Preview Content\"";
DEVELOPMENT_TEAM = TCQ6328PP6;
DEVELOPMENT_TEAM = SY87JRZ2T6;
ENABLE_HARDENED_RUNTIME = YES;
ENABLE_PREVIEWS = YES;
GENERATE_INFOPLIST_FILE = YES;
Expand Down
10 changes: 3 additions & 7 deletions Mochi Diffusion/Model/SDImage.swift
Original file line number Diff line number Diff line change
Expand Up @@ -111,13 +111,9 @@ extension SDImage {
1,
nil
) else { return nil }
let iptc = await [
kCGImagePropertyIPTCCaptionAbstract: metadata(),
kCGImagePropertyIPTCOriginatingProgram: "Mochi Diffusion",
kCGImagePropertyIPTCProgramVersion: "\(await NSApplication.appVersion)"
]
let meta = [kCGImagePropertyIPTCDictionary: iptc]
CGImageDestinationAddImage(destination, image, meta as CFDictionary)
let metadata = CreateMetadata(positivePrompt: prompt, negativePrompt: negativePrompt, width: width, height: height, seed: seed, guidanceScale: Float(guidanceScale), scheduler: scheduler, stepCount: steps, currentModel: model, upscaler: upscaler, currentStyle: "", computeUnits: mlComputeUnit ?? .cpuAndGPU)

CGImageDestinationAddImage(destination, image, metadata)
guard CGImageDestinationFinalize(destination) else { return nil }
return data as Data
}
Expand Down
95 changes: 95 additions & 0 deletions Mochi Diffusion/Support/MetadataHelper.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,95 @@
//
// MetadataHelper.swift
// Mochi Diffusion
//
// Created by Jones on 18/02/2024.
//

import Foundation
import CoreGraphics
import AppKit
import CoreML
// import GuernikaKit
import StableDiffusion

func CreateMetadata (positivePrompt: String, negativePrompt: String, width: Int, height: Int, seed: UInt32, guidanceScale: Float, scheduler: Scheduler, stepCount: Int, currentModel: String, upscaler: String, currentStyle: String, computeUnits: MLComputeUnits)-> NSMutableDictionary{
var schedulerString = ""
var mlComputeUnit = ""
/// Schedulers for ML-Stable-diffusion
if scheduler == .pndmScheduler{
schedulerString = "PNDM"}
if scheduler == .dpmSolverMultistepScheduler{
schedulerString = "DPM-Solver++"}

/// Schedulers for Guernika kit
// if Scheduler == .ddim{
// SchedulerString = "DDIM"}
// if Scheduler == .dpmSolverMultistep{
// SchedulerString = "DPM++ 2M"}
// if Scheduler == .dpmSolverMultistepKarras{
// SchedulerString = "DPM++ 2M Karras"}
// if Scheduler == .dpmSolverSinglestep{
// SchedulerString = "DPM++ SDE"}
// if Scheduler == .dpmSolverSinglestepKarras{
// SchedulerString = "DPM++ SDE Karras"}
// if Scheduler == .dpm2{
// SchedulerString = "DPM2"}
// if Scheduler == .dpm2Karras{
// SchedulerString = "DPM2 Karras"}
// if Scheduler == .eulerDiscrete{
// SchedulerString = "Euler"}
// if Scheduler == .eulerDiscreteKarras{
// SchedulerString = "Euler Karras"}
// if Scheduler == .eulerAncenstralDiscrete{
// SchedulerString = "Euler Ancenstral"}
// if Scheduler == .lcm{
// SchedulerString = "LCM"}
// if Scheduler == .pndm{
// SchedulerString = "PNDM"}


if computeUnits == .cpuOnly{
mlComputeUnit = "cpuOnly"}
if computeUnits == .cpuAndGPU{
mlComputeUnit = "cpuAndGPU"}
if computeUnits == .cpuAndNeuralEngine{
mlComputeUnit = "cpuAndNeuralEngine"}

// Add metadata
let meta = NSMutableDictionary()
// EXIF Metadata (User Comment) - used by most other apps "c" and "uc" for prompts are established naming conventions kept for compatability.
let exifMetadata = NSMutableDictionary()
exifMetadata[kCGImagePropertyExifUserComment as String] = "{\"c\":\"\(positivePrompt)\", \"uc\":\"\(negativePrompt)\", \"seed\":\(seed), \"guidance_scale\":\(guidanceScale), \"sampler\":\"\(schedulerString)\", \"steps\":\(stepCount), \"model\":\"\(currentModel)\", \"Upscaler\":\"\(upscaler)\", \"styles\":\"\(currentStyle)\"}"
meta[kCGImagePropertyExifDictionary as String] = exifMetadata
// IPTC Metadata (Caption)
let iptcMetadata = NSMutableDictionary()

iptcMetadata[kCGImagePropertyIPTCCaptionAbstract as String] = """
\(Metadata.includeInImage.rawValue): \(positivePrompt); \
\(Metadata.excludeFromImage.rawValue): \(negativePrompt); \
\(Metadata.model.rawValue): \(currentModel); \
\(Metadata.steps.rawValue): \(stepCount); \
\(Metadata.guidanceScale.rawValue): \(guidanceScale); \
\(Metadata.seed.rawValue): \(seed); \
\(Metadata.size.rawValue): \(width)x\(height);
"""
+
(!upscaler.isEmpty ? " \(Metadata.upscaler.rawValue): \(upscaler); " : " ")
+
"""
\(Metadata.scheduler.rawValue): \(scheduler.rawValue); \
\(Metadata.mlComputeUnit.rawValue): \(mlComputeUnit); \
\(Metadata.generator.rawValue): Mochi Diffusion \(NSApplication.appVersion)
"""

meta[kCGImagePropertyIPTCCaptionAbstract as String] = iptcMetadata
meta[kCGImagePropertyIPTCOriginatingProgram as String] = "Mochi Diffusion"
meta[kCGImagePropertyIPTCProgramVersion as String] = "\(NSApplication.appVersion)"
meta[kCGImagePropertyIPTCDictionary as String] = iptcMetadata

// TIFF Metadata (Image Description)
let tiffMetadata = NSMutableDictionary()
tiffMetadata[kCGImagePropertyTIFFImageDescription as String] = "c:\(positivePrompt), uc:\(negativePrompt), seed:\(seed), Guidance Scale:\(guidanceScale), Sampler:\(schedulerString), Steps:\(stepCount), Model:\(currentModel), upscaler:\(upscaler), Size:\(width)x\(height)"
meta[kCGImagePropertyTIFFDictionary as String] = tiffMetadata
return meta
}