diff --git a/src-tauri/src/main.rs b/src-tauri/src/main.rs index 4e379d9..174c002 100644 --- a/src-tauri/src/main.rs +++ b/src-tauri/src/main.rs @@ -95,6 +95,17 @@ fn unix_timestamp_millis() -> u128 { .unwrap_or(0) } +fn env_truthy(name: &str) -> bool { + std::env::var(name) + .map(|value| { + matches!( + value.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) + .unwrap_or(false) +} + fn log_slug(prompt: &str) -> String { let mut slug = String::new(); let mut last_was_separator = false; @@ -377,6 +388,7 @@ async fn generate_map( let vae_model = cfg .vae_model .unwrap_or_else(|| "madebyollin/sdxl-vae-fp16-fix".to_string()); + let allow_downloads = env_truthy("SKYMAP_ALLOW_DOWNLOADS"); let active_generation_pid = state.active_generation_pid.clone(); let cancel_generation_requested = state.cancel_generation_requested.clone(); let window = window.clone(); @@ -451,10 +463,14 @@ async fn generate_map( &mut log_file, format!(" program: {}", python.to_string_lossy()), ); - write_log_line( - &mut log_file, - " env: PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 DIFFUSERS_OFFLINE=1", - ); + if allow_downloads { + write_log_line(&mut log_file, " env: PYTHONUNBUFFERED=1 SKYMAP_ALLOW_DOWNLOADS=1"); + } else { + write_log_line( + &mut log_file, + " env: PYTHONUNBUFFERED=1 HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 DIFFUSERS_OFFLINE=1", + ); + } for arg in &args { write_log_line(&mut log_file, format!(" arg: {arg}")); } @@ -463,11 +479,15 @@ async fn generate_map( let mut cmd = Command::new(&python); cmd.current_dir(&root) .env("PYTHONUNBUFFERED", "1") - .env("HF_HUB_OFFLINE", "1") - .env("TRANSFORMERS_OFFLINE", "1") - .env("DIFFUSERS_OFFLINE", "1") .stdout(Stdio::piped()) .stderr(Stdio::piped()); + if allow_downloads { + cmd.env("SKYMAP_ALLOW_DOWNLOADS", "1"); + } else { + cmd.env("HF_HUB_OFFLINE", "1") + .env("TRANSFORMERS_OFFLINE", "1") + .env("DIFFUSERS_OFFLINE", "1"); + } for arg in &args { cmd.arg(arg); }