mirror of
https://github.com/zenlm/engine.git
synced 2026-10-30 15:00:05 +00:00
Eric BuehlerandGitHub
efbe4a69fc
* Sketch the diffusion pipeline, loader, model, and FLUX model * Connect response part to openai server * Sketch the flux loader * Handle 256 limitiation of timestep distilliation model * Work on loading * Handle raw prompts for image gen models * Should build now * Clippy * Add diffusion plain model * Typos * Sketch a server example * Add parse_diffusion_arch * Add debug bound * Use as_any * Some fixes * Fix config * Loads now * Avoid empty prompt err * Fix overflow * Fix overflow * Use the correct processor * Add debugs * Cast to correct dtype * Squeeze * Use sharded and downcast repo * Better loading * Dont modify cache * Fix response * Use big t5 for now * Only load relevant tensors * Clamp for f16 * Do some things * Use sharded and downcast repo * Diffusion model case * Oops * Oops * Hotload t5xxl * Use only encoder parts * Don't cast it * Fix files * Format * Cast every layer * Updated resolution and prompt * Fix issue with dual running * Reduce lifetime of hotloaded t5 xxl * Offloading t5 layers to cpu * Add offloading config * Add Python, Rust support, docs * Add configurable height and width
39 lines
916 B
Rust
39 lines
916 B
Rust
use std::time::Instant;
|
|
|
|
use anyhow::Result;
|
|
use mistralrs::{
|
|
DiffusionGenerationParams, DiffusionLoaderType, DiffusionModelBuilder,
|
|
ImageGenerationResponseFormat,
|
|
};
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
let model = DiffusionModelBuilder::new(
|
|
"black-forest-labs/FLUX.1-schnell",
|
|
DiffusionLoaderType::FluxOffloaded,
|
|
)
|
|
.with_logging()
|
|
.build()
|
|
.await?;
|
|
|
|
let start = Instant::now();
|
|
|
|
let response = model
|
|
.generate_image(
|
|
"A vibrant sunset in the mountains, 4k, high quality.".to_string(),
|
|
ImageGenerationResponseFormat::Url,
|
|
DiffusionGenerationParams::default(),
|
|
)
|
|
.await?;
|
|
|
|
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()
|
|
);
|
|
|
|
Ok(())
|
|
}
|