Flux2 Updates to Improve Image Quality

python/src/coreai_models/diffusion/flux2.py

@@ -201,7 +201,14 @@ def __init__(self, vae: torch.nn.Module) -> None:
         self.vae = self.vae.to(next(vae.parameters()).dtype)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return cast(torch.Tensor, self.vae.encode(x).latent_dist.parameters)
+        # diffusers encodes img2img reference images with
+        # `retrieve_latents(..., sample_mode="argmax")` -> `latent_dist.mode()`,
+        # i.e. the distribution MEAN (first `latent_channels` channels), not the
+        # raw `parameters` tensor (which is mean concat logvar = 2x channels).
+        # Returning `.parameters` would emit 64 channels where the pipeline
+        # expects 32, corrupting the img2img latents. `.mode()` is deterministic,
+        # so it is also the correct choice for a traced/exported graph.
+        return cast(torch.Tensor, self.vae.encode(x).latent_dist.mode())
 
 
 # ---------------------------------------------------------------------------

swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift

@@ -37,10 +37,44 @@ public struct Flux2Pipeline: DiffusionPipeline {
     private static let latentChannels = 128
     private static let textSeqLen = 512
     private static let defaultRopeTheta: Float = 2000.0
-    private static let schedulerBaseShift: Float = 0.5
-    private static let schedulerMaxShift: Float = 1.15
-    private static let schedulerBaseSeqLen: Float = 256
-    private static let schedulerMaxSeqLen: Float = 4096
+
+    /// FLUX.2 flow-matching timestep shift.
+    ///
+    /// Mirrors diffusers `compute_empirical_mu` (diffusers 0.37.1):
+    ///   pipelines/flux2/pipeline_flux2_klein.py:63-78
+    ///   (copied from pipelines/flux2/pipeline_flux2.py)
+    /// Call site — pipeline_flux2_klein.py:810-811:
+    ///   image_seq_len = latents.shape[1]
+    ///   mu = compute_empirical_mu(image_seq_len=image_seq_len, num_steps=num_inference_steps)
+    ///
+    /// Reference implementation:
+    ///   def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
+    ///       a1, b1 = 8.73809524e-05, 1.89833333
+    ///       a2, b2 = 0.00016927, 0.45666666
+    ///       if image_seq_len > 4300:
+    ///           mu = a2 * image_seq_len + b2
+    ///           return float(mu)
+    ///       m_200 = a2 * image_seq_len + b2
+    ///       m_10 = a1 * image_seq_len + b1
+    ///       a = (m_200 - m_10) / 190.0
+    ///       b = m_200 - 200.0 * a
+    ///       mu = a * num_steps + b
+    ///       return float(mu)
+    private static func computeEmpiricalMu(imageSeqLen: Int, numSteps: Int) -> Float {
+        let a1: Float = 8.73809524e-05
+        let b1: Float = 1.89833333
+        let a2: Float = 0.00016927
+        let b2: Float = 0.45666666
+        let seq = Float(imageSeqLen)
+        if imageSeqLen > 4300 {
+            return a2 * seq + b2
+        }
+        let m200 = a2 * seq + b2
+        let m10 = a1 * seq + b1
+        let a = (m200 - m10) / 190.0
+        let b = m200 - 200.0 * a
+        return a * Float(numSteps) + b
+    }
 
     /// Image size is determined by the mode selected at init.
     public var defaultImageSize: (width: Int, height: Int) {
@@ -121,11 +155,7 @@ public struct Flux2Pipeline: DiffusionPipeline {
         // For img2img: the schedule covers [strength → 0] using all steps, so every requested step
         // contributes to denoising and the noising sigma matches the first scheduled step exactly.
         // For txt2img: the schedule covers [1.0 → 0] as usual.
-        let baseShift: Float = Self.schedulerBaseShift
-        let maxShift: Float = Self.schedulerMaxShift
-        let baseSeqLen: Float = Self.schedulerBaseSeqLen
-        let maxSeqLen: Float = Self.schedulerMaxSeqLen
-        let mu = baseShift + (maxShift - baseShift) * (Float(seqLen) - baseSeqLen) / (maxSeqLen - baseSeqLen)
+        let mu = Self.computeEmpiricalMu(imageSeqLen: seqLen, numSteps: steps)
         let isActuallyImg2Img = configuration.isImageToImage && encoder != nil && configuration.startingImage != nil
         let sigmaMax: Float = isActuallyImg2Img ? configuration.strength : 1.0
         let scheduler = DiscreteFlowScheduler(
@@ -285,13 +315,26 @@ public struct Flux2Pipeline: DiffusionPipeline {
     private func encodeText(_ text: String) async throws -> [Float] {
         let seqLen = Self.textSeqLen
 
-        // Tokenize using Qwen3 chat template
+        // Tokenize using Qwen3 chat template.
+        //
+        // Must match diffusers `_get_qwen3_prompt_embeds`
+        // (diffusers 0.37.1, pipelines/flux2/pipeline_flux2_klein.py), which builds the
+        // input as:
+        //     messages = [{"role": "user", "content": single_prompt}]
+        //     text = tokenizer.apply_chat_template(
+        //         messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
+        //
+        // `enable_thinking=False` is significant for the Qwen3 template: it appends an
+        // empty `<think>\n\n</think>\n\n` block after the assistant prompt. Leaving it
+        // undefined omits that block, changing the trailing conditioning tokens and
+        // hurting prompt adherence. Pass it via additionalContext to match the reference.
         var ids: [Int]
         let messages: [[String: String]] = [["role": "user", "content": text]]
         do {
             ids = try tokenizer.applyChatTemplate(
                 messages: messages, chatTemplate: nil,
-                addGenerationPrompt: true, truncation: true, maxLength: seqLen, tools: nil
+                addGenerationPrompt: true, truncation: true, maxLength: seqLen, tools: nil,
+                additionalContext: ["enable_thinking": false]
             )
         } catch {
             let tokens = tokenizer.tokenize(text: text)
@@ -303,7 +346,14 @@ public struct Flux2Pipeline: DiffusionPipeline {
         }
 
         let realTokenCount = ids.count
-        let padTokenId = tokenizer.eosTokenId ?? 151643
+        // diffusers pads with the tokenizer's pad_token, not the eos_token. For
+        // FLUX.2 klein's Qwen tokenizer these differ: pad_token is <|endoftext|>
+        // (151643) while eos_token is <|im_end|> (151645). The reference builds
+        // input_ids via `tokenizer(text, padding="max_length", max_length=512)`
+        // (diffusers 0.37.1, pipeline_flux2_klein.py `_get_qwen3_prompt_embeds`),
+        // which uses pad_token. These ~490 padding tokens are fed to the DiT
+        // UNMASKED, so the id must match the reference exactly.
+        let padTokenId = tokenizer.convertTokenToId("<|endoftext|>") ?? 151643
 
         while ids.count < seqLen {
             ids.append(padTokenId)

swift/Sources/CoreAIDiffusionPipeline/Schedulers/DiscreteFlowScheduler.swift

@@ -37,7 +37,17 @@ public final class DiscreteFlowScheduler {
         self.mu = mu
         self.counter = 0
 
-        let sigmaMin: Float = (mu != nil) ? (1.0 / trainSteps) : (1.0 / Float(stepCount))
+        // Lower bound of the pre-shift sigma linspace. Diffusers builds
+        //   sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+        // (pipeline_flux2_klein.py) and passes it to
+        // FlowMatchEulerDiscreteScheduler.set_timesteps, which uses the provided
+        // sigmas as-is and only applies the mu/shift transform — it does NOT
+        // recompute the endpoints from num_train_timesteps. So the floor must be
+        // 1/stepCount for BOTH the dynamic-shift (mu) and static-shift paths.
+        // Using 1/trainStepCount here collapsed the final sigma to ~0 at low step
+        // counts (e.g. 4-step Klein), wasting the last step and misplacing all
+        // intermediate noise levels.
+        let sigmaMin: Float = 1.0 / Float(stepCount)
         var inferSigmas: [Float] = linspace(sigmaMax, sigmaMin, stepCount)
 
         if let mu {