Method encode

cake-core/src/models/flux/text_encoder.rs:378–393

Run the full encoder forward pass, returning hidden states. FLUX.2-klein extracts hidden states from layers [9, 18, 27] (0-indexed) and concatenates them: 3 × 2560 = 7680 = joint_attention_dim. Encode with attention mask. attn_mask: (1, seq) with 1=real, 0=padding.

(&self, token_ids: &Tensor, attn_mask: Option<&Tensor>)

    /// and concatenates them: 3 × 2560 = 7680 = joint_attention_dim.
    /// Encode with attention mask. attn_mask: (1, seq) with 1=real, 0=padding.
    pub fn encode(&self, token_ids: &Tensor, attn_mask: Option<&Tensor>) -> Result<Tensor> {
        const OUTPUT_LAYERS: [usize; 3] = [8, 17, 26];

        let mut x = self.backend.embedding(token_ids, &self.embeddings_weight)?;
        let mut layer_outputs: Vec<Tensor> = Vec::new();

        for (i, block) in self.blocks.iter().enumerate() {
            x = block.forward_with_mask(&x, attn_mask)?;
            if OUTPUT_LAYERS.contains(&i) {
                layer_outputs.push(x.clone());
            }
        }

        // Concatenate along feature dimension: (b, seq, 2560*3) = (b, seq, 7680)
        Tensor::cat(&layer_outputs, candle_core::D::Minus1)
    }
}