From c750f1609fb33a22de6e0aae362eaffacc0d5f2e Mon Sep 17 00:00:00 2001 From: tomsun28 Date: Thu, 8 Jan 2026 15:37:57 +0800 Subject: [PATCH 1/3] feat: support async image generation API --- examples/async_image_generator.py | 99 +++++++++++++++++++++++++++ src/zai/api_resource/images/images.py | 74 +++++++++++++++++++- src/zai/types/image/__init__.py | 3 +- src/zai/types/image/image.py | 21 ++++++ 4 files changed, 195 insertions(+), 2 deletions(-) create mode 100644 examples/async_image_generator.py diff --git a/examples/async_image_generator.py b/examples/async_image_generator.py new file mode 100644 index 0000000..5d32a7d --- /dev/null +++ b/examples/async_image_generator.py @@ -0,0 +1,99 @@ +import asyncio +import time + +from zai import ZaiClient + + +class AsyncImageGenerator: + def __init__(self): + self.client = ZaiClient() + + async def generate_image( + self, + prompt: str, + model: str = 'glm-image', + size: str = None, + quality: str = None, + max_wait_time: int = 300, + ): + """ + Asynchronous image generation method + + Args: + prompt: Image generation prompt + model: Model to use for image generation + size: Size of the generated image + quality: Quality level of the generated image + max_wait_time: Maximum wait time (seconds) + + Returns: + Image generation result + """ + try: + # Submit generation task + print('Submitting image generation task...') + response = self.client.images.async_generations( + model=model, + prompt=prompt, + size=size, + quality=quality, + ) + + print(f'Task submitted successfully, Task ID: {response.id}') + print(f'Initial response: {response}') + + # Asynchronously wait for task completion + task_id = response.id + start_time = time.time() + + while True: + # Check for timeout + if time.time() - start_time > max_wait_time: + raise TimeoutError(f'Image generation timeout, exceeded {max_wait_time} seconds') + + # Get task result + print(f'Querying task status... (waited {int(time.time() - start_time)} seconds)') + result = self.client.images.retrieve_images_result(id=task_id) + + print(f'Task status: {result.task_status}') + + # Check if task is completed + if result.task_status == 'SUCCESS': + print('Image generation completed!') + return result + elif result.task_status == 'FAIL': + raise Exception(f'Image generation failed: {result}') + elif result.task_status in ['PROCESSING', 'SUBMITTED']: + print(f'Task in progress, status: {result.task_status}') + else: + print(f'Unknown status: {result.task_status}') + + # Wait 3 seconds before querying again + await asyncio.sleep(30) + + except Exception as e: + print(f'Error occurred during image generation: {e}') + raise + + +# Usage example +async def main(): + generator = AsyncImageGenerator() + try: + result = await generator.generate_image( + prompt='A beautiful sunset over the ocean with colorful clouds', + model='glm-image', + ) + print('\n=== Final Result ===') + print(f'Task Status: {result.task_status}') + print(f'Request ID: {result.request_id}') + if result.image_result: + for i, img in enumerate(result.image_result): + print(f'Image {i + 1} URL: {img.url}') + + except Exception as e: + print(f'Generation failed: {e}') + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/src/zai/api_resource/images/images.py b/src/zai/api_resource/images/images.py index dce31fd..f3c239b 100644 --- a/src/zai/api_resource/images/images.py +++ b/src/zai/api_resource/images/images.py @@ -5,7 +5,7 @@ import httpx from zai.core import NOT_GIVEN, BaseAPI, Body, Headers, NotGiven, make_request_options -from zai.types.image import ImagesResponded +from zai.types.image import AsyncImagesResponded, ImagesResponded from zai.types.sensitive_word_check import SensitiveWordCheckRequest if TYPE_CHECKING: @@ -84,3 +84,75 @@ def generations( cast_type=_cast_type, stream=False, ) + + def async_generations( + self, + *, + prompt: str, + model: str | NotGiven = NOT_GIVEN, + quality: Optional[str] | NotGiven = NOT_GIVEN, + size: Optional[str] | NotGiven = NOT_GIVEN, + request_id: Optional[str] | NotGiven = NOT_GIVEN, + user_id: Optional[str] | NotGiven = NOT_GIVEN, + extra_headers: Headers | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + watermark_enabled: Optional[bool] | NotGiven = NOT_GIVEN, + ) -> AsyncImagesResponded: + """ + Asynchronously generate images from text prompts. + Use retrieve_images_result() to poll for the result. + + Arguments: + prompt (str): Text description of the desired image + model (str): The model to use for image generation + quality (Optional[str]): Quality level of the generated images + size (Optional[str]): Size of the generated images + request_id (Optional[str]): Unique identifier for the request + user_id (Optional[str]): User identifier + extra_headers (Headers): Additional headers to send + extra_body (Body): Additional body parameters + timeout (float | httpx.Timeout): Request timeout + watermark_enabled (Optional[bool]): Whether to enable watermark on generated images + """ + return self._post( + '/async/images/generations', + body={ + 'prompt': prompt, + 'model': model, + 'quality': quality, + 'size': size, + 'user_id': user_id, + 'request_id': request_id, + 'watermark_enabled': watermark_enabled, + }, + options=make_request_options(extra_headers=extra_headers, extra_body=extra_body, timeout=timeout), + cast_type=AsyncImagesResponded, + stream=False, + ) + + def retrieve_images_result( + self, + id: str, + *, + extra_headers: Headers | None = None, + extra_body: Body | None = None, + timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, + ) -> AsyncImagesResponded: + """ + Retrieve the result of an async image generation operation + + Arguments: + id (str): Unique identifier for the image generation task + extra_headers (Headers): Additional headers to send + extra_body (Body): Additional body parameters + timeout (float | httpx.Timeout): Request timeout + """ + if not id: + raise ValueError('`id` must be provided.') + + return self._get( + f'/async-result/{id}', + options=make_request_options(extra_headers=extra_headers, extra_body=extra_body, timeout=timeout), + cast_type=AsyncImagesResponded, + ) diff --git a/src/zai/types/image/__init__.py b/src/zai/types/image/__init__.py index 8046c7c..bf63dae 100644 --- a/src/zai/types/image/__init__.py +++ b/src/zai/types/image/__init__.py @@ -1,6 +1,7 @@ -from .image import GeneratedImage, ImagesResponded +from .image import AsyncImagesResponded, GeneratedImage, ImagesResponded __all__ = [ 'GeneratedImage', 'ImagesResponded', + 'AsyncImagesResponded', ] diff --git a/src/zai/types/image/image.py b/src/zai/types/image/image.py index 36ae3ab..d15b065 100644 --- a/src/zai/types/image/image.py +++ b/src/zai/types/image/image.py @@ -31,3 +31,24 @@ class ImagesResponded(BaseModel): created: int data: List[GeneratedImage] + + +class AsyncImagesResponded(BaseModel): + """ + Async image generation response + + Attributes: + id (Optional[str]): The task order number generated by the Z.ai open platform. + Please use this order number when calling the request result interface. + model (str): Model name + request_id (str): The task number submitted by the user when requesting on the client + or the task number generated by the platform. + task_status (str): Processing status, PROCESSING (processing), SUCCESS (successful), FAIL (failed). + image_result (Optional[List[GeneratedImage]]): Image generation result, available when task_status is SUCCESS + """ + + id: Optional[str] = None + model: str + request_id: str + task_status: str + image_result: Optional[List[GeneratedImage]] = None From d1d5c79ca6db5b1a144d038bb64365560d28bc1c Mon Sep 17 00:00:00 2001 From: tomsun28 Date: Thu, 8 Jan 2026 15:43:19 +0800 Subject: [PATCH 2/3] fix: remove SUBMITTED status from polling check in async image generator The SUBMITTED status should not trigger the "Task in progress" message during polling, as it's only returned immediately after task submission. --- examples/async_image_generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/async_image_generator.py b/examples/async_image_generator.py index 5d32a7d..f230b81 100644 --- a/examples/async_image_generator.py +++ b/examples/async_image_generator.py @@ -63,7 +63,7 @@ async def generate_image( return result elif result.task_status == 'FAIL': raise Exception(f'Image generation failed: {result}') - elif result.task_status in ['PROCESSING', 'SUBMITTED']: + elif result.task_status in ['PROCESSING']: print(f'Task in progress, status: {result.task_status}') else: print(f'Unknown status: {result.task_status}') From 112eef55c3ac765e5b843731c0332000976e0ae9 Mon Sep 17 00:00:00 2001 From: tomsun28 Date: Thu, 8 Jan 2026 16:44:03 +0800 Subject: [PATCH 3/3] update --- src/zai/api_resource/images/images.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/zai/api_resource/images/images.py b/src/zai/api_resource/images/images.py index f3c239b..fa5f957 100644 --- a/src/zai/api_resource/images/images.py +++ b/src/zai/api_resource/images/images.py @@ -100,7 +100,7 @@ def async_generations( watermark_enabled: Optional[bool] | NotGiven = NOT_GIVEN, ) -> AsyncImagesResponded: """ - Asynchronously generate images from text prompts. + Asynchronously generate images from text prompts. Only support glm-image model Use retrieve_images_result() to poll for the result. Arguments: