| 81 | impl UNet { |
| 82 | #[allow(clippy::too_many_arguments)] |
| 83 | pub fn load_model( |
| 84 | name: Option<String>, |
| 85 | use_flash_attn: bool, |
| 86 | version: StableDiffusionVersion, |
| 87 | use_f16: bool, |
| 88 | device: &Device, |
| 89 | dtype: DType, |
| 90 | cache_dir: String, |
| 91 | config: &StableDiffusionConfig, |
| 92 | ) -> anyhow::Result<Box<Self>> |
| 93 | where |
| 94 | Self: Sized, |
| 95 | { |
| 96 | let unet_weights = ModelFile::Unet.get(name, version, use_f16, cache_dir)?; |
| 97 | let unet = config.build_unet(unet_weights, device, 4, use_flash_attn, dtype)?; |
| 98 | |
| 99 | info!("Loading UNet model..."); |
| 100 | |
| 101 | Ok(Box::new(Self { unet_model: unet })) |
| 102 | } |
| 103 | |
| 104 | pub async fn forward_unpacked( |
| 105 | forwarder: &mut Box<dyn Forwarder>, |