darkmedia-x-api / engine /src /optimization.rs
cybermedia's picture
Upload folder using huggingface_hub
343eed9 verified
//! Module d'optimisation pour ressources limitées
//! Stratégie: Single-threaded, streaming, minimal memory footprint
use std::path::{Path, PathBuf};
use crate::error::Result;
/// Configuration pour optimisation mémoire
#[derive(Debug, Clone, Copy)]
pub struct MemoryOptimization {
/// Vider le cache après chaque image (vs garder en mémoire)
pub aggressive_cleanup: bool,
/// Taille max du cache d'images (0 = aucun cache)
pub max_cache_mb: usize,
/// Réduire la qualité si pas assez de RAM
pub adaptive_quality: bool,
/// Utiliser les fichiers temporaires plutôt que la RAM
pub use_temp_files: bool,
}
impl Default for MemoryOptimization {
fn default() -> Self {
Self {
aggressive_cleanup: true,
max_cache_mb: 100,
adaptive_quality: true,
use_temp_files: true,
}
}
}
/// Détecteur de ressources disponibles
pub struct ResourceMonitor {
pub optimization: MemoryOptimization,
}
impl ResourceMonitor {
pub fn new() -> Self {
Self {
optimization: MemoryOptimization::default(),
}
}
/// Vérifier la RAM disponible (estimation)
pub fn get_available_memory_mb() -> usize {
// Estimation simple: on suppose ~500MB dispo pour l'app
// En production, utiliser sysinfo crate
500
}
/// Ajuster les paramètres selon les ressources
pub fn auto_configure() -> MemoryOptimization {
let available = Self::get_available_memory_mb();
MemoryOptimization {
aggressive_cleanup: available < 1000,
max_cache_mb: (available / 5).min(200),
adaptive_quality: available < 2000,
use_temp_files: available < 1500,
}
}
/// Nettoyer les ressources
pub fn cleanup(temp_dir: &Path) -> std::io::Result<()> {
if temp_dir.exists() {
std::fs::remove_dir_all(temp_dir)?;
}
Ok(())
}
}
/// Cache avec limite de taille
pub struct LimitedCache {
max_size_mb: usize,
current_size_mb: usize,
}
impl LimitedCache {
pub fn new(max_size_mb: usize) -> Self {
Self {
max_size_mb,
current_size_mb: 0,
}
}
pub fn can_cache(&self, size_mb: usize) -> bool {
self.current_size_mb + size_mb <= self.max_size_mb
}
pub fn add(&mut self, size_mb: usize) {
self.current_size_mb += size_mb;
}
pub fn clear(&mut self) {
self.current_size_mb = 0;
}
}
/// Gestionnaire de génération en streaming pour traiter images une à une
pub struct StreamingGenerator {
pub temp_dir: PathBuf,
pub use_temp: bool,
}
impl StreamingGenerator {
pub fn new(work_dir: &Path, use_temp: bool) -> Result<Self> {
let temp_dir = work_dir.join(".temp");
if use_temp && !temp_dir.exists() {
std::fs::create_dir_all(&temp_dir)?;
}
Ok(Self { temp_dir, use_temp })
}
/// Traiter une image en streaming (charge -> traite -> décharge)
pub async fn process_image_streaming<F>(
&self,
input: &Path,
output: &Path,
processor: F,
) -> Result<()>
where
F: Fn(image::DynamicImage) -> Result<image::DynamicImage>,
{
// 1. Charger
let img = image::open(input).map_err(|e| {
crate::error::GeneratorError::ImageGenerationFailed(format!("Load failed: {}", e))
})?;
// 2. Traiter (reste en mémoire)
let processed = processor(img)?;
// 3. Sauvegarder
processed.save(output).map_err(|e| {
crate::error::GeneratorError::ImageGenerationFailed(format!("Save failed: {}", e))
})?;
// 4. Nettoyer explicitement
drop(processed);
Ok(())
}
/// Cleanup
pub fn cleanup(&self) -> std::io::Result<()> {
if self.use_temp && self.temp_dir.exists() {
std::fs::remove_dir_all(&self.temp_dir)?;
}
Ok(())
}
}
/// Checkpointing pour reprendre depuis le dernier succès
pub struct CheckpointManager {
pub checkpoint_file: PathBuf,
}
impl CheckpointManager {
pub fn new(work_dir: &Path) -> Self {
Self {
checkpoint_file: work_dir.join(".generation_checkpoint.json"),
}
}
/// Sauvegarder le progrès
pub fn save_checkpoint(&self, story_id: &str, scene_id: usize) -> Result<()> {
#[derive(serde::Serialize)]
struct Checkpoint {
story_id: String,
last_scene: usize,
}
let checkpoint = Checkpoint {
story_id: story_id.to_string(),
last_scene: scene_id,
};
let json = serde_json::to_string(&checkpoint)?;
std::fs::write(&self.checkpoint_file, json)?;
Ok(())
}
/// Charger le dernier checkpoint
pub fn load_checkpoint(&self) -> Result<Option<(String, usize)>> {
if !self.checkpoint_file.exists() {
return Ok(None);
}
#[derive(serde::Deserialize)]
struct Checkpoint {
story_id: String,
last_scene: usize,
}
let json = std::fs::read_to_string(&self.checkpoint_file)?;
let checkpoint: Checkpoint = serde_json::from_str(&json)?;
Ok(Some((checkpoint.story_id, checkpoint.last_scene)))
}
/// Nettoyer après succès
pub fn clear(&self) -> std::io::Result<()> {
if self.checkpoint_file.exists() {
std::fs::remove_file(&self.checkpoint_file)?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_limited_cache() {
let mut cache = LimitedCache::new(100);
assert!(cache.can_cache(50));
cache.add(50);
assert!(cache.can_cache(49));
assert!(!cache.can_cache(51));
cache.clear();
assert!(cache.can_cache(100));
}
}