Aurora Rust Engine (aurora-rust-engine)

Pure Rust inference engine for modern image generation and diffusion architectures

---

⚡ Overview

aurora-rust-engine is a standalone, lightweight, and memory-efficient AI inference engine written entirely in pure Rust using Candle, native FlashAttention-2 CUDA kernels, and hardware acceleration.

👉 Looking for full documentation? See the complete User & Developer Guide (USER_GUIDE.md) for SDK examples, REST API payloads, scheduler configurations, and VRAM optimization tips.

It provides a robust, zero-Python alternative for running generative diffusion models (Stable Diffusion XL, Pony XL, and future DiT/Flux architectures) with deterministic execution, in-memory zero-overhead LoRA weight merging, and sub-8GB VRAM footprint.

---

✨ Features

• Pure Rust Native Inference: Zero Python dependencies, zero PyTorch overhead, compiled directly to a native standalone executable.
• Unified HuggingFace AutoModel Facade: Standard AutoModel::from_local and AutoModel::from_pretrained interface supporting both generative diffusion pipelines and autoregressive text models.
• CausalLM Text-to-Text Generation (Pure Rust): Autoregressive text generation with GPU KV-Cache supporting Llama 3 / DeepSeek, Qwen 2.5/3.5, Gemma 2/3, and Mistral in both GGUF (quantized Q4_K_M, Q8_0) and SafeTensors formats.
• Flux.1 & Flux.2 Family Full Support: Native implementation of Multimodal Diffusion Transformers (MMDiT) for Flux.1 [dev/schnell], Flux.2-Klein-4B, Flux.2-Klein-9B, and Flux.2-Dev Scaled with exact 3D/4D Rotary Position Embeddings (RoPE).
• Sub-7.5GB VRAM Flux Sequential Block Streaming: Executes massive MMDiT models (3.88B to 12B parameters) with on-demand per-block GPU streaming and zero WDDM paging.
• FLUX.2 Image-to-Image (Img2Img) & Inpainting: Contextual ODE transformation and sharp mask boundary preservation with flow-matching background re-injection.
• Pure Rust 32-Channel & 16-Channel FluxVaeEncoder & FluxVaeDecoder: Bit-exact VAE encoding and decoding with BatchNorm latent standardization.
• SDXL & Pony XL Full Support: Seamless support for all .safetensors single-file checkpoints from Civitai and Hugging Face.
• Native FlashAttention-2 Acceleration: Up to 9.5x faster attention computation with fused CUDA kernels under Windows MSVC and Linux.
• Zero-Overhead In-Memory LoRA Merging: Instant hot-patching of UNet and CLIP weights directly in GPU/CPU memory with 0 MB extra VRAM overhead.
• Exact Penultimate Text Parity: Custom penultimate hidden state extractors for CLIP-L, OpenCLIP-bigG, and multi-layer concat for Qwen3-4B (Layers 9/18/27).
• Seamless $C^\infty$ Cosine Tiled VAE: 4-quadrant $72\times 72$ latent decoding with 128px smooth cosine cross-fade eliminating all tile seams.
• Deterministic Schedulers: Continuous Euler Discrete, Flow Matching Rectified Euler ODE (step_at support for arbitrary start step), and DPM-Solver++ 2M Karras.

---

🚀 Quick Start

1. Prerequisites

• Rust 1.80+ (cargo)
• NVIDIA GPU with CUDA Toolkit 12.x installed
• MSVC Build Tools (Windows) or GCC/Clang (Linux)

2. Build with FlashAttention-2 Acceleration

cargo build --release --features cuda,flash-attn

3. Launch the Pure Rust Interactive Studio (Grio Web UI)

cargo run --release --bin aurora_studio --features cuda,flash-attn,ui

Open http://127.0.0.1:7860 to access the complete pure Rust Diffusion Studio with a multi-model dropdown, per-model generation defaults, real-time progressive latent preview streaming, session history gallery, and GPU telemetry powered by Grio.

Models are declared in aurora_studio.json (no hard-coded paths) — see the end-user guide: 👉 docs/AURORA_STUDIO_GUIDE.md.

4. Run SOTA Grand Benchmark (All Optimizations Active)

cargo run --release --bin grand_benchmark --features cuda,flash-attn

5. Generate an Image via CLI

cargo run --release --bin test_single_gen --features cuda,flash-attn

6. Test LoRA Hot Weight Merging

cargo run --release --bin test_lora --features cuda,flash-attn

7. Run Comprehensive 15-Model Benchmark

cargo run --release --bin stress_test --features cuda,flash-attn

---

🧬 LoRA Integration Example

use candle_core::Device;
use aurora_rust_engine::{StableDiffusionXLPipeline, DiffusionParams};

fn main() -> anyhow::Result<()> {
    let device = Device::new_cuda(0)?;
    let mut pipeline = StableDiffusionXLPipeline::from_safetensors("checkpoint.safetensors", &device)?;

    // Hot-merge LoRA directly into model weights (< 10 seconds, 0 MB extra VRAM)
    pipeline.load_lora("style_lora.safetensors", 0.85)?;

    let params = DiffusionParams {
        prompt: "masterpiece, 1girl, cyberpunk city, vivid colors",
        negative_prompt: Some("blurry, low quality"),
        num_steps: 25,
        guidance_scale: 6.0,
        width: 1024,
        height: 1024,
        seed: 42,
    };

    let image = pipeline.generate(params, None)?;
    image.save("output_lora.png")?;

    // Unload LoRA to restore base checkpoint weights
    pipeline.unload_all_loras()?;

    Ok(())
}

