136 lines
3.7 KiB
Rust
136 lines
3.7 KiB
Rust
//! # RTX Diffuse
|
|
//!
|
|
//! High-performance diffusion models for RustyTorch++.
|
|
//! Implements state-of-the-art architectures including UNet and DiT (Diffusion Transformers)
|
|
//! with CUDA acceleration and kernel fusion optimization.
|
|
//!
|
|
//! ## Features
|
|
//!
|
|
//! - Multiple noise scheduling strategies (Linear, Cosine, Scaled Linear)
|
|
//! - Advanced samplers (DDIM, DPM++, UniPC, Euler Ancestral, DDPM)
|
|
//! - UniPC sampler with fast 5-10 step sampling and adaptive order selection
|
|
//! - UNet architecture for image generation
|
|
//! - DiT (Diffusion Transformer) for scalable generation
|
|
//! - Kernel fusion optimization via rtx-polygraph
|
|
//! - Memory-efficient implementations
|
|
//!
|
|
//! ## Example
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_diffuse::{
|
|
//! NoiseGenerator, NoiseSchedule, DiffusionScheduler, SchedulerType,
|
|
//! UniPCSampler, UniPCConfig, PredictionType
|
|
//! };
|
|
//!
|
|
//! # fn example() -> rtx_diffuse::Result<()> {
|
|
//! // Create noise generator with cosine schedule
|
|
//! let mut noise_gen = NoiseGenerator::new(
|
|
//! NoiseSchedule::Cosine { s: 0.008 },
|
|
//! 1000,
|
|
//! Some(42)
|
|
//! )?;
|
|
//!
|
|
//! // Create UniPC scheduler for fast sampling
|
|
//! let scheduler = DiffusionScheduler::new(
|
|
//! SchedulerType::UniPC {
|
|
//! predictor_order: 2,
|
|
//! corrector_order: 1,
|
|
//! use_corrector: true
|
|
//! },
|
|
//! noise_gen.clone(),
|
|
//! 10 // Fast 10-step sampling
|
|
//! )?;
|
|
//!
|
|
//! // Or use standalone UniPC sampler for advanced features
|
|
//! let mut unipc = UniPCSampler::new(
|
|
//! UniPCConfig {
|
|
//! predictor_order: 3,
|
|
//! adaptive_order: true,
|
|
//! variance_reduction: true,
|
|
//! prediction_type: PredictionType::Epsilon,
|
|
//! ..Default::default()
|
|
//! },
|
|
//! noise_gen
|
|
//! )?;
|
|
//! # Ok(())
|
|
//! # }
|
|
//! ```
|
|
|
|
pub mod conditioning;
|
|
pub mod controlnet;
|
|
pub mod ddim;
|
|
pub mod dpm_solver_pp;
|
|
pub mod error;
|
|
pub mod guidance;
|
|
pub mod ip_adapter;
|
|
pub mod ldm;
|
|
pub mod models;
|
|
pub mod noise;
|
|
pub mod scheduler;
|
|
pub mod t2i_adapter;
|
|
|
|
#[cfg(test)]
|
|
pub mod dpm_solver_pp_validation;
|
|
|
|
#[cfg(test)]
|
|
pub mod controlnet_tests;
|
|
|
|
#[cfg(test)]
|
|
pub mod controlnet_minimal_test;
|
|
|
|
#[cfg(test)]
|
|
pub mod ldm_tests;
|
|
|
|
pub mod lcm_sampler;
|
|
pub mod lora_diffusion;
|
|
pub mod physics_conditioner;
|
|
pub mod unipc_sampler;
|
|
|
|
#[cfg(test)]
|
|
pub mod lcm_demo;
|
|
|
|
// Re-export commonly used items
|
|
pub use ddim::{DDIMConfig, DDIMSampler};
|
|
pub use dpm_solver_pp::{
|
|
DPMSolverConfig, DPMSolverPP, DPMSolverStats, PredictionType as DPMPredictionType,
|
|
};
|
|
pub use error::{DiffusionError, Result};
|
|
pub use guidance::{CFGConfig, ClassifierFreeGuidance, GuidanceScale};
|
|
pub use ip_adapter::{
|
|
CLIPImageEncoder, DecoupledCrossAttention, IPAdapter, IPAdapterConfig, ProjectionLayer,
|
|
};
|
|
pub use lcm_sampler::{LCMConfig, LCMPredictionType, LCMSampler, LCMStats};
|
|
pub use lora_diffusion::{
|
|
AttentionLoRALayer, LayerTargeting, LoRAAdapter, LoRAConfig, LoRAManager, LoRAUNet,
|
|
};
|
|
pub use models::{DiT, DiTConfig, UNet, UNetConfig};
|
|
pub use noise::{NoiseGenerator, NoiseSchedule};
|
|
pub use physics_conditioner::{
|
|
BoundaryMode, PDEType, PhysicsConditioner, PhysicsConfig, PhysicsResidual, PhysicsSchedule,
|
|
ResidualNormalization,
|
|
};
|
|
pub use scheduler::{DiffusionScheduler, SchedulerType};
|
|
pub use t2i_adapter::{
|
|
AdapterBlock, ConditionEncoder, ConditionType, FeatureAligner, T2IAdapter, T2IAdapterConfig,
|
|
};
|
|
pub use unipc_sampler::{PredictionType, UniPCConfig, UniPCSampler, UniPCStats};
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_crate_integration() {
|
|
// Basic integration test
|
|
let noise_gen = NoiseGenerator::new(
|
|
NoiseSchedule::Linear {
|
|
beta_start: 0.0001,
|
|
beta_end: 0.02,
|
|
},
|
|
1000,
|
|
Some(42),
|
|
);
|
|
assert!(noise_gen.is_ok());
|
|
}
|
|
}
|