add image to image generation (#4)

Author: mzbac

Sources/FLUX.swift

@@ -143,11 +143,11 @@ open class FLUX {
   }
 }
 
-public class Flux1Schnell: FLUX, TextToImageGenerator, @unchecked Sendable {
+public class Flux1Schnell: FLUX, TextToImageGenerator, ImageToImageGenerator, @unchecked Sendable {
   let clipTokenizer: CLIPTokenizer
   let t5Tokenizer: any Tokenizer
-  let transformer: MultiModalDiffusionTransformer
-  let vae: VAE
+  public let transformer: MultiModalDiffusionTransformer
+  public let vae: VAE
   let t5Encoder: T5Encoder
   let clipEncoder: CLIPEncoder
 
@@ -186,27 +186,7 @@ public class Flux1Schnell: FLUX, TextToImageGenerator, @unchecked Sendable {
     return transformer
   }
 
-  public func generateLatents(parameters: EvaluateParameters) -> DenoiseIterator {
-    let latentsShape = [1, (parameters.height / 16) * (parameters.width / 16), 64]
-    let latents: MLXArray
-    if let seed = parameters.seed {
-      latents = MLXRandom.normal(latentsShape, key: MLXRandom.key(seed))
-    } else {
-      latents = MLXRandom.normal(latentsShape)
-    }
-    let (promptEmbeddings, pooledPromptEmbeddings) = conditionText(prompt: parameters.prompt)
-
-    return DenoiseIterator(
-      steps: parameters.numInferenceSteps,
-      promptEmbeddings: promptEmbeddings,
-      pooledPromptEmbeddings: pooledPromptEmbeddings,
-      latents: latents,
-      evaluateParameters: parameters,
-      transformer: transformer
-    )
-  }
-
-  func conditionText(prompt: String) -> (MLXArray, MLXArray) {
+  public func conditionText(prompt: String) -> (MLXArray, MLXArray) {
     let t5Tokens = t5Tokenizer.encode(text: prompt, addSpecialTokens: true)
     let paddedT5Tokens =
       Array(t5Tokens.prefix(256))
@@ -231,19 +211,19 @@ public class Flux1Schnell: FLUX, TextToImageGenerator, @unchecked Sendable {
   public func detachedDecoder() -> ImageDecoder {
     let autoencoder = self.vae
     func decode(xt: MLXArray) -> MLXArray {
-      var x = autoencoder.decode(latents: xt)
+      var x = autoencoder.decode(xt)
       x = clip(x / 2 + 0.5, min: 0, max: 1)
       return x
     }
     return decode(xt:)
   }
 }
 
-public class Flux1Dev: FLUX, TextToImageGenerator, @unchecked Sendable {
+public class Flux1Dev: FLUX, TextToImageGenerator, ImageToImageGenerator, @unchecked Sendable {
   let clipTokenizer: CLIPTokenizer
   let t5Tokenizer: any Tokenizer
-  let transformer: MultiModalDiffusionTransformer
-  let vae: VAE
+  public let transformer: MultiModalDiffusionTransformer
+  public let vae: VAE
   let t5Encoder: T5Encoder
   let clipEncoder: CLIPEncoder
 
@@ -283,27 +263,7 @@ public class Flux1Dev: FLUX, TextToImageGenerator, @unchecked Sendable {
     return transformer
   }
 
-  public func generateLatents(parameters: EvaluateParameters) -> DenoiseIterator {
-    let latentsShape = [1, (parameters.height / 16) * (parameters.width / 16), 64]
-    let latents: MLXArray
-    if let seed = parameters.seed {
-      latents = MLXRandom.normal(latentsShape, key: MLXRandom.key(seed))
-    } else {
-      latents = MLXRandom.normal(latentsShape)
-    }
-    let (promptEmbeddings, pooledPromptEmbeddings) = conditionText(prompt: parameters.prompt)
-
-    return DenoiseIterator(
-      steps: parameters.numInferenceSteps,
-      promptEmbeddings: promptEmbeddings,
-      pooledPromptEmbeddings: pooledPromptEmbeddings,
-      latents: latents,
-      evaluateParameters: parameters,
-      transformer: transformer
-    )
-  }
-
-  func conditionText(prompt: String) -> (MLXArray, MLXArray) {
+  public func conditionText(prompt: String) -> (MLXArray, MLXArray) {
     let t5Tokens = t5Tokenizer.encode(text: prompt, addSpecialTokens: true)
     let paddedT5Tokens =
       Array(t5Tokens.prefix(512))
@@ -328,7 +288,7 @@ public class Flux1Dev: FLUX, TextToImageGenerator, @unchecked Sendable {
   public func detachedDecoder() -> ImageDecoder {
     let autoencoder = self.vae
     func decode(xt: MLXArray) -> MLXArray {
-      var x = autoencoder.decode(latents: xt)
+      var x = autoencoder.decode(xt)
       x = clip(x / 2 + 0.5, min: 0, max: 1)
       return x
     }
@@ -349,7 +309,87 @@ public protocol ImageGenerator {
 }
 
 public protocol TextToImageGenerator: ImageGenerator, Sendable {
-  func generateLatents(parameters: EvaluateParameters) -> DenoiseIterator
+  var transformer: MultiModalDiffusionTransformer { get }
+  func conditionText(prompt: String) -> (MLXArray, MLXArray)
+}
+
+extension TextToImageGenerator {
+  public func generateLatents(parameters: EvaluateParameters) -> DenoiseIterator {
+    let latentsShape = [1, (parameters.height / 16) * (parameters.width / 16), 64]
+    let latents: MLXArray
+    if let seed = parameters.seed {
+      latents = MLXRandom.normal(latentsShape, key: MLXRandom.key(seed))
+    } else {
+      latents = MLXRandom.normal(latentsShape)
+    }
+    let (promptEmbeddings, pooledPromptEmbeddings) = conditionText(prompt: parameters.prompt)
+    return DenoiseIterator(
+      steps: parameters.numInferenceSteps,
+      promptEmbeddings: promptEmbeddings,
+      pooledPromptEmbeddings: pooledPromptEmbeddings,
+      latents: latents,
+      evaluateParameters: parameters,
+      transformer: transformer
+    )
+  }
+}
+
+/// Public interface for transforming a text prompt into an image.
+///
+/// Steps:
+///
+/// - ``generateLatents(image:parameters:strength:)``
+/// - evaluate each of the latents from the iterator
+/// - ``ImageGenerator/decode(xt:)`` or ``ImageGenerator/detachedDecoder()`` to convert the final latent into an image
+/// - use ``Image`` to save the image
+public protocol ImageToImageGenerator: ImageGenerator, Sendable {
+  var transformer: MultiModalDiffusionTransformer { get }
+  var vae: VAE { get }
+  func conditionText(prompt: String) -> (MLXArray, MLXArray)
+  func generateLatents(image: MLXArray, parameters: EvaluateParameters, strength: Float)
+    -> DenoiseIterator
+}
+
+extension ImageToImageGenerator {
+  internal func packLatents(latents: MLXArray, height: Int, width: Int) -> MLXArray {
+    let reshaped = latents.reshaped(1, height / 16, 2, width / 16, 2, 16)
+    let transposed = reshaped.transposed(0, 1, 3, 5, 2, 4)
+    return transposed.reshaped(1, (height / 16) * (width / 16), 64)
+  }
+
+  public func generateLatents(image: MLXArray, parameters: EvaluateParameters, strength: Float)
+    -> DenoiseIterator
+  {
+    if let seed = parameters.seed {
+      MLXRandom.seed(seed)
+    }
+    let noise = MLXRandom.normal([1, (parameters.height / 16) * (parameters.width / 16), 64])
+
+    // Calculate the start step and number of steps based on strength
+    let strength = max(0.0, min(1.0, strength))
+
+    let startStep = max(1, Int(Float(parameters.numInferenceSteps) * strength))
+
+    var latents = vae.encode(image[.newAxis])
+
+    latents = packLatents(latents: latents, height: parameters.height, width: parameters.width)
+
+    let sigma = parameters.sigmas[startStep]
+
+    latents = (latents * (1 - sigma) + sigma * noise)
+
+    let (promptEmbeddings, pooledPromptEmbeddings) = conditionText(prompt: parameters.prompt)
+
+    return DenoiseIterator(
+      startStep: startStep,
+      steps: parameters.numInferenceSteps,
+      promptEmbeddings: promptEmbeddings,
+      pooledPromptEmbeddings: pooledPromptEmbeddings,
+      latents: latents,
+      evaluateParameters: parameters,
+      transformer: transformer
+    )
+  }
 }
 
 public typealias ImageDecoder = (MLXArray) -> MLXArray
@@ -364,14 +404,15 @@ public struct DenoiseIterator: Sequence, IteratorProtocol {
   let transformer: MultiModalDiffusionTransformer
 
   init(
-    steps: Int, promptEmbeddings: MLXArray, pooledPromptEmbeddings: MLXArray, latents: MLXArray,
+    startStep: Int = 0, steps: Int, promptEmbeddings: MLXArray, pooledPromptEmbeddings: MLXArray,
+    latents: MLXArray,
     evaluateParameters: EvaluateParameters, transformer: MultiModalDiffusionTransformer
   ) {
     self.steps = steps
     self.promptEmbeddings = promptEmbeddings
     self.pooledPromptEmbeddings = pooledPromptEmbeddings
     self.latents = latents
-    self.i = 0
+    self.i = startStep
     self.evaluateParameters = evaluateParameters
     self.transformer = transformer
   }

Sources/FluxConfiguration.swift

@@ -146,6 +146,12 @@ public struct FluxConfiguration: Sendable {
     try factory(hub, self, configuration) as? TextToImageGenerator
   }
 
+  public func ImageToImageGenerator(hub: HubApi = HubApi(), configuration: LoadConfiguration)
+    throws -> ImageToImageGenerator?
+  {
+    try factory(hub, self, configuration) as? ImageToImageGenerator
+  }
+
   public static let flux1Schnell = FluxConfiguration(
     id: "black-forest-labs/FLUX.1-schnell",
     files: [