studio_worker/engine/
video.rs1use crate::engine::render_procedural;
17use crate::engine::{Engine, EngineCapabilities};
18use crate::types::*;
19use anyhow::{Context, Result};
20use gif::{Encoder, Frame, Repeat};
21use std::collections::BTreeMap;
22use std::io::Cursor;
23use std::time::Instant;
24use tracing::{debug, warn};
25
26const TRACE_TARGET: &str = "studio_worker::engine::video";
30
31pub struct VideoEngine;
32
33impl VideoEngine {
34 pub fn new() -> Self {
35 Self
36 }
37}
38
39impl Default for VideoEngine {
40 fn default() -> Self {
41 Self::new()
42 }
43}
44
45const MODEL_ID: &str = "procedural-gif";
46
47impl Engine for VideoEngine {
48 fn name(&self) -> &'static str {
49 "video"
50 }
51
52 fn capabilities(&self) -> EngineCapabilities {
53 let mut map: BTreeMap<TaskKind, Vec<String>> = BTreeMap::new();
54 map.insert(TaskKind::Video, vec![MODEL_ID.to_string()]);
55 EngineCapabilities {
56 supported_models_per_kind: map,
57 }
58 }
59
60 fn dispatch(&self, model: &str, task: Task) -> Result<TaskResult> {
61 let kind = task.kind();
62 let started = Instant::now();
63 let params = match task {
64 Task::Video(p) => p,
65 other => {
66 warn!(
67 target: TRACE_TARGET,
68 op = "dispatch",
69 kind = kind.as_str(),
70 model,
71 "unsupported task kind"
72 );
73 return Err(crate::engine::UnsupportedTask::new("video", other.kind()).into());
74 }
75 };
76 let result = render_gif(¶ms);
77 let elapsed_ms = started.elapsed().as_millis() as u64;
78 match &result {
79 Ok(bytes) => debug!(
80 target: TRACE_TARGET,
81 op = "dispatch",
82 kind = kind.as_str(),
83 model,
84 seconds = params.seconds,
85 width = params.width,
86 height = params.height,
87 bytes = bytes.len(),
88 elapsed_ms,
89 "ok"
90 ),
91 Err(e) => warn!(
92 target: TRACE_TARGET,
93 op = "dispatch",
94 kind = kind.as_str(),
95 model,
96 elapsed_ms,
97 error = %e,
98 "failed"
99 ),
100 }
101 let bytes = result?;
102 Ok(TaskResult::Video {
103 bytes,
104 ext: "gif".into(),
105 })
106 }
107}
108
109pub fn render_gif(params: &VideoParams) -> Result<Vec<u8>> {
113 let width = params.width.clamp(64, 1024) as u16;
114 let height = params.height.clamp(64, 1024) as u16;
115 let fps: u32 = 10;
116 let n_frames = (params.seconds.max(0.1) * fps as f32).round().max(1.0) as u32;
117 let mut out = Cursor::new(Vec::<u8>::new());
118 {
119 let mut encoder =
120 Encoder::new(&mut out, width, height, &[]).context("creating GIF encoder")?;
121 encoder.set_repeat(Repeat::Infinite)?;
122 for i in 0..n_frames {
123 let frame_prompt = format!("{} #{i}", params.prompt);
124 let png_bytes = render_procedural(&frame_prompt, "png")?;
125 let img = image::load_from_memory(&png_bytes)?
126 .resize_exact(
127 u32::from(width),
128 u32::from(height),
129 image::imageops::FilterType::Triangle,
130 )
131 .to_rgba8();
132 let mut buf = img.into_raw();
133 let mut frame = Frame::from_rgba_speed(width, height, &mut buf, 10);
134 frame.delay = 10;
136 encoder.write_frame(&frame).context("writing GIF frame")?;
137 }
138 }
139 Ok(out.into_inner())
140}
141
142#[cfg(test)]
143mod tests {
144 use super::*;
145
146 fn task(seconds: f32) -> Task {
147 Task::Video(VideoParams {
148 prompt: "a tiny dragon".into(),
149 seconds,
150 width: 128,
151 height: 128,
152 ext: "gif".into(),
153 ..Default::default()
154 })
155 }
156
157 #[test]
158 fn engine_advertises_video_kind() {
159 let engine = VideoEngine::new();
160 let caps = engine.capabilities();
161 assert_eq!(
162 caps.supported_models_per_kind[&TaskKind::Video],
163 vec![MODEL_ID.to_string()]
164 );
165 assert_eq!(engine.name(), "video");
166 }
167
168 #[test]
169 fn engine_default_constructs() {
170 let _ = VideoEngine;
171 }
172
173 #[test]
174 fn dispatch_rejects_non_video_tasks() {
175 let engine = VideoEngine::new();
176 let err = engine
177 .dispatch(
178 MODEL_ID,
179 Task::Llm(LlmParams {
180 messages: vec![],
181 max_tokens: 1,
182 temperature: 0.0,
183 ..Default::default()
184 }),
185 )
186 .unwrap_err();
187 assert!(err.to_string().contains("cannot serve llm"));
188 }
189
190 #[test]
191 fn dispatch_produces_valid_animated_gif() {
192 let engine = VideoEngine::new();
193 let result = engine.dispatch(MODEL_ID, task(1.0)).unwrap();
194 let (bytes, ext) = match result {
195 TaskResult::Video { bytes, ext } => (bytes, ext),
196 other => panic!("expected video, got {:?}", other.kind()),
197 };
198 assert_eq!(ext, "gif");
199 assert!(&bytes[..3] == b"GIF");
201 let mut decoder_opts = gif::DecodeOptions::new();
203 decoder_opts.set_color_output(gif::ColorOutput::RGBA);
204 let mut decoder = decoder_opts
205 .read_info(Cursor::new(bytes.clone()))
206 .expect("decode gif");
207 let mut frames = 0;
208 while decoder.next_frame_info().expect("frame info").is_some() {
209 frames += 1;
210 let mut tmp = vec![0u8; decoder.buffer_size()];
212 decoder.read_into_buffer(&mut tmp).unwrap();
213 }
214 assert_eq!(frames, 10, "expected 10 frames, got {frames}");
216 }
217
218 #[test]
219 fn shorter_seconds_produces_fewer_frames() {
220 let small = render_gif(&VideoParams {
221 prompt: "a".into(),
222 seconds: 0.2,
223 width: 64,
224 height: 64,
225 ext: "gif".into(),
226 ..Default::default()
227 })
228 .unwrap();
229 let big = render_gif(&VideoParams {
230 prompt: "a".into(),
231 seconds: 2.0,
232 width: 64,
233 height: 64,
234 ext: "gif".into(),
235 ..Default::default()
236 })
237 .unwrap();
238 assert!(big.len() > small.len());
239 }
240}