---

🖼️ Image-to-Image (Img2Img) Example

use candle_core::Device;
use aurora_rust_engine::{StableDiffusionXLPipeline, Img2ImgParams};

fn main() -> anyhow::Result<()> {
    let device = Device::new_cuda(0)?;
    let mut pipeline = StableDiffusionXLPipeline::from_safetensors("checkpoint.safetensors", &device)?;

    let input_img = image::open("input.png")?.to_rgb8();

    let params = Img2ImgParams {
        prompt: "masterpiece, 1girl, golden radiant armor, fiery glowing orange hair",
        negative_prompt: Some("blurry, low quality"),
        image: input_img,
        strength: 0.60, // 0.0 = identity, 1.0 = full re-generation
        num_steps: 30,
        guidance_scale: 6.5,
        seed: 42,
    };

    let result = pipeline.generate_img2img(params, None)?;
    result.save("output_img2img.png")?;

    Ok(())
}

Inpainting & Mask-Guided Diffusion

use aurora_rust_engine::{InpaintParams, StableDiffusionXLPipeline, select_device};

fn main() -> anyhow::Result<()> {
    let device = select_device()?;
    let mut pipeline = StableDiffusionXLPipeline::from_single_file("sdxl_base.safetensors", device)?;

    let base_image = image::open("input.png")?.to_rgb8();
    let mask_image = image::open("mask.png")?.to_luma8(); // White = edit, Black = keep

    let params = InpaintParams {
        prompt: "a wizard hat with golden stars",
        negative_prompt: Some("low quality, blurry"),
        image: base_image,
        mask: mask_image,
        mask_blur: 8,
        strength: 0.95,
        num_steps: 30,
        guidance_scale: 7.0,
        seed: 42,
    };

    let result = pipeline.generate_inpaint(params, None)?;
    result.save("output_inpaint.png")?;

    Ok(())
}

Multi-ControlNet Spatial Guidance

use aurora_rust_engine::{compute_canny_edge_map, ControlNetModel, ControlNetParams, MultiControlNet, StableDiffusionXLPipeline, select_device};

fn main() -> anyhow::Result<()> {
    let device = select_device()?;
    let mut pipeline = StableDiffusionXLPipeline::from_single_file("sdxl_base.safetensors", device.clone())?;

    // 1. Extract Canny edge map in Pure Rust (< 12ms)
    let source_img = image::open("input.png")?.to_rgb8();
    let edge_map = compute_canny_edge_map(&source_img, 100.0, 200.0);

    // 2. Load ControlNet model and configure MultiControlNet container
    let cnet = ControlNetModel::from_safetensors("controlnet_canny_sdxl.safetensors", &device, candle_core::DType::F16)?;
    let mut multi_controlnet = MultiControlNet::new();
    multi_controlnet.add(cnet, 0.85); // 0.85 conditioning strength

    // 3. Generate with spatial edge alignment
    let params = ControlNetParams::new("cyberpunk warrior, masterpiece, highly detailed", edge_map);
    let result = pipeline.generate_controlnet(params, &multi_controlnet, None)?;
    result.save("output_controlnet.png")?;

    Ok(())
}

High-Resolution Disentangled Profiling

let (image, metrics) = pipeline.generate_with_metrics(params, None)?;
println!("{}", metrics.summary_report());
// Output: ⏱️ [Telemetry] UNet: 15.37s (30 steps, 512.42 ms/step, 1.95 it/s) | VAE: 4.73s | Text: 2.33s | Total: 22.57s

Production REST & WebSocket Inference Server

Start the standalone async inference microservice:

cargo run --release --bin server --features cuda,flash-attn

• Health Check: GET http://127.0.0.1:8080/api/v1/health
• Text-to-Image Generation: POST http://127.0.0.1:8080/api/v1/generate

  {
    "prompt": "futuristic cyberpunk pilot, 8k masterpiece",
    "steps": 30,
    "guidance_scale": 6.5,
    "width": 1024,
    "height": 1024
  }
  

• WebSocket Streaming: ws://127.0.0.1:8080/api/v1/ws

---

📊 Benchmark Summary (RTX 4070 Ti 12GB)

---

🗺️ Project Roadmap

See ROADMAP.md for full technical specifications and development milestones:

• Milestone 1: SDXL Core Pipeline & Conditioning Parity
• Milestone 2: FlashAttention-2 Windows MSVC Kernel Fusion
• Milestone 3: LoRA & LyCORIS Engine & In-Memory Hot Weight Merging
• Milestone 4: Image-to-Image (Img2Img) Pipeline
• Milestone 5: Inpainting & Outpainting Pipeline
• Milestone 6: Multi-ControlNet (OpenPose, Depth, Canny) & IP-Adapter Conditioners
• Milestone 7: Telemetry Profiler, Parameterized Kernel Dispatch & Adaptive VAE
• Milestone 8: Production Async Axum Server & WebSocket Streaming
• Milestone 9: Flux.1 MMDiT Diffusion Transformers (Dev/Schnell)
• Milestone 10: Flux.2-Klein-4B MMDiT & Quality Parity
• Milestone 11: Scaling to FLUX.2-Klein-9B & FLUX.2-Dev
• Milestone 12: FLUX.2 Img2Img & Inpainting Pipeline (Pure Rust 16/32-ch VAE)
• Milestone 16: CausalLM LLM Text Generation & AutoModel Facade (Llama, DeepSeek, Qwen, Gemma, Mistral in GGUF/Safetensors)

---

📄 License

Licensed under Apache-2.0 / MIT.