
mistral.rs 扩散模型图像生成实战用 Rust 驱动 FLUX.1-schnell 完成文生图推理【免费下载链接】mistral.rsFast, flexible LLM inference项目地址: https://gitcode.com/GitHub_Trending/mi/mistral.rs本文基于 mistral.rs 官方示例mistralrs/examples/models/diffusion/main.rs对应文档 docs/src/content/docs/examples/rust/models/diffusion.md讲解如何用 Rust 调用DiffusionModelBuilder加载 FLUX 扩散模型并生成图像。读者将掌握扩散模型加载器类型、DiffusionGenerationParams参数、ImageGenerationResponseFormat输出格式等核心 API并能直接运行一个可复现的文生图程序。示例概览与运行方式该示例的目标非常直接加载black-forest-labs/FLUX.1-schnell模型输入一段文本提示词输出一张 720x1280 的图像。运行命令为cargo run --release --example diffusion -p mistralrs从源码结构看示例位于mistralrs/examples/models/diffusion/main.rs与文档中展示的代码完全一致。整段代码仅约 40 行覆盖了扩散模型推理的三个核心环节构建模型builder→ 发起生成请求 → 打印结果。代码开头引入了三个关键类型use mistralrs::{ DiffusionGenerationParams, DiffusionLoaderType, DiffusionModelBuilder, ImageGenerationResponseFormat, };DiffusionModelBuilder用于配置并加载扩散模型DiffusionLoaderType指定扩散模型架构的加载方式DiffusionGenerationParams控制生成图像的分辨率ImageGenerationResponseFormat控制生成结果的返回格式。模型加载DiffusionModelBuilder 与 DiffusionLoaderType构建与加载流程示例通过以下代码完成模型加载let model DiffusionModelBuilder::new( black-forest-labs/FLUX.1-schnell, DiffusionLoaderType::FluxOffloaded, ) .with_logging() .build() .await?;DiffusionModelBuilder::new接收两个必填参数参数说明model_idHugging Face 上的模型仓库标识例如black-forest-labs/FLUX.1-schnellloader_type扩散模型加载架构当前支持DiffusionLoaderType::Flux与DiffusionLoaderType::FluxOffloadedDiffusionModelBuilder的实现位于 mistralrs/src/diffusion_model.rs。其构造时会应用一组默认值dtype 为ModelDType::Auto自动选择、token 来源为TokenSource::CacheToken读取~/.cache/huggingface/token、最大并发序列数为 32、不强制 CPU、关闭日志。build()内部会调用build_diffusion_pipeline完成 pipeline 的组装再通过build_model_from_pipeline返回一个可直接使用的Model实例。可选的配置方法DiffusionModelBuilder还提供以下链式配置方法见 diffusion_model.rs方法作用with_dtype(dtype)以指定精度如ModelDType::F32、ModelDType::BF16加载模型默认Autowith_force_cpu()强制使用 CPU 设备注释明确指出不要与 PagedAttention 同时使用with_token_source(source)指定 Hugging Face token 来源如TokenSource::CacheToken、TokenSource::Env、TokenSource::Literalwith_hf_revision(rev)指定 Hugging Face 远端模型的 revisionwith_max_num_seqs(n)设置最多同时运行的序列数默认 32with_logging()开启日志输出Flux 与 FluxOffloaded 的区别DiffusionLoaderType定义于 mistralrs-core/src/pipeline/loaders/diffusion_loaders.rs目前包含两个变体Flux常规加载Transformer 与 VAE 全部放入 GPUFluxOffloaded将 FLUX Transformer 部分 offload 到 CPUVAE 仍保留在 GPU 上。其底层机制可以从 loader 的force_cpu_vb()方法看到vec![self.offload, false]—— 第一个元素对应 FLUX 权重offload 时强制 CPU VarBuilder第二个元素对应 VAE 权重始终为 false即留在 GPU。这在显存不足的环境中非常实用。自动检测逻辑DiffusionLoaderType::auto_detect_from_files会根据仓库文件列表自动识别是否为 FLUX 模型要求存在transformer/config.json、vae/config.json、ae.safetensors且存在匹配^flux\d-(schnell|dev)\.safetensors$正则的权重文件。这一逻辑意味着后续新增其他扩散架构如 SDXL时只需扩展该检测函数即可。在load()实现中loader 会分别下载 FLUX 权重文件匹配flux\d-(schnell|dev).safetensors与ae.safetensors自编码器即 VAE并校验 FLUX 与 VAE 的 dtype 一致后构造FluxStepper完成采样。生成请求generate_image 与生成参数请求调用let response model .generate_image( A vibrant sunset in the mountains, 4k, high quality..to_string(), ImageGenerationResponseFormat::Url, DiffusionGenerationParams::default(), None, ) .await?;Model::generate_image定义于 mistralrs/src/model.rs签名如下pub async fn generate_image( self, prompt: impl ToString, response_format: ImageGenerationResponseFormat, generation_params: DiffusionGenerationParams, save_file: OptionPathBuf, ) - crate::error::ResultImageGenerationResponse四个参数分别是提示词文本、响应格式、生成参数、可选的文件保存路径。此外还有generate_image_with_model变体多一个model_id: Optionstr参数用于在加载了多个模型时指定使用哪个模型生成None表示使用默认模型。生成参数 DiffusionGenerationParamsDiffusionGenerationParams定义于 mistralrs-core/src/diffusion_models/mod.rspub struct DiffusionGenerationParams { pub height: usize, pub width: usize, }目前只有height和width两个字段控制输出图像的分辨率。Default实现固定为720x1280竖版比例。如果需要生成横向或其他比例的图像可以手动构造use mistralrs::DiffusionGenerationParams; let params DiffusionGenerationParams { height: 1024, width: 1024, };响应格式 ImageGenerationResponseFormat该枚举定义于 mistralrs-core/src/request.rs包含两个变体变体说明ImageGenerationResponseFormat::Url图像保存为文件后返回 URL 字符串ImageGenerationResponseFormat::B64Json返回 Base64 编码的图像数据JSON 格式在Request::Normal的RequestMessage::ImageGeneration变体request.rs中请求会携带prompt、format、generation_params与可选的save_file字段随后进入 pipeline 由DiffusionModel::forward真正执行采样。结果处理耗时统计与输出生成完成后示例统计耗时并打印图像保存位置let finished Instant::now(); println!( Done! Took {} s. Image saved at: {}, finished.duration_since(start).as_secs_f32(), response.data[0].url.as_ref().unwrap() );response.data[0].url是OptionString当使用Url格式且生成成功时包含图像的本地文件路径若使用B64Json格式则相应字段存放 Base64 数据。使用as_ref().unwrap()之前建议先判空避免生成失败时 panic。底层原理FluxStepper 与采样参数生成过程的核心是 mistralrs-core/src/diffusion_models/flux/stepper.rs 中的FluxStepper。它会同时加载两个文本编码器T5 与 CLIP对提示词编码再由 FLUX Transformer 与 VAE 自编码器完成去噪与解码。采样步数由FluxStepperConfig::default_for_guidancestepper.rs根据模型是否使用 guidance 自动决定有 guidance 的模型如 FLUX.1-dev默认 50 步并启用FluxStepperShiftbase_shift: 0.5、max_shift: 1.15、guidance_scale: 4.0无 guidance 的模型如 FLUX.1-schnell默认仅4 步不启用 guidance。这也是 FLUX.1-schnell 主打快速推理的原因——以更少的采样步数换取更快的生成速度。FluxStepper::new后续会调用flux::sampling::get_schedule根据num_steps生成时间步调度。Python 等价实现mistral.rs 同时提供 Python 绑定mistralrs-pyo3对应示例为 examples/python/flux.py接口语义与 Rust 版一一对应from mistralrs import ( Runner, Which, DiffusionArchitecture, ImageGenerationResponseFormat, ) runner Runner( whichWhich.DiffusionPlain( model_idblack-forest-labs/FLUX.1-schnell, archDiffusionArchitecture.FluxOffloaded, ), ) res runner.generate_image( A vibrant sunset in the mountains, 4k, high quality., ImageGenerationResponseFormat.Url, ) print(res.data[0].url)其中DiffusionArchitecture.FluxOffloaded与 Rust 侧的DiffusionLoaderType::FluxOffloaded对应generate_image的默认分辨率同样是 720x1280。注意事项显存FLUX 系列模型体积较大显存有限时优先使用FluxOffloaded加载方式将 Transformer 部分 offload 至 CPUtoken 认证下载 gated 模型需要有效的 Hugging Face token默认从~/.cache/huggingface/token读取可通过with_token_source更换来源分辨率DiffusionGenerationParams::default()为 720x1280修改分辨率会影响生成耗时与显存占用保存文件generate_image的save_file参数可指定图像落盘路径方便在批处理场景下管理输出。相关参考文件示例源码mistralrs/examples/models/diffusion/main.rsBuilder 实现mistralrs/src/diffusion_model.rs加载器与DiffusionLoaderTypemistralrs-core/src/pipeline/loaders/diffusion_loaders.rs生成参数定义mistralrs-core/src/diffusion_models/mod.rs采样核心mistralrs-core/src/diffusion_models/flux/stepper.rsPython 等价示例examples/python/flux.py【免费下载链接】mistral.rsFast, flexible LLM inference项目地址: https://gitcode.com/GitHub_Trending/mi/mistral.rs创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考