Skip to main content

zai_rs/model/gen_image/
image_request.rs

1use serde::Serialize;
2use validator::Validate;
3
4use super::super::traits::*;
5
6/// Request body for image generation
7#[derive(Debug, Clone, Serialize, Validate)]
8pub struct ImageGenBody<N>
9where
10    N: ModelName + ImageGen + Serialize,
11{
12    /// The model to use for image generation
13    pub model: N,
14    /// Text description of the desired image
15    #[serde(skip_serializing_if = "Option::is_none")]
16    #[validate(length(min = 1, max = 4000))]
17    pub prompt: Option<String>,
18    /// Image generation quality
19    /// - HD: generates more refined and detailed images with higher
20    ///   consistency, takes ~20 seconds
21    /// - Standard: fast image generation, suitable for scenarios requiring
22    ///   speed, takes ~5-10 seconds This parameter only supports
23    ///   cogview-4-250304
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub quality: Option<ImageQuality>,
26    /// Image size
27    /// Recommended values: 1024x1024 (default), 768x1344, 864x1152, 1344x768,
28    /// 1152x864, 1440x720, 720x1440 Custom dimensions: width and height
29    /// must be between 512-2048px, divisible by 16, and total pixels must
30    /// not exceed 2^21 (2,097,152)
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub size: Option<ImageSize>,
33    /// Whether to add watermark to AI generated images
34    #[serde(skip_serializing_if = "Option::is_none")]
35    pub watermark_enabled: Option<bool>,
36    /// Unique ID of the end user to help platform intervene against violations,
37    /// illegal content generation, or other abusive behaviors
38    #[serde(skip_serializing_if = "Option::is_none")]
39    #[validate(length(min = 6, max = 128))]
40    pub user_id: Option<String>,
41}
42
43/// Image generation quality options
44#[derive(Debug, Clone, Serialize)]
45pub enum ImageQuality {
46    /// High quality - more refined and detailed images
47    #[serde(rename = "hd")]
48    Hd,
49    /// Standard quality - faster generation
50    #[serde(rename = "standard")]
51    Standard,
52}
53
54/// Image size options
55///
56/// Recommended sizes:
57/// - 1024x1024 (default)
58/// - 768x1344
59/// - 864x1152
60/// - 1344x768
61/// - 1152x864
62/// - 1440x720
63/// - 720x1440
64///
65/// Custom sizes must satisfy:
66/// - Width and height between 512-2048px
67/// - Both dimensions divisible by 16
68/// - Total pixels <= 2^21 (2,097,152)
69#[derive(Debug, Clone)]
70pub enum ImageSize {
71    /// 1024x1024 pixels
72    Size1024x1024,
73    /// 768x1344 pixels
74    Size768x1344,
75    /// 864x1152 pixels
76    Size864x1152,
77    /// 1344x768 pixels
78    Size1344x768,
79    /// 1152x864 pixels
80    Size1152x864,
81    /// 1440x720 pixels
82    Size1440x720,
83    /// 720x1440 pixels
84    Size720x1440,
85    /// Custom dimensions in width x height format
86    Custom { width: u32, height: u32 },
87}
88
89impl serde::Serialize for ImageSize {
90    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
91    where
92        S: serde::Serializer,
93    {
94        let s = match self {
95            ImageSize::Size1024x1024 => "1024x1024".to_string(),
96            ImageSize::Size768x1344 => "768x1344".to_string(),
97            ImageSize::Size864x1152 => "864x1152".to_string(),
98            ImageSize::Size1344x768 => "1344x768".to_string(),
99            ImageSize::Size1152x864 => "1152x864".to_string(),
100            ImageSize::Size1440x720 => "1440x720".to_string(),
101            ImageSize::Size720x1440 => "720x1440".to_string(),
102            ImageSize::Custom { width, height } => format!("{}x{}", width, height),
103        };
104        serializer.serialize_str(&s)
105    }
106}
107
108impl ImageSize {
109    /// Validate image size constraints
110    ///
111    /// Returns true if the size meets all requirements:
112    /// - Dimensions between 512-2048px
113    /// - Both dimensions divisible by 16
114    /// - Total pixels <= 2^21 (2,097,152)
115    pub fn is_valid(&self) -> bool {
116        match self {
117            ImageSize::Custom { width, height } => {
118                // Check dimension range
119                if *width < 512 || *width > 2048 || *height < 512 || *height > 2048 {
120                    return false;
121                }
122
123                // Check divisibility by 16
124                if width % 16 != 0 || height % 16 != 0 {
125                    return false;
126                }
127
128                // Check total pixels limit (2^21 = 2,097,152)
129                let total_pixels = (*width as u64) * (*height as u64);
130                total_pixels <= 2_097_152
131            },
132            _ => true, // Predefined sizes are already valid
133        }
134    }
135
136    /// Get the dimensions as (width, height)
137    pub fn dimensions(&self) -> (u32, u32) {
138        match self {
139            ImageSize::Size1024x1024 => (1024, 1024),
140            ImageSize::Size768x1344 => (768, 1344),
141            ImageSize::Size864x1152 => (864, 1152),
142            ImageSize::Size1344x768 => (1344, 768),
143            ImageSize::Size1152x864 => (1152, 864),
144            ImageSize::Size1440x720 => (1440, 720),
145            ImageSize::Size720x1440 => (720, 1440),
146            ImageSize::Custom { width, height } => (*width, *height),
147        }
148    }
149}