(
&self,
x: &Tensor,
index_pos: usize, // used as real_len for attention mask
_block_idx: usize,
ctx: &mut Context,
)
| 268 | } |
| 269 | |
| 270 | async fn forward( |
| 271 | &self, |
| 272 | x: &Tensor, |
| 273 | index_pos: usize, // used as real_len for attention mask |
| 274 | _block_idx: usize, |
| 275 | ctx: &mut Context, |
| 276 | ) -> anyhow::Result<Tensor> { |
| 277 | info!("FluxTextEncoder forwarding..."); |
| 278 | |
| 279 | // x is token IDs: (batch, seq_len) — padded to max_length |
| 280 | // index_pos carries the real (unpadded) token count for the attention mask |
| 281 | let seq_len = x.dim(1)?; |
| 282 | let attn_mask = if index_pos > 0 && index_pos < seq_len { |
| 283 | // Build attention mask: 1 for real tokens, 0 for padding |
| 284 | let mut mask_data = vec![1.0f32; index_pos]; |
| 285 | mask_data.resize(seq_len, 0.0); |
| 286 | Some(Tensor::new(mask_data.as_slice(), x.device())?.unsqueeze(0)?) |
| 287 | } else { |
| 288 | None |
| 289 | }; |
| 290 | |
| 291 | let hidden = self.encode(x, attn_mask.as_ref())?; |
| 292 | |
| 293 | let tensors = vec![hidden]; |
| 294 | let packed = pack_tensors(tensors, &ctx.device)?; |
| 295 | Ok(packed) |
| 296 | } |
| 297 | |
| 298 | async fn forward_mut( |
| 299 | &mut self, |
no test coverage detected