Fix invalid mask implementation

apple · Sep 23, 2024 · 4361735 · 4361735
1 parent cab24e9
commit 4361735
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/swift/StableDiffusion/pipeline/TextEncoderT5.swift b/swift/StableDiffusion/pipeline/TextEncoderT5.swift
@@ -91,9 +91,9 @@ public struct TextEncoderT5: TextEncoderT5Model {
 
         let attentionMaskName = "attention_mask"
         var attentionMask: [Int] = inputIds.map { token in
-            token == padToken ? maskToken : padToken
+            token == padToken ? padToken : maskToken
         }
-        attentionMask[0] = bosToken
+        attentionMask[0] = eosToken
 
         let floatIds = inputIds.map { Float32($0) }
         let floatMask = attentionMask.map { Float32($0) }