Skip to main content

novelai_bridge/
types.rs

1use serde::{Deserialize, Serialize};
2
3const DEFAULT_STEPS: u32 = 23;
4const DEFAULT_SCALE: f32 = 5.0;
5const QUALITY_TAGS: &str = ", very aesthetic, masterpiece, no text";
6const UC_PRESET_STRONG: &str = ", lowres, artistic error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, dithering, halftone, screentone, multiple views, logo, too many watermarks, negative space, blank page, ";
7const UC_PRESET_LIGHT: &str = ", lowres, artistic error, scan artifacts, worst quality, bad quality, jpeg artifacts, multiple views, very displeasing, too many watermarks, negative space, blank page, ";
8const UC_PRESET_FURRY_FOCUS: &str = ", {worst quality}, distracting watermark, unfinished, bad quality, {widescreen}, upscale, {sequence}, {{grandfathered content}}, blurred foreground, chromatic aberration, sketch, everyone, [sketch background], simple, [flat colors], ych (character), outline, multiple scenes, [[horror (theme)]], comic, ";
9const UC_PRESET_HUMAN_FOCUS: &str = ", lowres, artistic error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, dithering, halftone, screentone, multiple views, logo, too many watermarks, negative space, blank page, @_@, mismatched pupils, glowing eyes, bad anatomy, ";
10
11/// Selects how the API key is resolved.
12#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
13#[serde(rename_all = "snake_case")]
14pub enum ApiKeySource {
15    /// Read the API key from the `NOVELAI_API_KEY` environment variable.
16    Environment,
17    /// Use an explicitly provided API key value.
18    Inline { value: String },
19}
20
21/// Supported NovelAI image models.
22#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
23pub enum Model {
24    /// `nai-diffusion-4-5-full`
25    #[serde(rename = "nai-diffusion-4-5-full")]
26    #[default]
27    NaiDiffusion45Full,
28    /// `nai-diffusion-4-5-curated`
29    #[serde(rename = "nai-diffusion-4-5-curated")]
30    NaiDiffusion45Curated,
31    /// `nai-diffusion-4-full`
32    #[serde(rename = "nai-diffusion-4-full")]
33    NaiDiffusion4Full,
34    /// `nai-diffusion-4-curated`
35    #[serde(rename = "nai-diffusion-4-curated")]
36    NaiDiffusion4Curated,
37    /// `nai-diffusion-3`
38    #[serde(rename = "nai-diffusion-3")]
39    NaiDiffusion3,
40    /// `nai-diffusion-3-furry`
41    #[serde(rename = "nai-diffusion-3-furry")]
42    NaiDiffusion3Furry,
43}
44
45impl Model {
46    /// Returns the API model identifier.
47    pub fn as_str(&self) -> &'static str {
48        match self {
49            Self::NaiDiffusion45Full => "nai-diffusion-4-5-full",
50            Self::NaiDiffusion45Curated => "nai-diffusion-4-5-curated",
51            Self::NaiDiffusion4Full => "nai-diffusion-4-full",
52            Self::NaiDiffusion4Curated => "nai-diffusion-4-curated",
53            Self::NaiDiffusion3 => "nai-diffusion-3",
54            Self::NaiDiffusion3Furry => "nai-diffusion-3-furry",
55        }
56    }
57
58    /// Returns `true` when the model uses the V4/V4.5 request shape.
59    pub fn is_v4(&self) -> bool {
60        matches!(
61            self,
62            Self::NaiDiffusion45Full
63                | Self::NaiDiffusion45Curated
64                | Self::NaiDiffusion4Full
65                | Self::NaiDiffusion4Curated
66        )
67    }
68
69    /// Returns `true` when the model is a V4.5 variant.
70    pub fn is_v45(&self) -> bool {
71        matches!(self, Self::NaiDiffusion45Full | Self::NaiDiffusion45Curated)
72    }
73
74    /// Returns the model key used by vibe metadata payloads.
75    pub fn vibe_model_key(&self) -> &'static str {
76        match self {
77            Self::NaiDiffusion45Full => "v4-5full",
78            Self::NaiDiffusion45Curated => "v4-5curated",
79            Self::NaiDiffusion4Full => "v4full",
80            Self::NaiDiffusion4Curated => "v4curated",
81            Self::NaiDiffusion3 => "v3",
82            Self::NaiDiffusion3Furry => "v3furry",
83        }
84    }
85}
86
87/// Supported image samplers.
88#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
89pub enum Sampler {
90    /// `k_euler`
91    #[serde(rename = "k_euler")]
92    KEuler,
93    /// `k_euler_ancestral`
94    #[serde(rename = "k_euler_ancestral")]
95    #[default]
96    KEulerAncestral,
97    /// `k_dpm_2`
98    #[serde(rename = "k_dpm_2")]
99    KDpm2,
100    /// `k_dpm_2_ancestral`
101    #[serde(rename = "k_dpm_2_ancestral")]
102    KDpm2Ancestral,
103    /// `k_dpmpp_2m`
104    #[serde(rename = "k_dpmpp_2m")]
105    KDpmpp2m,
106    /// `k_dpmpp_2s_ancestral`
107    #[serde(rename = "k_dpmpp_2s_ancestral")]
108    KDpmpp2sAncestral,
109    /// `k_dpmpp_sde`
110    #[serde(rename = "k_dpmpp_sde")]
111    KDpmppSde,
112    /// `ddim`
113    #[serde(rename = "ddim")]
114    Ddim,
115}
116
117impl Sampler {
118    /// Returns the API sampler identifier.
119    pub fn as_str(&self) -> &'static str {
120        match self {
121            Self::KEuler => "k_euler",
122            Self::KEulerAncestral => "k_euler_ancestral",
123            Self::KDpm2 => "k_dpm_2",
124            Self::KDpm2Ancestral => "k_dpm_2_ancestral",
125            Self::KDpmpp2m => "k_dpmpp_2m",
126            Self::KDpmpp2sAncestral => "k_dpmpp_2s_ancestral",
127            Self::KDpmppSde => "k_dpmpp_sde",
128            Self::Ddim => "ddim",
129        }
130    }
131}
132
133/// Supported noise schedules.
134#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
135pub enum NoiseSchedule {
136    /// `karras`
137    #[serde(rename = "karras")]
138    #[default]
139    Karras,
140    /// `exponential`
141    #[serde(rename = "exponential")]
142    Exponential,
143    /// `polyexponential`
144    #[serde(rename = "polyexponential")]
145    Polyexponential,
146}
147
148impl NoiseSchedule {
149    /// Returns the API noise schedule identifier.
150    pub fn as_str(&self) -> &'static str {
151        match self {
152            Self::Karras => "karras",
153            Self::Exponential => "exponential",
154            Self::Polyexponential => "polyexponential",
155        }
156    }
157}
158
159/// Built-in UC presets supported by NovelAI.
160#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
161#[serde(rename_all = "snake_case")]
162pub enum UcPreset {
163    /// Strong default UC preset.
164    Strong,
165    /// Light default UC preset.
166    #[default]
167    Light,
168    /// Furry-focused UC preset.
169    FurryFocus,
170    /// Human-focused UC preset.
171    HumanFocus,
172    /// No preset text.
173    None,
174}
175
176impl UcPreset {
177    /// Returns the integer value expected by the API.
178    pub fn as_api_value(&self) -> i32 {
179        match self {
180            Self::Strong => 0,
181            Self::Light => 1,
182            Self::FurryFocus => 2,
183            Self::HumanFocus => 3,
184            Self::None => 4,
185        }
186    }
187
188    /// Returns the preset text appended to the negative prompt.
189    pub fn preset_text(&self) -> &'static str {
190        match self {
191            Self::Strong => UC_PRESET_STRONG,
192            Self::Light => UC_PRESET_LIGHT,
193            Self::FurryFocus => UC_PRESET_FURRY_FOCUS,
194            Self::HumanFocus => UC_PRESET_HUMAN_FOCUS,
195            Self::None => "",
196        }
197    }
198}
199
200/// Supported image stream modes.
201#[derive(Copy, Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)]
202#[serde(rename_all = "snake_case")]
203pub enum StreamMode {
204    /// Server-sent events.
205    #[default]
206    Sse,
207}
208
209impl StreamMode {
210    /// Returns the API stream identifier.
211    pub fn as_str(&self) -> &'static str {
212        match self {
213            Self::Sse => "sse",
214        }
215    }
216}
217
218/// Supported generated image output formats.
219#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
220pub enum ImageFormat {
221    /// PNG output.
222    #[serde(rename = "png")]
223    Png,
224    /// WebP output.
225    #[serde(rename = "webp")]
226    Webp,
227}
228
229impl ImageFormat {
230    /// Returns the API output format identifier.
231    pub fn as_str(&self) -> &'static str {
232        match self {
233            Self::Png => "png",
234            Self::Webp => "webp",
235        }
236    }
237}
238
239/// Image dimensions in pixels.
240#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
241pub struct ImageSize {
242    /// Width in pixels.
243    pub width: u32,
244    /// Height in pixels.
245    pub height: u32,
246}
247
248impl ImageSize {
249    /// Returns the default portrait canvas.
250    pub fn portrait() -> Self {
251        Self {
252            width: 832,
253            height: 1216,
254        }
255    }
256
257    /// Returns the default landscape canvas.
258    pub fn landscape() -> Self {
259        Self {
260            width: 1216,
261            height: 832,
262        }
263    }
264
265    /// Returns the default square canvas.
266    pub fn square() -> Self {
267        Self {
268            width: 1024,
269            height: 1024,
270        }
271    }
272
273    /// Returns the large portrait director/reference canvas.
274    pub fn large_portrait() -> Self {
275        Self {
276            width: 1024,
277            height: 1536,
278        }
279    }
280
281    /// Returns the large landscape director/reference canvas.
282    pub fn large_landscape() -> Self {
283        Self {
284            width: 1536,
285            height: 1024,
286        }
287    }
288}
289
290impl Default for ImageSize {
291    fn default() -> Self {
292        Self::portrait()
293    }
294}
295
296/// Normalized character position in the `[0, 1]` range.
297#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq)]
298pub struct CharacterPosition {
299    /// Horizontal position in the `[0, 1]` range.
300    pub x: f32,
301    /// Vertical position in the `[0, 1]` range.
302    pub y: f32,
303}
304
305impl Default for CharacterPosition {
306    fn default() -> Self {
307        Self { x: 0.5, y: 0.5 }
308    }
309}
310
311/// Character prompt input used by V4/V4.5 models.
312#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
313pub struct Character {
314    /// Positive prompt for the character.
315    pub prompt: String,
316    /// Optional negative prompt for the character.
317    pub negative_prompt: Option<String>,
318    /// Desired character position.
319    pub position: CharacterPosition,
320    /// Whether this character is enabled.
321    pub enabled: bool,
322}
323
324/// Character reference behavior used by V4.5 director references.
325#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
326pub enum CharacterReferenceType {
327    /// Character-only reference.
328    #[serde(rename = "character")]
329    Character,
330    /// Style-only reference.
331    #[serde(rename = "style")]
332    Style,
333    /// Combined character and style reference.
334    #[serde(rename = "character&style")]
335    CharacterAndStyle,
336}
337
338/// Character reference input used by V4.5 models.
339#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
340pub struct CharacterReference {
341    /// Base64 payload or image data URL for the reference image.
342    pub image: String,
343    /// How the reference should be interpreted.
344    pub reference_type: CharacterReferenceType,
345    /// Fidelity value used to derive secondary strength.
346    pub fidelity: f32,
347    /// Overall reference strength.
348    pub strength: f32,
349}
350
351/// Image-to-image input.
352#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
353pub struct Img2ImgRequest {
354    /// Base64 payload or image data URL for the source image.
355    pub image: String,
356    /// Strength value in the `(0, 1)` range.
357    pub strength: f32,
358    /// Noise value in the `[0, 0.99]` range.
359    pub noise: f32,
360    /// Optional base64 payload or image data URL for the inpaint mask.
361    pub mask: Option<String>,
362}
363
364/// One controlnet-style reference input.
365#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
366pub struct ControlNetInput {
367    /// Base64 payload or image data URL for the source image.
368    pub image: String,
369    /// Information-extracted value used for vibe encoding.
370    pub info_extracted: f32,
371    /// Per-image reference strength.
372    pub strength: f32,
373    /// Model used when encoding the vibe payload.
374    pub model: Model,
375    /// Optional precomputed base64 vibe payload.
376    pub vibe_data_cache: Option<String>,
377}
378
379/// Controlnet-style configuration.
380#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
381pub struct ControlNetConfig {
382    /// One or more controlnet images.
383    pub images: Vec<ControlNetInput>,
384    /// Overall controlnet strength.
385    pub strength: f32,
386}
387
388/// Request used by [`crate::Client::encode_vibe`].
389#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
390pub struct EncodeVibeRequest {
391    /// Base64 payload or image data URL for the source image.
392    pub image: String,
393    /// Information-extracted value in the `[0.01, 1.0]` range.
394    pub information_extracted: f32,
395    /// Target model for the encoded vibe payload.
396    pub model: Model,
397}
398
399/// Supported NovelAI director tools.
400#[derive(Copy, Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
401#[serde(rename_all = "kebab-case")]
402pub enum DirectorTool {
403    /// Line art extraction.
404    Lineart,
405    /// Sketch conversion.
406    Sketch,
407    /// Background removal.
408    BgRemoval,
409    /// Emotion transfer.
410    Emotion,
411    /// Declutter cleanup.
412    Declutter,
413    /// Colorization.
414    Colorize,
415}
416
417/// Request used by [`crate::Client::run_director_tool`].
418#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
419pub struct RunDirectorToolRequest {
420    /// Director tool to run.
421    pub tool: DirectorTool,
422    /// Base64 payload or image data URL for the source image.
423    pub image: String,
424    /// Optional prompt for tools that accept it.
425    pub prompt: Option<String>,
426    /// Optional defry value for tools that accept it.
427    pub defry: Option<u8>,
428}
429
430/// Request used by [`crate::Client::generate`].
431#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
432pub struct GenerateImageRequest {
433    /// Positive prompt text.
434    pub prompt: String,
435    /// Target image model.
436    pub model: Model,
437    /// Output canvas size.
438    pub size: ImageSize,
439    /// Optional negative prompt text.
440    pub negative_prompt: Option<String>,
441    /// Whether quality tags are appended to the positive prompt.
442    pub quality: bool,
443    /// Built-in UC preset appended to the negative prompt.
444    pub uc_preset: UcPreset,
445    /// Sampling step count in the `[1, 50]` range.
446    pub steps: u32,
447    /// CFG scale in the `[0, 10]` range.
448    pub scale: f32,
449    /// Sampler to use.
450    pub sampler: Sampler,
451    /// Noise schedule to use.
452    pub noise_schedule: NoiseSchedule,
453    /// Seed value. `0` requests a locally generated seed.
454    pub seed: i64,
455    /// Number of images to generate in the `[1, 4]` range.
456    pub n_samples: u32,
457    /// CFG rescale value in the `[0, 1]` range.
458    pub cfg_rescale: f32,
459    /// Enables NovelAI's variety boost behavior.
460    pub variety_boost: bool,
461    /// Optional image-to-image configuration.
462    pub i2i: Option<Img2ImgRequest>,
463    /// Optional controlnet-style configuration.
464    pub controlnet: Option<ControlNetConfig>,
465    /// Optional V4.5 character references.
466    pub character_references: Option<Vec<CharacterReference>>,
467    /// Optional V4/V4.5 character prompts.
468    pub characters: Option<Vec<Character>>,
469    /// Optional override for `use_coords`.
470    pub use_coords: Option<bool>,
471    /// Optional output image format.
472    pub image_format: Option<ImageFormat>,
473}
474
475impl Default for GenerateImageRequest {
476    fn default() -> Self {
477        Self {
478            prompt: String::new(),
479            model: Model::default(),
480            size: ImageSize::default(),
481            negative_prompt: None,
482            quality: true,
483            uc_preset: UcPreset::default(),
484            steps: DEFAULT_STEPS,
485            scale: DEFAULT_SCALE,
486            sampler: Sampler::default(),
487            noise_schedule: NoiseSchedule::default(),
488            seed: 0,
489            n_samples: 1,
490            cfg_rescale: 0.0,
491            variety_boost: false,
492            i2i: None,
493            controlnet: None,
494            character_references: None,
495            characters: None,
496            use_coords: None,
497            image_format: None,
498        }
499    }
500}
501
502impl GenerateImageRequest {
503    /// Returns the positive prompt after applying built-in quality tags.
504    pub fn processed_prompt(&self) -> String {
505        if self.quality {
506            format!("{}{}", self.prompt, QUALITY_TAGS)
507        } else {
508            self.prompt.clone()
509        }
510    }
511
512    /// Returns the negative prompt after applying the selected UC preset text.
513    pub fn processed_negative_prompt(&self) -> String {
514        format!(
515            "{}{}",
516            self.negative_prompt.clone().unwrap_or_default(),
517            self.uc_preset.preset_text()
518        )
519    }
520}
521
522/// Request used by [`crate::Client::generate_stream`].
523#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
524pub struct GenerateImageStreamRequest {
525    /// Base image-generation request.
526    pub base: GenerateImageRequest,
527    /// Stream delivery mode.
528    pub stream: StreamMode,
529}
530
531/// Generated image payload.
532#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
533pub struct GeneratedImage {
534    /// Raw image bytes.
535    pub bytes: Vec<u8>,
536    /// Best-effort MIME type inferred from the response.
537    pub mime_type: Option<String>,
538    /// Seed associated with the generated image when known.
539    pub seed: Option<i64>,
540}
541
542/// Parsed stream chunk emitted by NovelAI image streaming.
543#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
544pub struct ImageStreamChunk {
545    /// Stream event type.
546    pub event_type: String,
547    /// Sample index.
548    pub samp_ix: u32,
549    /// Optional step index.
550    pub step_ix: Option<u32>,
551    /// Generation identifier.
552    pub gen_id: u32,
553    /// Optional sigma value.
554    pub sigma: Option<f32>,
555    /// Base64-encoded image payload fragment.
556    pub image: String,
557}
558
559impl ImageStreamChunk {
560    /// Parses a single SSE line into an [`ImageStreamChunk`].
561    pub fn from_sse_line(line: &str) -> Result<Self, serde_json::Error> {
562        let payload = line
563            .strip_prefix("data:")
564            .map(str::trim_start)
565            .unwrap_or(line);
566        serde_json::from_str(payload)
567    }
568}
569
570/// Subscription summary returned by [`crate::Client::get_subscription`].
571#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
572pub struct SubscriptionInfo {
573    /// Combined fixed and purchased Anlas balance.
574    pub anlas_balance: i64,
575    /// Whether the account is effectively on Opus tier.
576    pub is_opus: bool,
577    /// Numeric tier value.
578    pub tier: i32,
579    /// Lowercase tier name.
580    pub tier_name: String,
581    /// Optional expiration timestamp in milliseconds since the Unix epoch.
582    pub expires_at_ms: Option<u64>,
583}