diff --git a/examples/async_image_generator.py b/examples/async_image_generator.py new file mode 100644 index 0000000..f230b81 --- /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']: + 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..fa5f957 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. Only support glm-image model + 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