(
&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_tensors = unpack_tensors(x)?; |
| 51 | let latent_model_input = &unpacked_tensors[0].to_dtype(ctx.dtype)?; |
| 52 | let text_embeddings = &unpacked_tensors[1].to_dtype(ctx.dtype)?; |
| 53 | |
| 54 | let timestep_tensor = &unpacked_tensors[2]; |
| 55 | let timestep_vec = timestep_tensor.to_vec1()?; |
| 56 | let timestep_f32: &f32 = timestep_vec.first().expect("Error retrieving timestep"); |
| 57 | |
| 58 | info!("UNet model forwarding..."); |
| 59 | |
| 60 | Ok(self |
| 61 | .unet_model |
| 62 | .forward(latent_model_input, *timestep_f32 as f64, text_embeddings) |
| 63 | .expect("Error running UNet forward")) |
| 64 | } |
| 65 | |
| 66 | async fn forward_mut( |
| 67 | &mut self, |
no test coverage detected