working through the plan + UI/ UX
This commit is contained in:
itsamejms
@@ -0,0 +1,203 @@
|
||||
use crate::llm::AppState;
|
||||
use serde::Deserialize;
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use tauri::Manager;
|
||||
|
||||
// ponytail: DefaultHasher is fine for a cache filename — not crypto, just a stable key.
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ImageRequest {
|
||||
pub prompt: String,
|
||||
/// Override the configured image model for this call.
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
/// One line of Ollama's NDJSON image-generation response.
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct OllamaImageLine {
|
||||
#[serde(default)]
|
||||
done: bool,
|
||||
#[serde(default)]
|
||||
image: Option<String>,
|
||||
}
|
||||
|
||||
/// Generate (or fetch from disk cache) an image for `prompt` via the configured
|
||||
/// Ollama image model. Returns a `data:image/png;base64,...` URL ready for `<img src>`.
|
||||
///
|
||||
/// Ollama image models are macOS-only today; on other platforms we return an error
|
||||
/// so the front-end can fall back to a placeholder instead of a confusing timeout.
|
||||
#[tauri::command]
|
||||
pub async fn generate_image(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, AppState>,
|
||||
req: ImageRequest,
|
||||
) -> Result<String, String> {
|
||||
if cfg!(not(target_os = "macos")) {
|
||||
return Err("image generation is macOS-only via Ollama (for now)".into());
|
||||
}
|
||||
|
||||
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
|
||||
let model = req.model.unwrap_or(config.image_model.clone());
|
||||
|
||||
let cache_dir = app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| e.to_string())?
|
||||
.join("dm-toolkit")
|
||||
.join("images");
|
||||
std::fs::create_dir_all(&cache_dir).map_err(|e| e.to_string())?;
|
||||
|
||||
// Stable cache key over model + prompt. Regenerating with a tweaked prompt
|
||||
// produces a new file; identical prompt reuses the cached PNG.
|
||||
let mut hasher = DefaultHasher::new();
|
||||
model.hash(&mut hasher);
|
||||
req.prompt.hash(&mut hasher);
|
||||
let cache_path = cache_dir.join(format!("{:016x}.png", hasher.finish()));
|
||||
|
||||
// Cache hit: return the stored PNG without calling the model.
|
||||
if cache_path.exists() {
|
||||
let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
|
||||
return Ok(data_url(&bytes));
|
||||
}
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/api/generate", config.api_url.trim_end_matches('/'));
|
||||
let body = serde_json::json!({
|
||||
"model": model,
|
||||
"prompt": req.prompt,
|
||||
"stream": false,
|
||||
});
|
||||
|
||||
let res = client
|
||||
.post(&url)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| format!("image request failed: {e}"))?;
|
||||
|
||||
if !res.status().is_success() {
|
||||
let status = res.status();
|
||||
let text = res.text().await.unwrap_or_default();
|
||||
return Err(format!("image error {status}: {text}"));
|
||||
}
|
||||
|
||||
// Ollama returns newline-delimited JSON even with stream:false for image models;
|
||||
// the final line with `done: true` carries the singular `image` base64 field.
|
||||
let body_text = res.text().await.map_err(|e| format!("image read error: {e}"))?;
|
||||
let mut png_b64: Option<String> = None;
|
||||
for line in body_text.lines() {
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Ok(parsed) = serde_json::from_str::<OllamaImageLine>(line) {
|
||||
if let Some(b64) = parsed.image {
|
||||
png_b64 = Some(b64);
|
||||
if parsed.done {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let b64 = png_b64.ok_or_else(|| "no image data in Ollama response".to_string())?;
|
||||
|
||||
// Decode + persist to cache, then return a data URL.
|
||||
let png_bytes = base64_decode(&b64)?;
|
||||
std::fs::write(&cache_path, &png_bytes).map_err(|e| e.to_string())?;
|
||||
Ok(data_url(&png_bytes))
|
||||
}
|
||||
|
||||
fn data_url(png: &[u8]) -> String {
|
||||
// ponytail: no base64 crate dep — a 30-line encoder.
|
||||
let table: [u8; 64] = *b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
let mut out = String::with_capacity((png.len() + 2) / 3 * 4);
|
||||
let mut chunks = png.chunks_exact(3);
|
||||
for c in &mut chunks {
|
||||
let n = (c[0] as usize) << 16 | (c[1] as usize) << 8 | c[2] as usize;
|
||||
out.push(table[(n >> 18) & 63] as char);
|
||||
out.push(table[(n >> 12) & 63] as char);
|
||||
out.push(table[(n >> 6) & 63] as char);
|
||||
out.push(table[n & 63] as char);
|
||||
}
|
||||
let rem = chunks.remainder();
|
||||
match rem.len() {
|
||||
1 => {
|
||||
let n = (rem[0] as usize) << 16;
|
||||
out.push(table[(n >> 18) & 63] as char);
|
||||
out.push(table[(n >> 12) & 63] as char);
|
||||
out.push('=');
|
||||
out.push('=');
|
||||
}
|
||||
2 => {
|
||||
let n = (rem[0] as usize) << 16 | (rem[1] as usize) << 8;
|
||||
out.push(table[(n >> 18) & 63] as char);
|
||||
out.push(table[(n >> 12) & 63] as char);
|
||||
out.push(table[(n >> 6) & 63] as char);
|
||||
out.push('=');
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
format!("data:image/png;base64,{out}")
|
||||
}
|
||||
|
||||
/// Minimal standard base64 decoder (no extra dependency).
|
||||
fn base64_decode(input: &str) -> Result<Vec<u8>, String> {
|
||||
fn val(c: u8) -> Option<u8> {
|
||||
match c {
|
||||
b'A'..=b'Z' => Some(c - b'A'),
|
||||
b'a'..=b'z' => Some(c - b'a' + 26),
|
||||
b'0'..=b'9' => Some(c - b'0' + 52),
|
||||
b'+' => Some(62),
|
||||
b'/' => Some(63),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
let input = input.trim();
|
||||
let bytes: Vec<u8> = input.bytes().filter(|b| !b.is_ascii_whitespace()).collect();
|
||||
if bytes.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut out = Vec::with_capacity(bytes.len() * 3 / 4);
|
||||
let mut buf: u32 = 0;
|
||||
let mut bits: u32 = 0;
|
||||
for &b in &bytes {
|
||||
if b == b'=' {
|
||||
break;
|
||||
}
|
||||
let v = val(b).ok_or_else(|| format!("invalid base64 char: {b}"))? as u32;
|
||||
buf = (buf << 6) | v;
|
||||
bits += 6;
|
||||
if bits >= 8 {
|
||||
bits -= 8;
|
||||
out.push((buf >> bits) as u8);
|
||||
buf &= (1 << bits) - 1;
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn roundtrip_base64() {
|
||||
let data = b"hello world \x00\xff\x10";
|
||||
let url = data_url(data);
|
||||
let b64 = url.strip_prefix("data:image/png;base64,").unwrap();
|
||||
let decoded = base64_decode(b64).unwrap();
|
||||
assert_eq!(decoded, data);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cache_key_is_stable() {
|
||||
// sanity: same inputs → same filename shape (16 hex digits)
|
||||
let mut h = DefaultHasher::new();
|
||||
"x/flux2-klein:4b".hash(&mut h);
|
||||
"prompt".hash(&mut h);
|
||||
let s = format!("{:016x}", h.finish());
|
||||
assert_eq!(s.len(), 16);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user