studio_worker/
types.rs

1//! Shared wire-format mirrors of the studio API.
2use serde::{Deserialize, Serialize};
3use std::collections::BTreeMap;
4
5// ---------------------------------------------------------------------------
6// Task kinds — every job claimed by the worker is one of these.
7// ---------------------------------------------------------------------------
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Hash)]
10#[serde(rename_all = "snake_case")]
11pub enum TaskKind {
12    Image,
13    Llm,
14    AudioStt,
15    AudioTts,
16    Video,
17}
18
19impl TaskKind {
20    pub const ALL: [TaskKind; 5] = [
21        TaskKind::Image,
22        TaskKind::Llm,
23        TaskKind::AudioStt,
24        TaskKind::AudioTts,
25        TaskKind::Video,
26    ];
27
28    pub fn as_str(&self) -> &'static str {
29        match self {
30            TaskKind::Image => "image",
31            TaskKind::Llm => "llm",
32            TaskKind::AudioStt => "audio_stt",
33            TaskKind::AudioTts => "audio_tts",
34            TaskKind::Video => "video",
35        }
36    }
37}
38
39impl std::fmt::Display for TaskKind {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.write_str(self.as_str())
42    }
43}
44
45// ---------------------------------------------------------------------------
46// Per-kind task parameters
47// ---------------------------------------------------------------------------
48
49#[derive(Debug, Clone, Default, Serialize, Deserialize)]
50#[serde(rename_all = "camelCase")]
51pub struct ImageParams {
52    pub prompt: String,
53    /// Negative prompt — what the model should steer AWAY from.
54    /// Optional: when `None`, `sd-cli` is invoked without
55    /// `--negative-prompt`.  Wire-format key: `negativePrompt`.
56    #[serde(default)]
57    pub negative_prompt: Option<String>,
58    /// HTTPS URL to a base image for image-to-image generation.
59    /// When set, the worker downloads the bytes to a local tempfile
60    /// and invokes `sd-cli --init-img <path>`.  Required for any
61    /// task profile that declares the `initImage` capability.
62    #[serde(default)]
63    pub init_image_url: Option<String>,
64    /// HTTPS URL to a black/white inpaint mask (white = the region the
65    /// model may repaint). When set alongside `init_image_url`, the
66    /// worker downloads it and invokes `sd-cli --mask <path>` so only
67    /// the masked region changes. Wire-format key: `maskUrl`.
68    #[serde(default)]
69    pub mask_url: Option<String>,
70    /// HTTPS URL to a reference image for instruction-edit models
71    /// (e.g. Qwen-Image-Edit / Flux Kontext).  When set, the worker
72    /// downloads it and invokes `sd-cli -r <path>` (reference mode)
73    /// instead of the `--init-img`/`--strength`/`--mask` img2img path:
74    /// the model regenerates the whole image from the reference per the
75    /// instruction prompt; per-region clipping happens in the studio
76    /// composite. Wire-format key: `refImageUrl`.
77    #[serde(default)]
78    pub ref_image_url: Option<String>,
79    /// Denoise / noise-strength for i2i (0.0 = keep init image
80    /// unchanged, 1.0 = full re-noise).  Maps to `sd-cli --strength`.
81    #[serde(default)]
82    pub denoise: Option<f32>,
83    /// Classifier-free guidance scale.  Per-job override; falls back
84    /// to `ModelSource.cliDefaults.cfgScale` when `None`.
85    #[serde(default)]
86    pub cfg_scale: Option<f32>,
87    /// Sampler choice (`euler`, `euler_a`, `dpm++2m`, ...).  Per-job
88    /// override; falls back to `ModelSource.cliDefaults.samplingMethod`.
89    #[serde(default)]
90    pub sampling_method: Option<String>,
91    #[serde(default = "default_image_dim")]
92    pub width: u32,
93    #[serde(default = "default_image_dim")]
94    pub height: u32,
95    #[serde(default = "default_steps")]
96    pub steps: u32,
97    #[serde(default)]
98    pub seed: Option<u64>,
99    #[serde(default = "default_image_ext")]
100    pub ext: String,
101}
102
103fn default_image_dim() -> u32 {
104    512
105}
106fn default_steps() -> u32 {
107    20
108}
109fn default_image_ext() -> String {
110    "webp".into()
111}
112
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct ChatMessage {
115    pub role: String,
116    pub content: String,
117}
118
119#[derive(Debug, Clone, Default, Serialize, Deserialize)]
120#[serde(rename_all = "camelCase")]
121pub struct LlmParams {
122    pub messages: Vec<ChatMessage>,
123    /// System prompt prepended to the conversation.  Engines that
124    /// don't accept a separate system role inline it as the first
125    /// chat turn.
126    #[serde(default)]
127    pub system: Option<String>,
128    #[serde(default = "default_max_tokens")]
129    pub max_tokens: u32,
130    #[serde(default = "default_temperature")]
131    pub temperature: f32,
132    #[serde(default)]
133    pub top_p: Option<f32>,
134    #[serde(default)]
135    pub stop: Option<Vec<String>>,
136    /// Strict-JSON output schema.  Engine-specific; passed through
137    /// verbatim to the backend (e.g. llama.cpp's `--grammar` JSON).
138    #[serde(default)]
139    pub json_schema: Option<serde_json::Value>,
140    /// Reasoning effort hint.  Currently honoured by Gemini-style
141    /// backends only; ignored elsewhere.
142    #[serde(default)]
143    pub reasoning: Option<String>,
144    /// Extra chat-template variables for this request (e.g.
145    /// `{"enable_thinking": false}`); overrides the model's own defaults.
146    #[serde(default, skip_serializing_if = "Option::is_none")]
147    pub chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
148}
149
150fn default_max_tokens() -> u32 {
151    512
152}
153fn default_temperature() -> f32 {
154    0.7
155}
156
157#[derive(Debug, Clone, Default, Serialize, Deserialize)]
158#[serde(rename_all = "camelCase")]
159pub struct AudioSttParams {
160    /// HTTPS URL to fetch the audio bytes from (e.g. R2 signed URL).
161    pub input_url: String,
162    #[serde(default)]
163    pub language: Option<String>,
164    /// Translate the audio to English (whisper `--translate`).
165    #[serde(default)]
166    pub translate: Option<bool>,
167    /// Initial prompt to bias the transcription.
168    #[serde(default)]
169    pub prompt: Option<String>,
170    /// Voice-activity detection.
171    #[serde(default)]
172    pub vad: Option<bool>,
173    /// Timestamp granularity: `"segment"` or `"word"`.
174    #[serde(default)]
175    pub timestamps: Option<String>,
176}
177
178#[derive(Debug, Clone, Default, Serialize, Deserialize)]
179#[serde(rename_all = "camelCase")]
180pub struct AudioTtsParams {
181    pub text: String,
182    #[serde(default = "default_voice")]
183    pub voice: String,
184    /// Playback speed multiplier (1.0 = natural pace).
185    #[serde(default)]
186    pub speed: Option<f32>,
187    /// Spoken-language hint (e.g. `"en"`, `"nl"`).
188    #[serde(default)]
189    pub language: Option<String>,
190    #[serde(default = "default_audio_ext")]
191    pub ext: String,
192}
193
194fn default_voice() -> String {
195    "default".into()
196}
197fn default_audio_ext() -> String {
198    "wav".into()
199}
200
201#[derive(Debug, Clone, Default, Serialize, Deserialize)]
202#[serde(rename_all = "camelCase")]
203pub struct VideoParams {
204    pub prompt: String,
205    #[serde(default)]
206    pub negative_prompt: Option<String>,
207    /// HTTPS URL to a base frame for image-to-video models.
208    #[serde(default)]
209    pub init_image_url: Option<String>,
210    #[serde(default = "default_video_seconds")]
211    pub seconds: f32,
212    /// Frame rate; defaults to backend-specific value when `None`.
213    #[serde(default)]
214    pub fps: Option<u32>,
215    #[serde(default = "default_image_dim")]
216    pub width: u32,
217    #[serde(default = "default_image_dim")]
218    pub height: u32,
219    #[serde(default = "default_video_ext")]
220    pub ext: String,
221}
222
223fn default_video_seconds() -> f32 {
224    2.0
225}
226fn default_video_ext() -> String {
227    "mp4".into()
228}
229
230#[derive(Debug, Clone, Serialize, Deserialize)]
231#[serde(tag = "kind", rename_all = "snake_case")]
232pub enum Task {
233    Image(ImageParams),
234    Llm(LlmParams),
235    AudioStt(AudioSttParams),
236    AudioTts(AudioTtsParams),
237    Video(VideoParams),
238}
239
240impl Task {
241    pub fn kind(&self) -> TaskKind {
242        match self {
243            Task::Image(_) => TaskKind::Image,
244            Task::Llm(_) => TaskKind::Llm,
245            Task::AudioStt(_) => TaskKind::AudioStt,
246            Task::AudioTts(_) => TaskKind::AudioTts,
247            Task::Video(_) => TaskKind::Video,
248        }
249    }
250}
251
252// ---------------------------------------------------------------------------
253// Per-kind task results
254// ---------------------------------------------------------------------------
255
256#[derive(Debug, Clone)]
257pub enum TaskResult {
258    /// Binary image (webp/png/...) with the chosen extension.
259    Image { bytes: Vec<u8>, ext: String },
260    /// JSON response from an LLM call (free-form to mirror common APIs).
261    Llm { json: serde_json::Value },
262    /// JSON transcript.
263    AudioStt { json: serde_json::Value },
264    /// Binary audio (wav/mp3/...) with the chosen extension.
265    AudioTts { bytes: Vec<u8>, ext: String },
266    /// Binary video (mp4/webm/...) with the chosen extension.
267    Video { bytes: Vec<u8>, ext: String },
268}
269
270impl TaskResult {
271    pub fn kind(&self) -> TaskKind {
272        match self {
273            TaskResult::Image { .. } => TaskKind::Image,
274            TaskResult::Llm { .. } => TaskKind::Llm,
275            TaskResult::AudioStt { .. } => TaskKind::AudioStt,
276            TaskResult::AudioTts { .. } => TaskKind::AudioTts,
277            TaskResult::Video { .. } => TaskKind::Video,
278        }
279    }
280}
281
282// ---------------------------------------------------------------------------
283// Worker capabilities + registration
284// ---------------------------------------------------------------------------
285
286#[derive(Debug, Clone, Serialize, Deserialize)]
287pub struct WorkerCapabilities {
288    #[serde(rename = "machineName")]
289    pub machine_name: String,
290    pub username: String,
291    #[serde(rename = "agentVersion")]
292    pub agent_version: String,
293    pub engine: String,
294    #[serde(rename = "vramTotalGb")]
295    pub vram_total_gb: f32,
296    #[serde(rename = "vramThresholdGb")]
297    pub vram_threshold_gb: f32,
298    #[serde(rename = "autoEnabled")]
299    pub auto_enabled: bool,
300    #[serde(rename = "autoStart")]
301    pub auto_start: bool,
302    /// Flat list of models, kept for backward compat with the existing studio
303    /// API that doesn't know about kinds yet.  Equivalent to
304    /// `supported_models_per_kind[Image]`.
305    #[serde(rename = "supportedModels")]
306    pub supported_models: Vec<String>,
307    /// New: task kinds this worker can serve.
308    #[serde(rename = "taskKinds", default)]
309    pub task_kinds: Vec<TaskKind>,
310    /// New: per-kind supported model ids.
311    #[serde(rename = "supportedModelsPerKind", default)]
312    pub supported_models_per_kind: BTreeMap<TaskKind, Vec<String>>,
313}
314
315// ---------------------------------------------------------------------------
316// Auto-register wire format — the only registration path.
317//
318// Worker POSTs `/workers/register-request` with hostname / username /
319// VRAM / supported models, gets a `requestId` back, and polls
320// `/workers/register-requests/:id` until the operator approves or
321// rejects from the studio dashboard.
322// ---------------------------------------------------------------------------
323
324#[derive(Debug, Clone, Serialize)]
325pub struct AutoRegisterRequest {
326    /// Per-install UUID stable across worker restarts on the same
327    /// machine.  The studio uses it to dedup re-submissions.
328    #[serde(rename = "installId")]
329    pub install_id: String,
330    /// SHA-256 hex of the worker-side `registration_secret`.  The
331    /// worker keeps the secret locally and presents it as a Bearer
332    /// token when polling for status; only the hash leaves the box.
333    #[serde(rename = "registrationSecretHash")]
334    pub registration_secret_hash: String,
335    /// Full capability snapshot — hostname, username, engine, VRAM,
336    /// supported models so the operator can decide.
337    pub capabilities: WorkerCapabilities,
338    /// `studio-worker/<version>` so the operator sees stale clients.
339    #[serde(rename = "userAgent")]
340    pub user_agent: String,
341}
342
343#[derive(Debug, Clone, Deserialize)]
344pub struct AutoRegisterRequestResponse {
345    #[serde(rename = "requestId")]
346    pub request_id: String,
347    /// Always `"pending"` on first response.  Idempotent dedup may
348    /// return the existing requestId for the same
349    /// `(installId, sourceIp)` tuple — status is still "pending".
350    pub status: String,
351}
352
353/// Tagged union returned by `GET /workers/register-requests/:id`.
354#[derive(Debug, Clone, Deserialize, PartialEq, Eq)]
355#[serde(rename_all = "snake_case", tag = "status")]
356pub enum RegisterStatus {
357    Pending,
358    Approved {
359        #[serde(rename = "workerId")]
360        worker_id: String,
361        #[serde(rename = "authToken")]
362        auth_token: String,
363    },
364    Rejected {
365        #[serde(default)]
366        reason: String,
367    },
368}
369
370#[derive(Debug, Clone, Serialize)]
371pub struct HeartbeatRequest {
372    pub capabilities: WorkerCapabilities,
373    #[serde(rename = "currentJobId", skip_serializing_if = "Option::is_none")]
374    pub current_job_id: Option<String>,
375}
376
377// ---------------------------------------------------------------------------
378// ModelSource — download spec the studio attaches to every offer so
379// the worker can fetch + run any model without per-model knowledge.
380// Mirrors `WorkerModelSource` on the TS side.
381// ---------------------------------------------------------------------------
382
383#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
384#[serde(rename_all = "kebab-case")]
385pub enum ModelFileRole {
386    DiffusionModel,
387    TextEncoder,
388    /// Vision tower / mmproj for a multimodal text encoder (e.g.
389    /// Qwen2.5-VL ViT). Maps to `sd-cli --llm_vision`; required by
390    /// instruction-edit models that condition on the reference image.
391    TextEncoderVision,
392    Vae,
393    Lora,
394    Model,
395}
396
397#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
398#[serde(rename_all = "kebab-case")]
399pub enum ModelEngine {
400    SdCpp,
401    LlamaCpp,
402    /// ONNX Runtime image engine (pykeio/ort).  Serves the LaMa
403    /// object-removal model used by Find-the-Differences removals.
404    Onnx,
405    /// Streaming speech-to-text (parakeet-rs).  Local-only: served by a
406    /// loaded model over the LAN streaming listener, never by a job.
407    Parakeet,
408    Synthetic,
409}
410
411#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
412#[serde(rename_all = "camelCase")]
413pub struct ModelFile {
414    pub role: ModelFileRole,
415    pub url: String,
416    pub filename: String,
417    #[serde(default, skip_serializing_if = "Option::is_none")]
418    pub approx_bytes: Option<u64>,
419    /// Hex sha256 of the file's bytes.  When present the worker
420    /// verifies the downloaded body against it before committing the
421    /// file to the cache; absent (legacy registry rows) means
422    /// Content-Length is the only integrity check.
423    #[serde(default, skip_serializing_if = "Option::is_none")]
424    pub sha256: Option<String>,
425}
426
427#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
428#[serde(rename_all = "camelCase")]
429pub struct ModelCliDefaults {
430    pub cfg_scale: f32,
431    pub steps: u32,
432    pub width: u32,
433    pub height: u32,
434    #[serde(default, skip_serializing_if = "Option::is_none")]
435    pub sampling_method: Option<String>,
436    /// `--flow-shift` for Flow models (SD3.x / WAN / Qwen-Image). Model-level constant.
437    #[serde(default, skip_serializing_if = "Option::is_none")]
438    pub flow_shift: Option<f32>,
439    /// Qwen-Image zero-cond-t, mandatory for Qwen-Image edit quality. Emitted as
440    /// `--model-args qwen_image_zero_cond_t=true` (sd.cpp master-9xx dropped the dedicated flag).
441    #[serde(default, skip_serializing_if = "Option::is_none")]
442    pub zero_cond_t: Option<bool>,
443    /// `--offload-to-cpu` — keep weights in RAM, stream to VRAM, so a large model fits a small card.
444    #[serde(default, skip_serializing_if = "Option::is_none")]
445    pub offload_to_cpu: Option<bool>,
446    /// LLM context window in tokens (the KV-cache size).
447    #[serde(default, skip_serializing_if = "Option::is_none")]
448    pub context_size: Option<u32>,
449    /// LLM chat-template variables every request gets unless it overrides
450    /// them (e.g. `{"enable_thinking": false}`).
451    #[serde(default, skip_serializing_if = "Option::is_none")]
452    pub chat_template_kwargs: Option<serde_json::Map<String, serde_json::Value>>,
453    /// `--mmap` — memory-map the weights, so offloaded models stream without a RAM copy.
454    #[serde(default, skip_serializing_if = "Option::is_none")]
455    pub mmap: Option<bool>,
456    /// `--max-vram <GiB>` — device budget for managed weights and buffers, so a large model shares
457    /// the card with other tenants (pairs with `offload_to_cpu`).
458    #[serde(default, skip_serializing_if = "Option::is_none")]
459    pub max_vram_gib: Option<f32>,
460}
461
462#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
463#[serde(rename_all = "camelCase")]
464pub struct ModelSource {
465    pub engine: ModelEngine,
466    pub files: Vec<ModelFile>,
467    pub cli_defaults: ModelCliDefaults,
468}
469
470// ---------------------------------------------------------------------------
471// JobClaim — every offer carries `task` + `model_source` populated by the
472// studio's registry resolver.  No legacy `prompt + ext` shadow fields.
473// ---------------------------------------------------------------------------
474
475#[derive(Debug, Clone, Deserialize)]
476#[serde(rename_all = "camelCase")]
477pub struct JobClaim {
478    pub job_id: String,
479    #[allow(dead_code)]
480    pub game_id: String,
481    pub asset_name: String,
482    pub model: String,
483    pub vram_gb_estimate: f32,
484    /// Structured task payload.  Required — the studio refuses to
485    /// promote a job without one.  Worker treats a missing `task`
486    /// as a protocol_violation.
487    pub task: Task,
488    /// Download + engine + CLI defaults the studio resolved from its
489    /// model registry.  Required — `synthetic` is just another engine
490    /// option, not a fallback for missing rows.
491    pub model_source: ModelSource,
492}
493
494#[derive(Debug, Clone, Serialize)]
495pub struct FailRequest {
496    pub error: String,
497    pub retryable: bool,
498}
499
500#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
501pub struct LogEntry {
502    pub ts: String,
503    pub level: String,
504    pub category: String,
505    pub message: String,
506    #[serde(rename = "jobId", default, skip_serializing_if = "Option::is_none")]
507    pub job_id: Option<String>,
508}
509
510#[derive(Debug, Clone, Serialize, Deserialize)]
511pub struct LogBatch {
512    pub entries: Vec<LogEntry>,
513}
514
515// ---------------------------------------------------------------------------
516// Release feed (auto-update)
517// ---------------------------------------------------------------------------
518
519/// Subset of the GitHub Releases API we care about.
520#[derive(Debug, Clone, Deserialize)]
521pub struct GithubRelease {
522    pub tag_name: String,
523    #[serde(default)]
524    pub prerelease: bool,
525    #[serde(default)]
526    pub draft: bool,
527    #[serde(default)]
528    pub assets: Vec<GithubReleaseAsset>,
529}
530
531#[derive(Debug, Clone, Deserialize)]
532pub struct GithubReleaseAsset {
533    pub name: String,
534    pub browser_download_url: String,
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540
541    fn synthetic_model_source_json() -> serde_json::Value {
542        serde_json::json!({
543            "engine": "synthetic",
544            "files": [],
545            "cliDefaults": {
546                "cfgScale": 1.0,
547                "steps": 8,
548                "width": 1024,
549                "height": 1024,
550            },
551        })
552    }
553
554    #[test]
555    fn job_claim_requires_task_and_model_source() {
556        // Without `task` or `modelSource` the offer must fail to
557        // deserialise — no silent fallback to an Option/Image default.
558        let bare = serde_json::json!({
559            "jobId": "j-1",
560            "gameId": "g-1",
561            "assetName": "g-1/creatures/x",
562            "model": "synthetic-image",
563            "vramGbEstimate": 1.0,
564        });
565        assert!(
566            serde_json::from_value::<JobClaim>(bare).is_err(),
567            "JobClaim must reject missing task + modelSource"
568        );
569    }
570
571    #[test]
572    fn job_claim_with_explicit_llm_task() {
573        let json = serde_json::json!({
574            "jobId": "j-2",
575            "gameId": "g-1",
576            "assetName": "g-1/conversations/x",
577            "model": "llama-3.1-8b",
578            "vramGbEstimate": 8.0,
579            "task": {
580                "kind": "llm",
581                "messages": [{"role": "user", "content": "hi"}],
582                "maxTokens": 32,
583                "temperature": 0.5,
584            },
585            "modelSource": synthetic_model_source_json(),
586        });
587        let claim: JobClaim = serde_json::from_value(json).unwrap();
588        match claim.task {
589            Task::Llm(p) => {
590                assert_eq!(p.messages.len(), 1);
591                assert_eq!(p.max_tokens, 32);
592            }
593            other => panic!("expected llm, got {:?}", other),
594        }
595    }
596
597    #[test]
598    fn job_claim_with_explicit_image_task() {
599        let json = serde_json::json!({
600            "jobId": "j-3",
601            "gameId": "g-1",
602            "assetName": "g-1/creatures/y",
603            "model": "synthetic-image",
604            "vramGbEstimate": 8.0,
605            "task": {
606                "kind": "image",
607                "prompt": "a koi",
608                "width": 1024,
609                "height": 1024,
610                "steps": 30,
611                "ext": "png",
612            },
613            "modelSource": synthetic_model_source_json(),
614        });
615        let claim: JobClaim = serde_json::from_value(json).unwrap();
616        match claim.task {
617            Task::Image(p) => {
618                assert_eq!(p.prompt, "a koi");
619                assert_eq!(p.width, 1024);
620                assert_eq!(p.ext, "png");
621            }
622            other => panic!("expected image, got {:?}", other),
623        }
624    }
625
626    #[test]
627    fn image_params_round_trips_with_new_fields() {
628        let json = serde_json::json!({
629            "kind": "image",
630            "prompt": "a stone golem",
631            "negativePrompt": "text, watermark, low quality",
632            "initImageUrl": "https://example.invalid/t2-golem-stone/latest.webp",
633            "denoise": 0.55,
634            "cfgScale": 7.5,
635            "samplingMethod": "dpm++2m",
636            "width": 768,
637            "height": 512,
638            "steps": 30,
639            "seed": 1234,
640            "ext": "webp",
641        });
642        let task: Task = serde_json::from_value(json).unwrap();
643        match task {
644            Task::Image(p) => {
645                assert_eq!(p.prompt, "a stone golem");
646                assert_eq!(
647                    p.negative_prompt.as_deref(),
648                    Some("text, watermark, low quality")
649                );
650                assert_eq!(
651                    p.init_image_url.as_deref(),
652                    Some("https://example.invalid/t2-golem-stone/latest.webp")
653                );
654                assert!((p.denoise.unwrap() - 0.55).abs() < 1e-6);
655                assert!((p.cfg_scale.unwrap() - 7.5).abs() < 1e-6);
656                assert_eq!(p.sampling_method.as_deref(), Some("dpm++2m"));
657                assert_eq!(p.width, 768);
658                assert_eq!(p.height, 512);
659                assert_eq!(p.steps, 30);
660                assert_eq!(p.seed, Some(1234));
661            }
662            other => panic!("expected image, got {:?}", other),
663        }
664    }
665
666    #[test]
667    fn image_params_defaults_when_optional_fields_absent() {
668        let json = serde_json::json!({
669            "kind": "image",
670            "prompt": "a fox"
671        });
672        let task: Task = serde_json::from_value(json).unwrap();
673        match task {
674            Task::Image(p) => {
675                assert_eq!(p.prompt, "a fox");
676                assert!(p.negative_prompt.is_none());
677                assert!(p.init_image_url.is_none());
678                assert!(p.denoise.is_none());
679                assert!(p.cfg_scale.is_none());
680                assert!(p.sampling_method.is_none());
681                assert_eq!(p.width, 512);
682                assert_eq!(p.height, 512);
683                assert_eq!(p.steps, 20);
684                assert_eq!(p.ext, "webp");
685            }
686            other => panic!("expected image, got {:?}", other),
687        }
688    }
689
690    #[test]
691    fn task_kinds_round_trip_via_json() {
692        for kind in TaskKind::ALL {
693            let s = serde_json::to_string(&kind).unwrap();
694            let back: TaskKind = serde_json::from_str(&s).unwrap();
695            assert_eq!(kind, back);
696        }
697    }
698}
studio_worker/types.rs

studio_worker/
types.rs
