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(Clone, Serialize, Validate)]
8#[validate(schema(function = "validate_image_body"))]
9pub struct ImageGenBody<N>
10where
11    N: ImageGen,
12{
13    /// The model to use for image generation
14    pub model: N,
15    /// Text description of the desired image
16    #[serde(skip_serializing_if = "Option::is_none")]
17    #[validate(length(min = 1))]
18    pub prompt: Option<String>,
19    /// Image-generation quality.
20    ///
21    /// `Hd` favors detail and consistency; `Standard` favors lower latency.
22    /// The `glm-image` model supports only [`ImageQuality::Hd`].
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub quality: Option<ImageQuality>,
25    /// Image size
26    /// Use an [`ImageSize`] preset or custom dimensions satisfying the limits
27    /// documented on that type.
28    #[serde(skip_serializing_if = "Option::is_none")]
29    pub size: Option<ImageSize>,
30    /// Whether to add watermark to AI generated images
31    #[serde(skip_serializing_if = "Option::is_none")]
32    pub watermark_enabled: Option<bool>,
33    /// Unique ID of the end user to help platform intervene against violations,
34    /// illegal content generation, or other abusive behaviors
35    #[serde(skip_serializing_if = "Option::is_none")]
36    #[validate(length(min = 6, max = 128))]
37    pub user_id: Option<String>,
38}
39
40impl<N> std::fmt::Debug for ImageGenBody<N>
41where
42    N: ImageGen + std::fmt::Debug,
43{
44    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        formatter
46            .debug_struct("ImageGenBody")
47            .field("model", &self.model)
48            .field("prompt_configured", &self.prompt.is_some())
49            .field("quality", &self.quality)
50            .field("size", &self.size)
51            .field("watermark_enabled", &self.watermark_enabled)
52            .field("user_id", &self.user_id.as_ref().map(|_| "[REDACTED]"))
53            .finish()
54    }
55}
56
57fn validate_image_body<N>(body: &ImageGenBody<N>) -> Result<(), validator::ValidationError>
58where
59    N: ImageGen,
60{
61    if body
62        .prompt
63        .as_deref()
64        .is_none_or(|prompt| prompt.trim().is_empty())
65    {
66        return Err(validator::ValidationError::new("prompt_required"));
67    }
68    let model = N::NAME;
69
70    if model == "glm-image" && matches!(body.quality, Some(ImageQuality::Standard)) {
71        return Err(validator::ValidationError::new(
72            "quality_not_supported_by_model",
73        ));
74    }
75    if body
76        .size
77        .as_ref()
78        .is_some_and(|size| !size.is_valid_for_model(model))
79    {
80        return Err(validator::ValidationError::new("invalid_image_size"));
81    }
82    Ok(())
83}
84
85impl<N> ImageGenBody<N>
86where
87    N: ImageGen,
88{
89    /// Create an image-generation body with all optional fields omitted.
90    pub fn new(model: N) -> Self {
91        Self {
92            model,
93            prompt: None,
94            quality: None,
95            size: None,
96            watermark_enabled: None,
97            user_id: None,
98        }
99    }
100}
101
102/// Image generation quality options
103#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
104pub enum ImageQuality {
105    /// High quality - more refined and detailed images
106    #[serde(rename = "hd")]
107    Hd,
108    /// Standard quality - faster generation
109    #[serde(rename = "standard")]
110    Standard,
111}
112
113/// Image size options
114///
115/// `glm-image` accepts dimensions from 1,024 through 2,048 pixels, divisible
116/// by 32, with at most 2²² total pixels. Other models accept dimensions from
117/// 512 through 2,048 pixels, divisible by 16, with at most 2²¹ total pixels.
118#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
119pub enum ImageSize {
120    /// 1280x1280 pixels (`glm-image` default)
121    Size1280x1280,
122    /// 1568x1056 pixels
123    Size1568x1056,
124    /// 1056x1568 pixels
125    Size1056x1568,
126    /// 1472x1088 pixels
127    Size1472x1088,
128    /// 1088x1472 pixels
129    Size1088x1472,
130    /// 1728x960 pixels
131    Size1728x960,
132    /// 960x1728 pixels
133    Size960x1728,
134    /// 1024x1024 pixels
135    Size1024x1024,
136    /// 768x1344 pixels
137    Size768x1344,
138    /// 864x1152 pixels
139    Size864x1152,
140    /// 1344x768 pixels
141    Size1344x768,
142    /// 1152x864 pixels
143    Size1152x864,
144    /// 1440x720 pixels
145    Size1440x720,
146    /// 720x1440 pixels
147    Size720x1440,
148    /// Custom dimensions in width x height format
149    Custom {
150        /// Custom width in pixels; limits depend on the selected model.
151        width: u32,
152        /// Custom height in pixels; limits depend on the selected model.
153        height: u32,
154    },
155}
156
157impl serde::Serialize for ImageSize {
158    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
159    where
160        S: serde::Serializer,
161    {
162        let dimensions = match self {
163            ImageSize::Size1280x1280 => "1280x1280",
164            ImageSize::Size1568x1056 => "1568x1056",
165            ImageSize::Size1056x1568 => "1056x1568",
166            ImageSize::Size1472x1088 => "1472x1088",
167            ImageSize::Size1088x1472 => "1088x1472",
168            ImageSize::Size1728x960 => "1728x960",
169            ImageSize::Size960x1728 => "960x1728",
170            ImageSize::Size1024x1024 => "1024x1024",
171            ImageSize::Size768x1344 => "768x1344",
172            ImageSize::Size864x1152 => "864x1152",
173            ImageSize::Size1344x768 => "1344x768",
174            ImageSize::Size1152x864 => "1152x864",
175            ImageSize::Size1440x720 => "1440x720",
176            ImageSize::Size720x1440 => "720x1440",
177            ImageSize::Custom { width, height } => {
178                return serializer.collect_str(&format_args!("{width}x{height}"));
179            },
180        };
181        serializer.serialize_str(dimensions)
182    }
183}
184
185impl ImageSize {
186    /// Validate the size against the non-`glm-image` constraints.
187    ///
188    /// This method retains the original SDK behavior. Request-body validation
189    /// automatically applies the correct rules for the selected model.
190    pub fn is_valid(&self) -> bool {
191        let (width, height) = self.dimensions();
192        dimensions_are_valid(width, height, 512, 16, 1 << 21)
193    }
194
195    fn is_valid_for_model(&self, model: &str) -> bool {
196        if model == "glm-image"
197            && matches!(
198                self,
199                Self::Size1280x1280
200                    | Self::Size1568x1056
201                    | Self::Size1056x1568
202                    | Self::Size1472x1088
203                    | Self::Size1088x1472
204                    | Self::Size1728x960
205                    | Self::Size960x1728
206            )
207        {
208            return true;
209        }
210        let (width, height) = self.dimensions();
211        if model == "glm-image" {
212            dimensions_are_valid(width, height, 1024, 32, 1 << 22)
213        } else {
214            dimensions_are_valid(width, height, 512, 16, 1 << 21)
215        }
216    }
217
218    /// Get the dimensions as (width, height)
219    pub fn dimensions(&self) -> (u32, u32) {
220        match self {
221            ImageSize::Size1280x1280 => (1280, 1280),
222            ImageSize::Size1568x1056 => (1568, 1056),
223            ImageSize::Size1056x1568 => (1056, 1568),
224            ImageSize::Size1472x1088 => (1472, 1088),
225            ImageSize::Size1088x1472 => (1088, 1472),
226            ImageSize::Size1728x960 => (1728, 960),
227            ImageSize::Size960x1728 => (960, 1728),
228            ImageSize::Size1024x1024 => (1024, 1024),
229            ImageSize::Size768x1344 => (768, 1344),
230            ImageSize::Size864x1152 => (864, 1152),
231            ImageSize::Size1344x768 => (1344, 768),
232            ImageSize::Size1152x864 => (1152, 864),
233            ImageSize::Size1440x720 => (1440, 720),
234            ImageSize::Size720x1440 => (720, 1440),
235            ImageSize::Custom { width, height } => (*width, *height),
236        }
237    }
238}
239
240fn dimensions_are_valid(
241    width: u32,
242    height: u32,
243    minimum: u32,
244    divisor: u32,
245    maximum_pixels: u64,
246) -> bool {
247    (minimum..=2048).contains(&width)
248        && (minimum..=2048).contains(&height)
249        && width.is_multiple_of(divisor)
250        && height.is_multiple_of(divisor)
251        && u64::from(width) * u64::from(height) <= maximum_pixels
252}
253
254#[cfg(test)]
255mod tests {
256    use super::*;
257    use crate::model::gen_image::{CogView4, GlmImage};
258
259    #[test]
260    fn body_validation_requires_prompt_and_valid_custom_size() {
261        assert!(ImageGenBody::new(CogView4 {}).validate().is_err());
262        let mut body = ImageGenBody::new(CogView4 {});
263        body.prompt = Some("image".into());
264        assert!(body.validate().is_ok());
265        body.size = Some(ImageSize::Custom {
266            width: 513,
267            height: 512,
268        });
269        assert!(body.validate().is_err());
270    }
271
272    #[test]
273    fn body_validation_applies_model_specific_size_and_quality_rules() {
274        let mut glm = ImageGenBody::new(GlmImage {});
275        glm.prompt = Some("image".into());
276        glm.size = Some(ImageSize::Size960x1728);
277        glm.quality = Some(ImageQuality::Hd);
278        assert!(glm.validate().is_ok());
279
280        glm.quality = Some(ImageQuality::Standard);
281        assert!(glm.validate().is_err());
282
283        let mut cogview = ImageGenBody::new(CogView4 {});
284        cogview.prompt = Some("image".into());
285        cogview.size = Some(ImageSize::Size1568x1056);
286        assert!(cogview.validate().is_ok());
287        cogview.size = Some(ImageSize::Size1280x1280);
288        assert!(cogview.validate().is_ok());
289    }
290
291    #[test]
292    fn debug_redacts_prompt_and_user_identifier() {
293        let mut body = ImageGenBody::new(GlmImage {});
294        body.prompt = Some("private image prompt".to_owned());
295        body.user_id = Some("private-user".to_owned());
296        let debug = format!("{body:?}");
297        assert!(!debug.contains("private image prompt"));
298        assert!(!debug.contains("private-user"));
299        assert!(debug.contains("prompt_configured: true"));
300    }
301}
zai_rs/model/gen_image/image_request.rs

zai_rs/model/gen_image/
image_request.rs
