Merge branch 'main' into sukru/input-handler-protocol

Author: stikves

swift/Sources/CoreAIDiffusionPipeline/Pipelines/Flux2Pipeline.swift

@@ -37,6 +37,7 @@ public struct Flux2Pipeline: DiffusionPipeline {
     private static let latentChannels = 128
     private static let textSeqLen = 512
     private static let defaultRopeTheta: Float = 2000.0
+    private static let qwen3PadTokenId = 151643
 
     /// FLUX.2 flow-matching timestep shift.
     ///
@@ -113,6 +114,12 @@ public struct Flux2Pipeline: DiffusionPipeline {
         self.batchNormMean = batchNormMean
         self.batchNormVar = batchNormVar
         self.batchNormEps = batchNormEps
+
+        if tokenizer.convertTokenToId("<|endoftext|>") == nil {
+            CLILogger.log(
+                "⚠️ Flux2Pipeline: tokenizer has no <|endoftext|> token, using Qwen3 fallback pad ID",
+                component: "Diffusion")
+        }
     }
 
     // MARK: - ResourceManaging
@@ -353,7 +360,7 @@ public struct Flux2Pipeline: DiffusionPipeline {
         // (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
+        let padTokenId = tokenizer.convertTokenToId("<|endoftext|>") ?? Self.qwen3PadTokenId
 
         while ids.count < seqLen {
             ids.append(padTokenId)

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAIPipelinedEngine.swift

@@ -688,6 +688,7 @@ private struct EngineImpl: ~Copyable {
     // MARK: - Sampler
 
     private mutating func getOrCreateSampler(for config: SamplingConfiguration) throws -> any MPSGraphSampler {
+        let config = config.normalized()
         let temperature = config.temperature
 
         if let existingSampler = cachedSampler, let existingTemp = cachedSamplerTemperature {

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialEngine.swift

@@ -538,7 +538,7 @@ extension CoreAISequentialEngine.GenerationSequence {
             generationToken: GenerationToken
         ) {
             self.engine = engine
-            self.samplingConfiguration = samplingConfiguration
+            self.samplingConfiguration = samplingConfiguration.normalized()
             self.returnsLogits = inferenceOptions.includeLogits
             self.forcedContinuation = inferenceOptions.forcedContinuation
             self.stopReasonStore = stopReasonStore

swift/Sources/CoreAILanguageModels/InferenceEngines/CoreAISequentialVLMEngine.swift

@@ -1009,7 +1009,7 @@ extension CoreAISequentialVLMEngine.GenerationSequence {
             generationToken: GenerationToken
         ) {
             self.engine = engine
-            self.samplingConfiguration = samplingConfiguration
+            self.samplingConfiguration = samplingConfiguration.normalized()
             self.returnsLogits = inferenceOptions.includeLogits
             self.forcedContinuation = inferenceOptions.forcedContinuation
             self.stopReasonStore = stopReasonStore

swift/Sources/CoreAILanguageModels/ModelShapeConfig.swift

@@ -308,31 +308,28 @@ public struct LogitsWriter {
     /// Handle all logits output based on provided options
     /// - Parameters:
     ///   - logits: Array of logits for each generated token
-    ///   - generatedText: The generated text
+    ///   - generatedTokenIds: Token IDs produced during generation (one per logits vector)
     ///   - tokenizer: Tokenizer for decoding token IDs
     ///   - saveLogitsLength: Specifies whether to save top-K or full logits
     ///   - saveJsonPath: Optional path to save logits as JSON
     ///   - printToConsole: Whether to print top-5 to console
     public static func handleOutput<T: BinaryFloatingPoint>(
         logits: [[T]],
-        generatedText: String,
+        generatedTokenIds: [Int],
         tokenizer: any Tokenizer,
         saveLogitsLength: LogitsLength,
         saveJsonPath: String?,
         printToConsole: Bool
     ) throws {
-        // Decode generated text to get token IDs
-        let generatedTokens = tokenizer.encode(text: generatedText)
-
         // Token count must match logits count - each generation step produces one token and one logits vector
-        guard generatedTokens.count == logits.count else {
+        guard generatedTokenIds.count == logits.count else {
             throw LogitsWriterError.tokenCountMismatch(
-                tokenCount: generatedTokens.count,
+                tokenCount: generatedTokenIds.count,
                 logitsCount: logits.count
             )
         }
 
-        let tokenCount = generatedTokens.count
+        let tokenCount = generatedTokenIds.count
 
         // Only build top-K array if we need it for printing or saving top-K
         let needsTopKArray = printToConsole || !saveLogitsLength.isFull
@@ -343,7 +340,7 @@ public struct LogitsWriter {
             if let jsonPath = saveJsonPath {
                 try saveFullJSON(
                     logits: logits,
-                    generatedTokens: generatedTokens,
+                    generatedTokens: generatedTokenIds,
                     tokenizer: tokenizer,
                     path: jsonPath
                 )
@@ -355,8 +352,8 @@ public struct LogitsWriter {
         let topK = saveLogitsLength.topKForConsole
         for index in 0..<tokenCount {
             let logitVector = logits[index]
-            let tokenId = Int32(generatedTokens[index])
-            let tokenText = tokenizer.decode(tokens: [generatedTokens[index]])
+            let tokenId = Int32(generatedTokenIds[index])
+            let tokenText = tokenizer.decode(tokens: [generatedTokenIds[index]])
 
             // Extract top-K logits
             let topLogits = extractTopK(

swift/Sources/CoreAILanguageModels/Output/LogitsWriter.swift

@@ -226,6 +226,14 @@ public struct SamplingConfiguration: Sendable, Equatable, Hashable {
     ///
     /// - Returns: A new configuration with redundant settings removed.
     public func normalized() -> SamplingConfiguration {
+        // topK=1 with temperature>0 is equivalent to greedy — promote it
+        if let k = topK, k == 1, temperature > 0 {
+            CLILogger.log(
+                "⚠️ SamplingConfiguration: topK=1 normalized to greedy (temperature=0)",
+                component: "Sampling")
+            return .greedy
+        }
+
         let effectiveTopK: Int?
         let effectiveTopP: Double?
         let effectiveMinP: Double?

swift/Sources/CoreAILanguageModels/Samplers/SamplingConfiguration.swift

@@ -65,13 +65,14 @@ public class TextGenerator {
     ///   - prompt: Input text prompt
     ///   - maxTokens: Maximum number of tokens to generate
     ///   - stopSequences: Stop token sequences that halt generation. If nil, uses tokenizer's EOS tokens and common fallbacks.
-    /// - Returns: Tuple containing generated text and array of logits for each token
+    /// - Returns: Tuple containing generated text, token IDs, and array of logits for each token
     public func generateWithLogits(
         input: Input,
         maxTokens: Int = 50,
         stopSequences: StopSequences? = nil
-    ) async throws -> (text: String, logits: [[LogitsScalarType]]) {
+    ) async throws -> (text: String, tokenIds: [Int], logits: [[LogitsScalarType]]) {
         var textParts: [String] = []
+        var allTokenIds: [Int] = []
         var allLogits: [[LogitsScalarType]] = []
 
         // Use provided stop sequences, or create default ones from tokenizer
@@ -89,11 +90,12 @@ public class TextGenerator {
         for try await result in resultStream {
             textParts.append(result.text)
             if let logits = result.rawLogits {
+                allTokenIds.append(Int(result.tokenId))
                 allLogits.append(logits)
             }
         }
 
-        return (text: textParts.joined(), logits: allLogits)
+        return (text: textParts.joined(), tokenIds: allTokenIds, logits: allLogits)
     }
 
     /// Evaluate continuation probability (no generation)

swift/Sources/CoreAILanguageModels/TextGeneration/TextGenerator.swift

@@ -18,8 +18,6 @@ public struct MelConfig: Sendable {
     public let nMelBins: Int
     public let nFrames: Int
 
-    public var nSamples: Int { Int(sampleRate) * (nFrames * hopLength / Int(sampleRate / 100)) }
-
     /// Whisper / Parakeet shared parameters.
     public static let whisper = MelConfig(
         sampleRate: 16_000, nFFT: 400, hopLength: 160, nMelBins: 128, nFrames: 3_000)

swift/Sources/CoreAISpeech/MelSpectrogram.swift

@@ -637,6 +637,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
             let needsLogits = saveLogits != nil || printLogits
             let generatedText: String
             var allLogits: [[LogitsScalarType]] = []
+            var allTokenIds: [Int] = []
 
             if needsLogits {
                 // Generate with logits
@@ -646,6 +647,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
                     stopSequences: stopSequences
                 )
                 generatedText = result.text
+                allTokenIds = result.tokenIds
                 allLogits = result.logits
             } else {
                 // Standard generation without logits
@@ -681,7 +683,7 @@ struct LLMRunner: AsyncParsableCommand, Sendable {
             if needsLogits && !allLogits.isEmpty {
                 try LogitsWriter.handleOutput(
                     logits: allLogits,
-                    generatedText: generatedOnly,
+                    generatedTokenIds: allTokenIds,
                     tokenizer: tokenizer,
                     saveLogitsLength: saveLogitsLength,
                     saveJsonPath: saveLogits,