mistralrs_core/pipeline/loaders/
diffusion_loaders.rs

1use std::{
2    fmt::Debug,
3    path::{Path, PathBuf},
4    str::FromStr,
5};
6
7use anyhow::{Context, Result};
8use candle_core::{Device, Tensor};
9
10use hf_hub::api::sync::ApiRepo;
11use mistralrs_quant::ShardedVarBuilder;
12#[cfg(feature = "pyo3_macros")]
13use pyo3::pyclass;
14
15use regex::Regex;
16use serde::Deserialize;
17
18use tracing::info;
19
20use super::{ModelPaths, NormalLoadingMetadata};
21use crate::{
22    api_dir_list, api_get_file,
23    diffusion_models::{
24        flux::{
25            self,
26            stepper::{FluxStepper, FluxStepperConfig},
27        },
28        DiffusionGenerationParams,
29    },
30    paged_attention::AttentionImplementation,
31    pipeline::paths::AdapterPaths,
32};
33
34pub trait DiffusionModel {
35    /// This returns a tensor of shape (bs, c, h, w), with values in [0, 255].
36    fn forward(
37        &mut self,
38        prompts: Vec<String>,
39        params: DiffusionGenerationParams,
40    ) -> candle_core::Result<Tensor>;
41    fn device(&self) -> &Device;
42    fn max_seq_len(&self) -> usize;
43}
44
45pub trait DiffusionModelLoader: Send + Sync {
46    /// If the model is being loaded with `load_model_from_hf` (so manual paths not provided), this will be called.
47    fn get_model_paths(&self, api: &ApiRepo, model_id: &Path) -> Result<Vec<PathBuf>>;
48    /// If the model is being loaded with `load_model_from_hf` (so manual paths not provided), this will be called.
49    fn get_config_filenames(&self, api: &ApiRepo, model_id: &Path) -> Result<Vec<PathBuf>>;
50    fn force_cpu_vb(&self) -> Vec<bool>;
51    // `configs` and `vbs` should be corresponding. It is up to the implementer to maintain this invaraint.
52    fn load(
53        &self,
54        configs: Vec<String>,
55        vbs: Vec<ShardedVarBuilder>,
56        normal_loading_metadata: NormalLoadingMetadata,
57        attention_mechanism: AttentionImplementation,
58        silent: bool,
59    ) -> Result<Box<dyn DiffusionModel + Send + Sync>>;
60}
61
62#[cfg_attr(feature = "pyo3_macros", pyclass(eq, eq_int))]
63#[derive(Clone, Debug, Deserialize, PartialEq)]
64/// The architecture to load the vision model as.
65pub enum DiffusionLoaderType {
66    #[serde(rename = "flux")]
67    Flux,
68    #[serde(rename = "flux-offloaded")]
69    FluxOffloaded,
70}
71
72impl FromStr for DiffusionLoaderType {
73    type Err = String;
74    fn from_str(s: &str) -> Result<Self, Self::Err> {
75        match s {
76            "flux" => Ok(Self::Flux),
77            "flux-offloaded" => Ok(Self::FluxOffloaded),
78            a => Err(format!(
79                "Unknown architecture `{a}`. Possible architectures: `flux`."
80            )),
81        }
82    }
83}
84
85#[derive(Clone, Debug)]
86pub struct DiffusionModelPathsInner {
87    pub config_filenames: Vec<PathBuf>,
88    pub filenames: Vec<PathBuf>,
89}
90
91#[derive(Clone, Debug)]
92pub struct DiffusionModelPaths(pub DiffusionModelPathsInner);
93
94impl ModelPaths for DiffusionModelPaths {
95    fn get_config_filename(&self) -> &PathBuf {
96        unreachable!("Use `std::any::Any`.")
97    }
98    fn get_tokenizer_filename(&self) -> &PathBuf {
99        unreachable!("Use `std::any::Any`.")
100    }
101    fn get_weight_filenames(&self) -> &[PathBuf] {
102        unreachable!("Use `std::any::Any`.")
103    }
104    fn get_template_filename(&self) -> &Option<PathBuf> {
105        unreachable!("Use `std::any::Any`.")
106    }
107    fn get_gen_conf_filename(&self) -> Option<&PathBuf> {
108        unreachable!("Use `std::any::Any`.")
109    }
110    fn get_preprocessor_config(&self) -> &Option<PathBuf> {
111        unreachable!("Use `std::any::Any`.")
112    }
113    fn get_processor_config(&self) -> &Option<PathBuf> {
114        unreachable!("Use `std::any::Any`.")
115    }
116    fn get_chat_template_explicit(&self) -> &Option<PathBuf> {
117        unreachable!("Use `std::any::Any`.")
118    }
119    fn get_adapter_paths(&self) -> &AdapterPaths {
120        unreachable!("Use `std::any::Any`.")
121    }
122}
123
124// ======================== Flux loader
125
126/// [`DiffusionLoader`] for a Flux Diffusion model.
127///
128/// [`DiffusionLoader`]: https://ericlbuehler.github.io/mistral.rs/mistralrs/struct.DiffusionLoader.html
129pub struct FluxLoader {
130    pub(crate) offload: bool,
131}
132
133impl DiffusionModelLoader for FluxLoader {
134    fn get_model_paths(&self, api: &ApiRepo, model_id: &Path) -> Result<Vec<PathBuf>> {
135        let regex = Regex::new(r"^flux\d+-(schnell|dev)\.safetensors$")?;
136        let flux_name = api_dir_list!(api, model_id)
137            .filter(|x| regex.is_match(x))
138            .nth(0)
139            .with_context(|| "Expected at least 1 .safetensors file matching the FLUX regex, please raise an issue.")?;
140        let flux_file = api_get_file!(api, &flux_name, model_id);
141        let ae_file = api_get_file!(api, "ae.safetensors", model_id);
142
143        // NOTE(EricLBuehler): disgusting way of doing this but the 0th path is the flux, 1 is ae
144        Ok(vec![flux_file, ae_file])
145    }
146    fn get_config_filenames(&self, api: &ApiRepo, model_id: &Path) -> Result<Vec<PathBuf>> {
147        let flux_file = api_get_file!(api, "transformer/config.json", model_id);
148        let ae_file = api_get_file!(api, "vae/config.json", model_id);
149
150        // NOTE(EricLBuehler): disgusting way of doing this but the 0th path is the flux, 1 is ae
151        Ok(vec![flux_file, ae_file])
152    }
153    fn force_cpu_vb(&self) -> Vec<bool> {
154        vec![self.offload, false]
155    }
156    fn load(
157        &self,
158        mut configs: Vec<String>,
159        mut vbs: Vec<ShardedVarBuilder>,
160        normal_loading_metadata: NormalLoadingMetadata,
161        _attention_mechanism: AttentionImplementation,
162        silent: bool,
163    ) -> Result<Box<dyn DiffusionModel + Send + Sync>> {
164        let (vae_cfg, vae_vb) = (configs.remove(1), vbs.remove(1));
165        let (flux_cfg, flux_vb) = (configs.remove(0), vbs.remove(0));
166
167        let vae_cfg: flux::autoencoder::Config = serde_json::from_str(&vae_cfg)?;
168        let flux_cfg: flux::model::Config = serde_json::from_str(&flux_cfg)?;
169
170        let flux_dtype = flux_vb.dtype();
171        if flux_dtype != vae_vb.dtype() {
172            anyhow::bail!(
173                "Expected VAE and FLUX model VBs to be the same dtype, got {:?} and {flux_dtype:?}",
174                vae_vb.dtype()
175            );
176        }
177
178        Ok(Box::new(FluxStepper::new(
179            FluxStepperConfig::default_for_guidance(flux_cfg.guidance_embeds),
180            (flux_vb, &flux_cfg),
181            (vae_vb, &vae_cfg),
182            flux_dtype,
183            &normal_loading_metadata.real_device,
184            silent,
185            self.offload,
186        )?))
187    }
188}
mistralrs_core/pipeline/loaders/diffusion_loaders.rs

mistralrs_core/pipeline/loaders/
diffusion_loaders.rs
