Fix SD text encoder crash: infer sequence length from model (#103)

* Fix SD text encoder crash: infer sequence length from model

StableDiffusionPipeline.encodeText() passed raw token count as the
sequence dimension, but CLIP text encoders are exported with a fixed
seq_len (77). When a prompt tokenizes to a different length,
resolvingDynamicDimensions crashes.

Fix: use CoreAITextEncoder.encode() which pads/truncates correctly,
and infer the sequence length from the model input descriptor at
load time instead of hardcoding 77.

Fixes #102.

* Add unit tests for text encoder token padding and truncation

Author: stikves

swift/Sources/CoreAIDiffusionPipeline/Components/CoreAIDiffusionModelFunction.swift

@@ -223,6 +223,17 @@ public actor CoreAIDiffusionModelFunction {
             return result
         }
     }
+
+    /// Infer the sequence length from the first input's shape (dim 1).
+    /// Returns nil if the model isn't loaded or has no rank-2 input.
+    public func inferSequenceLength() async throws -> Int? {
+        let descs = try await inputDescriptors
+        guard let desc = descs.values.first, desc.shape.count >= 2 else {
+            return nil
+        }
+        let dim = desc.shape[1]
+        return dim > 0 ? dim : nil
+    }
 }
 
 // MARK: - Errors

swift/Sources/CoreAIDiffusionPipeline/Pipelines/PipelineDescriptor+CoreAI.swift

@@ -113,13 +113,15 @@ extension PipelineDescriptor {
             throw PipelineLoadError.missingComponent("text_encoder")
         }
 
+        let textEncoderSeqLength = try await textEncoderFunction.inferSequenceLength() ?? 77
+
         let textEncoder = CoreAITextEncoder(
             function: textEncoderFunction,
             tokenize: { text in
-                let (_, ids) = tokenizer.tokenize(input: text, minCount: 77)
+                let (_, ids) = tokenizer.tokenize(input: text, minCount: textEncoderSeqLength)
                 return ids.map(Int32.init)
             },
-            maxLength: 77
+            maxLength: textEncoderSeqLength
         )
 
         let denoiser = CoreAIDenoiser(function: unetFunction)

swift/Sources/CoreAIDiffusionPipeline/Pipelines/StableDiffusionPipeline.swift

@@ -172,9 +172,10 @@ public struct StableDiffusionPipeline: DiffusionPipeline {
     // MARK: - Private Helpers
 
     private func encodeText(_ text: String) async throws -> [Float] {
-        let tokenize = components.textEncoder.tokenize
-        let ids = tokenize(text)
-        return try await components.textEncoder.function.run(intInputs: [(ids, [1, ids.count])])
+        let output = try await components.textEncoder.encode(text)
+        let shape = output.hiddenStates.shape
+        let count = shape.reduce(1, *)
+        return readNDArray(output.hiddenStates, as: Float.self, count: count)
     }
 
     private func runDenoiser(

swift/Tests/DiffusionPipelineTests/DiffusionPipelineTests.swift

@@ -45,3 +45,36 @@ struct DiffusionPipelineTests {
         }
     }
 }
+
+@Suite("TextEncoder tokenize padding")
+struct TextEncoderTokenizeTests {
+    @Test("Tokenize closure pads short input to maxLength")
+    func padsShortInput() {
+        let maxLength = 77
+        let tokenize: @Sendable (String) -> [Int32] = { _ in
+            var ids = [Int32](1...10)
+            if ids.count < maxLength {
+                ids += [Int32](repeating: 0, count: maxLength - ids.count)
+            }
+            return ids
+        }
+        let ids = tokenize("short prompt")
+        #expect(ids.count == 77)
+        #expect(ids.last == 0)
+    }
+
+    @Test("Tokenize closure truncates long input to maxLength")
+    func truncatesLongInput() {
+        let maxLength = 77
+        let tokenize: @Sendable (String) -> [Int32] = { _ in
+            var ids = [Int32](1...100)
+            if ids.count > maxLength {
+                ids = Array(ids.prefix(maxLength))
+            }
+            return ids
+        }
+        let ids = tokenize("very long prompt that produces many tokens")
+        #expect(ids.count == 77)
+        #expect(ids[76] == 77)
+    }
+}