auto-git:
[add] README.md [add] default.png [add] equirect_hdr_icon_512-Wiederhergestellt.png [add] generate_equirect.py [add] icon.png [add] index.html [add] package-lock.json [add] package.json [add] public/ [add] requirements.txt [add] run.sh [add] src-tauri/ [add] src/ [add] vite.config.js
This commit is contained in:
386
src-tauri/src/main.rs
Normal file
386
src-tauri/src/main.rs
Normal file
@@ -0,0 +1,386 @@
|
||||
#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")]
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use std::fs;
|
||||
use std::io::{BufRead, BufReader, Read};
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::{Command, Stdio};
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, Mutex};
|
||||
use tauri::{async_runtime, Env, State, Window};
|
||||
|
||||
#[derive(Default, Clone)]
|
||||
struct Paths {
|
||||
project_root: PathBuf,
|
||||
active_generation_pid: Arc<Mutex<Option<u32>>>,
|
||||
cancel_generation_requested: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct MapInfo {
|
||||
path: String,
|
||||
filename: String,
|
||||
modified: i64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct GenerateResult {
|
||||
output_path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize, Default)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GenSettings {
|
||||
steps: Option<u32>,
|
||||
width: Option<u32>,
|
||||
height: Option<u32>,
|
||||
guidance: Option<f32>,
|
||||
scheduler: Option<String>,
|
||||
upscale: Option<String>,
|
||||
seam_inpaint: Option<bool>,
|
||||
model_path: Option<String>,
|
||||
base_model: Option<String>,
|
||||
vae_model: Option<String>,
|
||||
}
|
||||
|
||||
fn discover_root() -> Result<PathBuf, String> {
|
||||
let cwd = std::env::current_dir().map_err(|e| format!("Failed to resolve current dir: {e}"))?;
|
||||
for ancestor in cwd.ancestors() {
|
||||
if ancestor.join("generate_equirect.py").exists() {
|
||||
return Ok(ancestor.to_path_buf());
|
||||
}
|
||||
}
|
||||
Err("Could not find project root (generate_equirect.py not found in ancestors)".into())
|
||||
}
|
||||
|
||||
fn output_dir(root: &Path) -> Result<PathBuf, String> {
|
||||
let mut dir = root.to_path_buf();
|
||||
dir.push("output");
|
||||
fs::create_dir_all(&dir).map_err(|e| format!("Failed to create output dir: {e}"))?;
|
||||
Ok(dir)
|
||||
}
|
||||
|
||||
fn script_path(root: &Path) -> Result<PathBuf, String> {
|
||||
let candidate = root.join("generate_equirect.py");
|
||||
if candidate.exists() {
|
||||
return Ok(candidate);
|
||||
}
|
||||
Err("generate_equirect.py not found".to_string())
|
||||
}
|
||||
|
||||
/// Pick the Python binary to run: prefer project-local venv to ensure deps are present.
|
||||
fn python_binary(root: &Path) -> PathBuf {
|
||||
let venv_py = root.join(".venv").join("bin").join("python");
|
||||
if venv_py.exists() {
|
||||
return venv_py;
|
||||
}
|
||||
PathBuf::from("python3")
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn list_maps(state: State<Paths>) -> Result<Vec<MapInfo>, String> {
|
||||
let dir = output_dir(&state.project_root)?;
|
||||
let mut maps = Vec::new();
|
||||
if !dir.exists() {
|
||||
return Ok(maps);
|
||||
}
|
||||
for entry in fs::read_dir(&dir).map_err(|e| format!("Failed to read output dir: {e}"))? {
|
||||
let entry = entry.map_err(|e| format!("Failed to read entry: {e}"))?;
|
||||
let path = entry.path();
|
||||
if path
|
||||
.extension()
|
||||
.and_then(|s| s.to_str())
|
||||
.map(|s| s.eq_ignore_ascii_case("png"))
|
||||
!= Some(true)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
let metadata = entry
|
||||
.metadata()
|
||||
.map_err(|e| format!("Metadata error: {e}"))?;
|
||||
let modified = metadata
|
||||
.modified()
|
||||
.map(|m| {
|
||||
m.duration_since(std::time::UNIX_EPOCH)
|
||||
.map(|d| d.as_secs() as i64)
|
||||
.unwrap_or(0)
|
||||
})
|
||||
.unwrap_or(0);
|
||||
let filename = path
|
||||
.file_name()
|
||||
.and_then(|s| s.to_str())
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
maps.push(MapInfo {
|
||||
path: path.to_string_lossy().to_string(),
|
||||
filename,
|
||||
modified,
|
||||
});
|
||||
}
|
||||
maps.sort_by(|a, b| b.modified.cmp(&a.modified));
|
||||
Ok(maps)
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn delete_map(path: String, state: State<Paths>) -> Result<(), String> {
|
||||
let dir = output_dir(&state.project_root)?
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Failed to resolve output dir: {e}"))?;
|
||||
let candidate = PathBuf::from(path);
|
||||
let resolved = candidate
|
||||
.canonicalize()
|
||||
.map_err(|e| format!("Failed to resolve map path: {e}"))?;
|
||||
|
||||
if !resolved.starts_with(&dir) {
|
||||
return Err("Refusing to delete a file outside the output directory".to_string());
|
||||
}
|
||||
if resolved
|
||||
.extension()
|
||||
.and_then(|s| s.to_str())
|
||||
.map(|s| s.eq_ignore_ascii_case("png"))
|
||||
!= Some(true)
|
||||
{
|
||||
return Err("Only PNG maps can be deleted".to_string());
|
||||
}
|
||||
|
||||
fs::remove_file(&resolved).map_err(|e| format!("Failed to delete map: {e}"))
|
||||
}
|
||||
|
||||
fn terminate_process(pid: u32) -> Result<(), String> {
|
||||
#[cfg(target_family = "unix")]
|
||||
{
|
||||
Command::new("kill")
|
||||
.arg("-TERM")
|
||||
.arg(pid.to_string())
|
||||
.status()
|
||||
.map_err(|e| format!("Failed to request cancellation: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(target_family = "windows")]
|
||||
{
|
||||
Command::new("taskkill")
|
||||
.args(["/PID", &pid.to_string(), "/T", "/F"])
|
||||
.status()
|
||||
.map_err(|e| format!("Failed to request cancellation: {e}"))?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
fn cancel_generation(state: State<Paths>) -> Result<(), String> {
|
||||
state
|
||||
.cancel_generation_requested
|
||||
.store(true, Ordering::SeqCst);
|
||||
let pid = *state
|
||||
.active_generation_pid
|
||||
.lock()
|
||||
.map_err(|_| "Failed to read active generation state".to_string())?;
|
||||
|
||||
match pid {
|
||||
Some(0) => Ok(()),
|
||||
Some(pid) => terminate_process(pid),
|
||||
None => Err("No generation is running".to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
#[tauri::command]
|
||||
async fn generate_map(
|
||||
window: Window,
|
||||
prompt: String,
|
||||
settings: Option<GenSettings>,
|
||||
state: State<'_, Paths>,
|
||||
) -> Result<GenerateResult, String> {
|
||||
if prompt.trim().is_empty() {
|
||||
return Err("Prompt must not be empty".into());
|
||||
}
|
||||
|
||||
let prompt_clone = prompt.clone();
|
||||
let root = state.project_root.clone();
|
||||
let cfg = settings.unwrap_or_default();
|
||||
let steps = cfg.steps.unwrap_or(25);
|
||||
let width = cfg.width.unwrap_or(1536);
|
||||
let height = cfg.height.unwrap_or(768);
|
||||
let guidance = cfg.guidance.unwrap_or(6.5);
|
||||
let scheduler = cfg.scheduler.unwrap_or_else(|| "dpmsolver-sde".to_string());
|
||||
let upscale = cfg.upscale.unwrap_or_else(|| "none".to_string());
|
||||
let seam_inpaint = cfg.seam_inpaint.unwrap_or(false);
|
||||
let model_path = cfg
|
||||
.model_path
|
||||
.unwrap_or_else(|| "proximasan/sdxl-360-diffusion".to_string());
|
||||
let base_model = cfg
|
||||
.base_model
|
||||
.unwrap_or_else(|| "stabilityai/stable-diffusion-xl-base-1.0".to_string());
|
||||
let vae_model = cfg
|
||||
.vae_model
|
||||
.unwrap_or_else(|| "madebyollin/sdxl-vae-fp16-fix".to_string());
|
||||
let active_generation_pid = state.active_generation_pid.clone();
|
||||
let cancel_generation_requested = state.cancel_generation_requested.clone();
|
||||
let window = window.clone();
|
||||
|
||||
{
|
||||
let mut active = active_generation_pid
|
||||
.lock()
|
||||
.map_err(|_| "Failed to update active generation state".to_string())?;
|
||||
if active.is_some() {
|
||||
return Err("Generation is already running".to_string());
|
||||
}
|
||||
*active = Some(0);
|
||||
cancel_generation_requested.store(false, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
async_runtime::spawn_blocking(move || {
|
||||
let result = (|| {
|
||||
let out_dir = output_dir(&root)?;
|
||||
let script = script_path(&root)?;
|
||||
let python = python_binary(&root);
|
||||
|
||||
let mut cmd = Command::new(python);
|
||||
cmd.current_dir(&root)
|
||||
.arg(script.as_os_str())
|
||||
.arg("--prompt")
|
||||
.arg(prompt_clone)
|
||||
.arg("--output-dir")
|
||||
.arg(&out_dir)
|
||||
.arg("--work-dir")
|
||||
.arg(root.to_string_lossy().to_string())
|
||||
.arg("--upscale")
|
||||
.arg(upscale)
|
||||
.arg("--steps")
|
||||
.arg(steps.to_string())
|
||||
.arg("--guidance")
|
||||
.arg(guidance.to_string())
|
||||
.arg("--width")
|
||||
.arg(width.to_string())
|
||||
.arg("--height")
|
||||
.arg(height.to_string())
|
||||
.arg("--scheduler")
|
||||
.arg(scheduler)
|
||||
.arg("--model-path")
|
||||
.arg(model_path)
|
||||
.arg("--base-model")
|
||||
.arg(base_model)
|
||||
.arg("--vae-model")
|
||||
.arg(vae_model)
|
||||
.stdout(Stdio::piped())
|
||||
.stderr(Stdio::piped());
|
||||
|
||||
if seam_inpaint {
|
||||
cmd.arg("--seam-inpaint");
|
||||
}
|
||||
|
||||
let mut child = cmd
|
||||
.spawn()
|
||||
.map_err(|e| format!("Failed to start generator: {e}"))?;
|
||||
|
||||
{
|
||||
let mut active = active_generation_pid
|
||||
.lock()
|
||||
.map_err(|_| "Failed to update active generation state".to_string())?;
|
||||
*active = Some(child.id());
|
||||
}
|
||||
|
||||
if cancel_generation_requested.load(Ordering::SeqCst) {
|
||||
let _ = child.kill();
|
||||
}
|
||||
|
||||
let stdout = child
|
||||
.stdout
|
||||
.take()
|
||||
.ok_or_else(|| "Failed to capture generator stdout".to_string())?;
|
||||
let stderr = child
|
||||
.stderr
|
||||
.take()
|
||||
.ok_or_else(|| "Failed to capture generator stderr".to_string())?;
|
||||
|
||||
// Consume stderr in a side thread to avoid blocking.
|
||||
let mut stderr_reader = BufReader::new(stderr);
|
||||
let stderr_handle = std::thread::spawn(move || {
|
||||
let mut buf = String::new();
|
||||
let _ = stderr_reader.read_to_string(&mut buf);
|
||||
buf
|
||||
});
|
||||
|
||||
let mut last_path: Option<String> = None;
|
||||
let mut reader = BufReader::new(stdout);
|
||||
let mut line = String::new();
|
||||
while reader
|
||||
.read_line(&mut line)
|
||||
.map_err(|e| format!("Failed to read generator output: {e}"))?
|
||||
> 0
|
||||
{
|
||||
let trimmed = line.trim_end().to_string();
|
||||
if let Some(json_str) = trimmed.strip_prefix("PROGRESS ") {
|
||||
if let Ok(val) = serde_json::from_str::<Value>(json_str) {
|
||||
let _ = window.emit("gen-progress", val);
|
||||
}
|
||||
} else if !trimmed.is_empty() {
|
||||
last_path = Some(trimmed.clone());
|
||||
}
|
||||
line.clear();
|
||||
}
|
||||
|
||||
let status = child
|
||||
.wait()
|
||||
.map_err(|e| format!("Failed to wait for generator: {e}"))?;
|
||||
let stderr_output = stderr_handle.join().unwrap_or_default();
|
||||
|
||||
if cancel_generation_requested.load(Ordering::SeqCst) {
|
||||
return Err("Generation cancelled".to_string());
|
||||
}
|
||||
|
||||
if !status.success() {
|
||||
return Err(format!("Generator failed: {stderr_output}"));
|
||||
}
|
||||
|
||||
let path_line =
|
||||
last_path.ok_or_else(|| "Generator did not return a path".to_string())?;
|
||||
let path = PathBuf::from(path_line.trim());
|
||||
let resolved = if path.is_absolute() {
|
||||
path
|
||||
} else {
|
||||
out_dir.join(path)
|
||||
};
|
||||
Ok(GenerateResult {
|
||||
output_path: resolved.to_string_lossy().to_string(),
|
||||
})
|
||||
})();
|
||||
|
||||
if let Ok(mut active) = active_generation_pid.lock() {
|
||||
*active = None;
|
||||
}
|
||||
cancel_generation_requested.store(false, Ordering::SeqCst);
|
||||
|
||||
result
|
||||
})
|
||||
.await
|
||||
.map_err(|e| format!("Generation task failed to join: {e}"))?
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let context = tauri::generate_context!();
|
||||
let env = Env::default();
|
||||
// Try to locate the project root by walking up from cwd; if missing, fall back to Tauri resource dir.
|
||||
let root = discover_root()
|
||||
.or_else(|_| {
|
||||
tauri::api::path::resource_dir(context.package_info(), &env)
|
||||
.ok_or_else(|| "Could not locate project root or resource dir".to_string())
|
||||
})
|
||||
.unwrap_or_else(|_| PathBuf::from("."));
|
||||
|
||||
let state = Paths {
|
||||
project_root: root,
|
||||
..Default::default()
|
||||
};
|
||||
tauri::Builder::default()
|
||||
.manage(state)
|
||||
.invoke_handler(tauri::generate_handler![
|
||||
list_maps,
|
||||
delete_map,
|
||||
cancel_generation,
|
||||
generate_map
|
||||
])
|
||||
.run(context)
|
||||
.expect("error while running tauri application");
|
||||
}
|
||||
Reference in New Issue
Block a user