1use crate::config::Config;
11use crate::types::*;
12use anyhow::Result;
13use image::{ImageBuffer, Rgb, RgbImage};
14use sha2::{Digest, Sha256};
15use std::collections::BTreeMap;
16use std::io::Cursor;
17use std::time::Instant;
18use tracing::{debug, info, warn};
19
20const TRACE_TARGET_SYNTHETIC: &str = "studio_worker::engine::synthetic";
23
24const TRACE_TARGET_BUILD: &str = "studio_worker::engine";
27
28fn log_engine_roster(engines: &[Box<dyn Engine>]) {
36 let names: Vec<&str> = engines.iter().map(|e| e.name()).collect();
37 info!(
38 target: TRACE_TARGET_BUILD,
39 op = "build",
40 engine_count = names.len(),
41 engines = %names.join(","),
42 "engine roster assembled"
43 );
44}
45
46#[derive(Debug, thiserror::Error)]
55#[error("{engine} engine cannot serve {kind} tasks")]
56pub struct UnsupportedTask {
57 pub engine: &'static str,
58 pub kind: &'static str,
59}
60
61impl UnsupportedTask {
62 pub fn new(engine: &'static str, kind: TaskKind) -> Self {
63 Self {
64 engine,
65 kind: kind.as_str(),
66 }
67 }
68}
69
70#[derive(Debug, Clone, Default)]
72pub struct EngineCapabilities {
73 pub supported_models_per_kind: BTreeMap<TaskKind, Vec<String>>,
76}
77
78impl EngineCapabilities {
79 pub fn supports(&self, kind: TaskKind, model: &str) -> bool {
80 self.supported_models_per_kind
81 .get(&kind)
82 .map(|ms| ms.iter().any(|m| m == model))
83 .unwrap_or(false)
84 }
85
86 pub fn kinds(&self) -> Vec<TaskKind> {
87 self.supported_models_per_kind.keys().copied().collect()
88 }
89
90 pub fn flat_models(&self) -> Vec<String> {
91 self.supported_models_per_kind
92 .values()
93 .flat_map(|ms| ms.iter().cloned())
94 .collect()
95 }
96}
97
98#[cfg(feature = "image-candle")]
99pub mod candle_image;
100pub mod download;
101#[cfg(all(feature = "llama", not(target_os = "windows")))]
104pub mod llama;
105pub mod multi;
106#[cfg(feature = "image-onnx")]
107pub mod onnx;
108#[cfg(feature = "image-onnx")]
109pub mod onnx_provision;
110pub mod sd_provision;
111pub mod sdcpp;
112#[cfg(feature = "tts")]
113pub mod tts;
114#[cfg(feature = "video")]
115pub mod video;
116#[cfg(feature = "whisper")]
117pub mod whisper;
118
119pub trait Engine: Send + Sync {
120 fn name(&self) -> &'static str;
121 fn capabilities(&self) -> EngineCapabilities;
122 fn dispatch(&self, model: &str, task: Task) -> Result<TaskResult>;
123
124 fn dispatch_with_source(
129 &self,
130 model: &str,
131 task: Task,
132 _source: &crate::types::ModelSource,
133 ) -> Result<TaskResult> {
134 self.dispatch(model, task)
135 }
136}
137
138pub fn build(cfg: &Config) -> Result<Box<dyn Engine>> {
149 #[allow(clippy::vec_init_then_push)]
154 let engines: Vec<Box<dyn Engine>> = {
155 let mut v: Vec<Box<dyn Engine>> = Vec::new();
156 #[cfg(all(feature = "llama", not(target_os = "windows")))]
157 v.push(Box::new(llama::LlamaEngine::new(cfg.models_root.clone())?));
158 #[cfg(feature = "whisper")]
159 v.push(Box::new(whisper::WhisperEngine::new(
160 cfg.models_root.clone(),
161 )));
162 #[cfg(feature = "image-candle")]
163 v.push(Box::new(candle_image::CandleImageEngine::new()));
164 #[cfg(feature = "image-onnx")]
169 v.push(Box::new(onnx::OnnxImageEngine::new(
170 cfg.models_root.clone(),
171 )));
172 #[cfg(feature = "video")]
173 v.push(Box::new(video::VideoEngine::new()));
174 #[cfg(feature = "tts")]
175 v.push(Box::new(tts::TtsEngine::new()));
176 v.push(Box::new(sdcpp::SdCppEngine::new(&cfg.models_root)));
182 v.push(Box::new(SyntheticEngine::new()));
183 v
184 };
185
186 log_engine_roster(&engines);
187 Ok(Box::new(multi::MultiEngine::new(engines)))
188}
189
190pub fn default_models_root() -> std::path::PathBuf {
193 crate::config::default_models_root()
194}
195
196pub struct SyntheticEngine;
202
203impl SyntheticEngine {
204 pub fn new() -> Self {
205 Self
206 }
207}
208
209impl Default for SyntheticEngine {
210 fn default() -> Self {
211 Self::new()
212 }
213}
214
215pub const MODEL_WILDCARD: &str = "*";
222
223const DEFAULT_IMAGE_MODELS: &[&str] = &["synthetic", "synthetic-image"];
227const DEFAULT_LLM_MODELS: &[&str] = &["synthetic", "synthetic-llm"];
228const DEFAULT_STT_MODELS: &[&str] = &["synthetic", "synthetic-stt"];
229const DEFAULT_TTS_MODELS: &[&str] = &["synthetic", "synthetic-tts"];
230const DEFAULT_VIDEO_MODELS: &[&str] = &["synthetic", "synthetic-video"];
231
232fn models(list: &[&str]) -> Vec<String> {
233 list.iter().map(|s| (*s).to_string()).collect()
234}
235
236impl Engine for SyntheticEngine {
237 fn name(&self) -> &'static str {
238 "synthetic"
239 }
240
241 fn capabilities(&self) -> EngineCapabilities {
242 let mut map: BTreeMap<TaskKind, Vec<String>> = BTreeMap::new();
243 map.insert(TaskKind::Image, models(DEFAULT_IMAGE_MODELS));
244 map.insert(TaskKind::Llm, models(DEFAULT_LLM_MODELS));
245 map.insert(TaskKind::AudioStt, models(DEFAULT_STT_MODELS));
246 map.insert(TaskKind::AudioTts, models(DEFAULT_TTS_MODELS));
247 map.insert(TaskKind::Video, models(DEFAULT_VIDEO_MODELS));
248 EngineCapabilities {
249 supported_models_per_kind: map,
250 }
251 }
252
253 fn dispatch(&self, model: &str, task: Task) -> Result<TaskResult> {
254 let kind = task.kind();
255 let started = Instant::now();
256 let result = match task {
257 Task::Image(p) => render_procedural(&p.prompt, &p.ext)
258 .map(|bytes| TaskResult::Image { bytes, ext: p.ext }),
259 Task::Llm(p) => {
260 let prompt = p
261 .messages
262 .iter()
263 .map(|m| format!("{}: {}", m.role, m.content))
264 .collect::<Vec<_>>()
265 .join("\n");
266 Ok(TaskResult::Llm {
267 json: synthetic_llm_response(&prompt),
268 })
269 }
270 Task::AudioStt(p) => Ok(TaskResult::AudioStt {
271 json: synthetic_stt_response(&p.input_url, p.language.as_deref()),
272 }),
273 Task::AudioTts(p) => render_wav(&p.text).map(|bytes| TaskResult::AudioTts {
274 bytes,
275 ext: "wav".into(),
276 }),
277 Task::Video(p) => {
278 render_animated_webp(&p.prompt, p.width, p.height, p.seconds).map(|bytes| {
282 TaskResult::Video {
283 bytes,
284 ext: "webp".into(),
285 }
286 })
287 }
288 };
289 let elapsed_ms = started.elapsed().as_millis() as u64;
290 match &result {
291 Ok(_) => debug!(
292 target: TRACE_TARGET_SYNTHETIC,
293 op = "dispatch",
294 kind = kind.as_str(),
295 model,
296 elapsed_ms,
297 "ok"
298 ),
299 Err(e) => warn!(
300 target: TRACE_TARGET_SYNTHETIC,
301 op = "dispatch",
302 kind = kind.as_str(),
303 model,
304 elapsed_ms,
305 error = %e,
306 "failed"
307 ),
308 }
309 result
310 }
311}
312
313pub fn render_procedural(prompt: &str, ext: &str) -> Result<Vec<u8>> {
319 let digest = sha256_bytes(prompt);
320 let palette = [
321 Rgb([digest[0], digest[1], digest[2]]),
322 Rgb([digest[3], digest[4], digest[5]]),
323 Rgb([digest[6], digest[7], digest[8]]),
324 Rgb([digest[9], digest[10], digest[11]]),
325 ];
326
327 let size: u32 = 512;
328 let mut img: RgbImage = ImageBuffer::new(size, size);
329 for (x, y, pixel) in img.enumerate_pixels_mut() {
330 let cx = size as f32 / 2.0;
331 let cy = size as f32 / 2.0;
332 let dx = (x as f32 - cx).abs();
333 let dy = (y as f32 - cy).abs();
334 let chebyshev = dx.max(dy) / cx;
335 let ring = (chebyshev * 6.0).floor() as usize;
336 let base = palette[ring.min(palette.len() - 1)];
337 let phase = ((x as f32 / 24.0).sin() + (y as f32 / 24.0).cos()) * 12.0;
338 *pixel = Rgb([
339 base.0[0].saturating_add(phase as i8 as u8),
340 base.0[1].saturating_add((phase * 0.7) as i8 as u8),
341 base.0[2].saturating_add((phase * 1.3) as i8 as u8),
342 ]);
343 }
344
345 let mut out = Cursor::new(Vec::<u8>::new());
346 let dyn_img = image::DynamicImage::ImageRgb8(img);
347 match ext {
348 "webp" => dyn_img.write_to(&mut out, image::ImageFormat::WebP)?,
349 _ => dyn_img.write_to(&mut out, image::ImageFormat::Png)?,
350 }
351 Ok(out.into_inner())
352}
353
354pub fn synthetic_llm_response(prompt: &str) -> serde_json::Value {
357 let hash = hex::encode(sha256_bytes(prompt));
358 serde_json::json!({
359 "object": "chat.completion",
360 "model": "synthetic-llm",
361 "choices": [{
362 "index": 0,
363 "message": {
364 "role": "assistant",
365 "content": format!("[synthetic] reply to prompt #{}", &hash[..16]),
366 },
367 "finish_reason": "stop",
368 }],
369 "usage": {
370 "prompt_tokens": prompt.split_whitespace().count(),
371 "completion_tokens": 8,
372 "total_tokens": prompt.split_whitespace().count() + 8,
373 },
374 })
375}
376
377pub fn synthetic_stt_response(input_url: &str, language: Option<&str>) -> serde_json::Value {
379 let hash = hex::encode(sha256_bytes(input_url));
380 serde_json::json!({
381 "text": format!("[synthetic] transcript of {}", &hash[..16]),
382 "language": language.unwrap_or("en"),
383 "duration": 1.0,
384 })
385}
386
387pub fn render_wav(text: &str) -> Result<Vec<u8>> {
390 use hound::{SampleFormat, WavSpec, WavWriter};
391 let digest = sha256_bytes(text);
392 let freq_hz = 220.0 + (digest[0] as f32) * (660.0 / 255.0); let sample_rate: u32 = 22_050;
394 let spec = WavSpec {
395 channels: 1,
396 sample_rate,
397 bits_per_sample: 16,
398 sample_format: SampleFormat::Int,
399 };
400
401 let mut buf = Cursor::new(Vec::<u8>::new());
402 {
403 let mut writer = WavWriter::new(&mut buf, spec)?;
404 let total_samples = sample_rate; for n in 0..total_samples {
406 let t = n as f32 / sample_rate as f32;
407 let amplitude = (t * 2.0 * std::f32::consts::PI * freq_hz).sin();
408 let s = (amplitude * 0.4 * i16::MAX as f32) as i16;
409 writer.write_sample(s)?;
410 }
411 writer.finalize()?;
412 }
413 Ok(buf.into_inner())
414}
415
416pub fn render_animated_webp(prompt: &str, _w: u32, _h: u32, seconds: f32) -> Result<Vec<u8>> {
420 let _ = seconds;
426 render_procedural(prompt, "webp")
427}
428
429fn sha256_bytes(input: &str) -> [u8; 32] {
430 let mut hasher = Sha256::new();
431 hasher.update(input.as_bytes());
432 let digest = hasher.finalize();
433 let mut out = [0u8; 32];
434 out.copy_from_slice(&digest);
435 out
436}
437
438#[cfg(test)]
443mod tests {
444 use super::*;
445 use std::io::Cursor;
446
447 #[test]
448 fn synthetic_image_round_trips_as_webp() {
449 let engine = SyntheticEngine::new();
450 let task = Task::Image(ImageParams {
451 prompt: "hello world".into(),
452 width: 512,
453 height: 512,
454 steps: 20,
455 ext: "webp".into(),
456 ..Default::default()
457 });
458 let result = engine.dispatch("synthetic", task).unwrap();
459 let (bytes, ext) = match result {
460 TaskResult::Image { bytes, ext } => (bytes, ext),
461 other => panic!("expected image, got {:?}", other.kind()),
462 };
463 assert_eq!(ext, "webp");
464 assert!(bytes.len() > 100);
465 let reader = image::ImageReader::new(Cursor::new(&bytes))
466 .with_guessed_format()
467 .unwrap();
468 assert_eq!(reader.format().unwrap(), image::ImageFormat::WebP);
469 }
470
471 #[test]
472 fn synthetic_image_round_trips_as_png() {
473 let engine = SyntheticEngine::new();
480 let task = Task::Image(ImageParams {
481 prompt: "hello world".into(),
482 width: 512,
483 height: 512,
484 steps: 20,
485 ext: "png".into(),
486 ..Default::default()
487 });
488 let result = engine.dispatch("synthetic", task).unwrap();
489 let (bytes, ext) = match result {
490 TaskResult::Image { bytes, ext } => (bytes, ext),
491 other => panic!("expected image, got {:?}", other.kind()),
492 };
493 assert_eq!(ext, "png", "the requested ext must be preserved");
494 assert!(bytes.len() > 100);
495 let reader = image::ImageReader::new(Cursor::new(&bytes))
496 .with_guessed_format()
497 .unwrap();
498 assert_eq!(reader.format().unwrap(), image::ImageFormat::Png);
499 }
500
501 #[test]
502 fn synthetic_llm_returns_chat_completion_shape() {
503 let engine = SyntheticEngine::new();
504 let task = Task::Llm(LlmParams {
505 messages: vec![ChatMessage {
506 role: "user".into(),
507 content: "what is the capital of france?".into(),
508 }],
509 max_tokens: 64,
510 temperature: 0.5,
511 ..Default::default()
512 });
513 let result = engine.dispatch("synthetic", task).unwrap();
514 let json = match result {
515 TaskResult::Llm { json } => json,
516 other => panic!("expected llm, got {:?}", other.kind()),
517 };
518 assert_eq!(json["object"], "chat.completion");
519 assert!(json["choices"][0]["message"]["content"]
520 .as_str()
521 .unwrap()
522 .starts_with("[synthetic]"));
523 }
524
525 #[test]
526 fn synthetic_stt_returns_whisper_shape() {
527 let engine = SyntheticEngine::new();
528 let task = Task::AudioStt(AudioSttParams {
529 input_url: "https://example.com/audio.wav".into(),
530 language: Some("nl".into()),
531 ..Default::default()
532 });
533 let result = engine.dispatch("synthetic", task).unwrap();
534 let json = match result {
535 TaskResult::AudioStt { json } => json,
536 other => panic!("expected stt, got {:?}", other.kind()),
537 };
538 assert_eq!(json["language"], "nl");
539 assert!(json["text"].as_str().unwrap().starts_with("[synthetic]"));
540 }
541
542 #[test]
543 fn synthetic_tts_produces_real_wav() {
544 let engine = SyntheticEngine::new();
545 let task = Task::AudioTts(AudioTtsParams {
546 text: "hello world".into(),
547 voice: "default".into(),
548 ext: "wav".into(),
549 ..Default::default()
550 });
551 let result = engine.dispatch("synthetic", task).unwrap();
552 let (bytes, ext) = match result {
553 TaskResult::AudioTts { bytes, ext } => (bytes, ext),
554 other => panic!("expected tts, got {:?}", other.kind()),
555 };
556 assert_eq!(ext, "wav");
557 let mut reader = hound::WavReader::new(Cursor::new(bytes)).expect("real WAV should decode");
559 let spec = reader.spec();
560 assert_eq!(spec.sample_rate, 22_050);
561 assert_eq!(spec.channels, 1);
562 let samples = reader
563 .samples::<i16>()
564 .collect::<std::result::Result<Vec<_>, _>>()
565 .expect("samples should decode");
566 assert_eq!(samples.len(), 22_050); }
568
569 #[test]
570 fn synthetic_video_emits_decodable_bytes() {
571 let engine = SyntheticEngine::new();
572 let task = Task::Video(VideoParams {
573 prompt: "a tiny dragon".into(),
574 seconds: 1.0,
575 width: 256,
576 height: 256,
577 ext: "mp4".into(), ..Default::default()
579 });
580 let result = engine.dispatch("synthetic", task).unwrap();
581 let (bytes, ext) = match result {
582 TaskResult::Video { bytes, ext } => (bytes, ext),
583 other => panic!("expected video, got {:?}", other.kind()),
584 };
585 assert_eq!(ext, "webp");
586 let reader = image::ImageReader::new(Cursor::new(&bytes))
587 .with_guessed_format()
588 .unwrap();
589 assert_eq!(reader.format().unwrap(), image::ImageFormat::WebP);
590 }
591
592 #[test]
593 fn synthetic_engine_advertises_all_kinds() {
594 let engine = SyntheticEngine::new();
595 let caps = engine.capabilities();
596 for k in TaskKind::ALL {
597 assert!(
598 caps.supported_models_per_kind.contains_key(&k),
599 "{} should be advertised",
600 k.as_str()
601 );
602 }
603 assert!(caps.supports(TaskKind::Image, "synthetic"));
604 assert!(
605 !caps.supports(TaskKind::Image, "*"),
606 "synthetic engine MUST NOT advertise the wildcard \
607 (it would happily fulfil real-model jobs with placeholder \
608 bytes, which is destructive on a live queue)"
609 );
610 }
611
612 #[test]
613 fn build_default_yields_multi_engine_with_synthetic_inside() {
614 let cfg = crate::config::Config::default();
618 let eng = build(&cfg).unwrap();
619 assert_eq!(eng.name(), "multi");
620 let caps = eng.capabilities();
621 for k in TaskKind::ALL {
622 assert!(caps.supported_models_per_kind.contains_key(&k));
623 }
624 assert!(caps.supports(TaskKind::Image, "synthetic"));
625 assert!(caps.supports(TaskKind::Llm, "synthetic"));
626 }
627
628 #[test]
629 fn build_emits_engine_roster_breadcrumb() {
630 let logs = crate::test_support::capture(|| {
637 let cfg = crate::config::Config::default();
638 let _ = build(&cfg).unwrap();
639 });
640 assert!(
641 logs.contains("studio_worker::engine"),
642 "expected engine target, got: {logs}"
643 );
644 assert!(logs.contains("op=\"build\""), "expected op=build: {logs}");
645 assert!(
646 logs.contains("engine roster assembled"),
647 "expected roster message: {logs}"
648 );
649 assert!(
650 logs.contains("synthetic"),
651 "expected synthetic in the roster: {logs}"
652 );
653 }
654
655 #[test]
656 fn log_engine_roster_reports_count_and_comma_joined_names() {
657 let logs = crate::test_support::capture(|| {
661 let engines: Vec<Box<dyn Engine>> = vec![
662 Box::new(SyntheticEngine::new()),
663 Box::new(SyntheticEngine::new()),
664 ];
665 log_engine_roster(&engines);
666 });
667 assert!(
668 logs.contains("engine_count=2"),
669 "expected engine_count=2, got: {logs}"
670 );
671 assert!(
672 logs.contains("engines=synthetic,synthetic"),
673 "expected comma-joined names, got: {logs}"
674 );
675 }
676
677 #[test]
678 fn synthetic_engine_is_deterministic_per_prompt() {
679 let engine = SyntheticEngine::new();
680 let task = || {
681 Task::Image(ImageParams {
682 prompt: "deterministic".into(),
683 width: 512,
684 height: 512,
685 steps: 20,
686 ext: "webp".into(),
687 ..Default::default()
688 })
689 };
690 let a = engine.dispatch("synthetic", task()).unwrap();
691 let b = engine.dispatch("synthetic", task()).unwrap();
692 match (a, b) {
693 (TaskResult::Image { bytes: a, .. }, TaskResult::Image { bytes: b, .. }) => {
694 assert_eq!(a, b);
695 }
696 _ => panic!("expected images"),
697 }
698 }
699}