(
&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 | info!("VAE model forwarding..."); |
| 51 | |
| 52 | let unpacked_tensors = unpack_tensors(x)?; |
| 53 | |
| 54 | let direction_tensor = &unpacked_tensors[0]; |
| 55 | let direction_vec = direction_tensor.to_vec1()?; |
| 56 | let direction_f32: f32 = *direction_vec |
| 57 | .first() |
| 58 | .expect("Error retrieving direction info"); |
| 59 | |
| 60 | let input = &unpacked_tensors[1].to_dtype(ctx.dtype)?; |
| 61 | |
| 62 | debug!("VAE tensors decoded."); |
| 63 | |
| 64 | if direction_f32 == 1.0 { |
| 65 | let dist = self.vae_model.encode(input)?; |
| 66 | Ok(dist.sample()?) |
| 67 | } else { |
| 68 | Ok(self.vae_model.decode(input)?) |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | async fn forward_mut( |
| 73 | &mut self, |
no test coverage detected