(
&self,
x: &Tensor,
_index_pos: usize,
_block_idx: usize,
_ctx: &mut Context,
)
| 41 | } |
| 42 | |
| 43 | async fn forward( |
| 44 | &self, |
| 45 | x: &Tensor, |
| 46 | _index_pos: usize, |
| 47 | _block_idx: usize, |
| 48 | _ctx: &mut Context, |
| 49 | ) -> anyhow::Result<Tensor> { |
| 50 | let unpacked = unpack_tensors(x)?; |
| 51 | |
| 52 | // New interface: [dims(h_half, w_half), patchified_latents] |
| 53 | // Old interface: [direction_scalar, latent_image] — direction=0 means decode |
| 54 | if unpacked[0].dims().len() == 1 && unpacked[0].dim(0)? == 2 && unpacked.len() == 2 { |
| 55 | // New packed format: BN denorm + unpatchify + decode |
| 56 | let dims = unpacked[0].to_vec1::<f32>()?; |
| 57 | let h_half = dims[0] as usize; |
| 58 | let w_half = dims[1] as usize; |
| 59 | let img = unpacked[1].to_dtype(DType::F32)?; |
| 60 | |
| 61 | // BN denormalization |
| 62 | let bn_eps = 0.0001_f64; |
| 63 | let bn_std = self.bn_running_var |
| 64 | .to_dtype(DType::F32)? |
| 65 | .broadcast_add(&Tensor::new(&[bn_eps as f32], self.bn_running_var.device())?)? |
| 66 | .sqrt()?; |
| 67 | let img = img |
| 68 | .broadcast_mul(&bn_std.unsqueeze(0)?.unsqueeze(0)?)? |
| 69 | .broadcast_add(&self.bn_running_mean.to_dtype(DType::F32)?.unsqueeze(0)?.unsqueeze(0)?)?; |
| 70 | |
| 71 | // Unpatchify: (b, h*w, 128) → (b, 32, h*2, w*2) |
| 72 | let img = img |
| 73 | .reshape((1, h_half, w_half, 128))? |
| 74 | .permute((0, 3, 1, 2))?; |
| 75 | let img = img |
| 76 | .reshape((1, 32, 2, 2, h_half, w_half))? |
| 77 | .permute((0, 1, 4, 2, 5, 3))? |
| 78 | .reshape((1, 32, h_half * 2, w_half * 2))?; |
| 79 | |
| 80 | Ok(self.model.decode(&img)?) |
| 81 | } else { |
| 82 | // Legacy interface: [direction, input] |
| 83 | let direction_vec = unpacked[0].to_vec1::<f32>()?; |
| 84 | let direction = direction_vec[0]; |
| 85 | let input = &unpacked[1].to_dtype(DType::F32)?; |
| 86 | |
| 87 | debug!("FluxVAE forwarding (direction={direction})..."); |
| 88 | |
| 89 | if direction == 1.0 { |
| 90 | anyhow::bail!("FluxVAE encode not implemented") |
| 91 | } else { |
| 92 | Ok(self.model.decode(input)?) |
| 93 | } |
| 94 | } |
| 95 | } |
| 96 | |
| 97 | async fn forward_mut( |
| 98 | &mut self, |
no test coverage detected