feat: ark image generation sdk

Author: fabianlimxuhuai

volcenginesdkarkruntime/_client.py

@@ -42,6 +42,7 @@ class Ark(SyncAPIClient):
     context: resources.Context
     multimodal_embeddings: resources.MultimodalEmbeddings
     content_generation: resources.ContentGeneration
+    images: resources.Images
     batch_chat: resources.BatchChat
     model_breaker_map: dict[str, ModelBreaker]
     model_breaker_lock: threading.Lock
@@ -71,7 +72,6 @@ def __init__(
             Returns:
                 ark client
         """
-
         if ak is None:
             ak = os.environ.get("VOLC_ACCESSKEY")
         if sk is None:
@@ -107,6 +107,7 @@ def __init__(
         self.context = resources.Context(self)
         self.multimodal_embeddings = resources.MultimodalEmbeddings(self)
         self.content_generation = resources.ContentGeneration(self)
+        self.images = resources.Images(self)
         self.batch_chat = resources.BatchChat(self)
         self.model_breaker_map = defaultdict(ModelBreaker)
         self.model_breaker_lock = threading.Lock()
@@ -161,6 +162,7 @@ class AsyncArk(AsyncAPIClient):
     context: resources.AsyncContext
     multimodal_embeddings: resources.AsyncMultimodalEmbeddings
     content_generation: resources.AsyncContentGeneration
+    images: resources.AsyncImages
     batch_chat: resources.AsyncBatchChat
     model_breaker_map: dict[str, ModelBreaker]
     model_breaker_lock: asyncio.Lock
@@ -225,6 +227,7 @@ def __init__(
         self.context = resources.AsyncContext(self)
         self.multimodal_embeddings = resources.AsyncMultimodalEmbeddings(self)
         self.content_generation = resources.AsyncContentGeneration(self)
+        self.images = resources.AsyncImages(self)
         self.batch_chat = resources.AsyncBatchChat(self)
         self.model_breaker_map = defaultdict(ModelBreaker)
         self.model_breaker_lock = asyncio.Lock()

volcenginesdkarkruntime/resources/__init__.py

@@ -6,6 +6,7 @@
 from .context import Context, AsyncContext
 from .multimodal_embeddings import MultimodalEmbeddings, AsyncMultimodalEmbeddings
 from .content_generation import ContentGeneration, AsyncContentGeneration
+from .images import Images, AsyncImages
 from .batch_chat import BatchChat, AsyncBatchChat
 
 __all__ = [
@@ -24,6 +25,8 @@
     "AsyncContext",
     "ContentGeneration",
     "AsyncContentGeneration",
+    "Images",
+    "AsyncImages",
     "BatchChat",
     "AsyncBatchChat",
 ]

volcenginesdkarkruntime/resources/images/__init__.py

@@ -0,0 +1,6 @@
+from volcenginesdkarkruntime.resources.images.images import (
+    Images,
+    AsyncImages,
+)
+
+__all__ = ["Images", "AsyncImages"]

volcenginesdkarkruntime/resources/images/images.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+import httpx
+
+from ..._base_client import make_request_options
+from ..._utils._utils import with_sts_token, async_with_sts_token
+from ..._resource import SyncAPIResource, AsyncAPIResource
+from ...types.images import ImagesResponse
+from ..._types import Body, Query, Headers
+
+
+class Images(SyncAPIResource):
+    @with_sts_token
+    def generate(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        response_format: str | None = None,
+        size: str | None = None,
+        seed: int | None = None,
+        guidance_scale: float | None = None,
+        watermark: bool | None = None,
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> ImagesResponse:
+        resp = self._post(
+            "/images/generations",
+            body={
+                "model": model,
+                "prompt": prompt,
+                "response_format": response_format,
+                "size": size,
+                "seed": seed,
+                "guidance_scale": guidance_scale,
+                "watermark": watermark,
+            },
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ImagesResponse,
+        )
+
+        return resp
+
+
+class AsyncImages(AsyncAPIResource):
+    @async_with_sts_token
+    async def generate(
+            self,
+            *,
+            model: str,
+            prompt: str,
+            response_format: str | None = None,
+            size: str | None = None,
+            seed: int | None = None,
+            guidance_scale: float | None = None,
+            watermark: bool | None = None,
+            extra_headers: Headers | None = None,
+            extra_query: Query | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None = None,
+    ) -> ImagesResponse:
+        return await self._post(
+            "/images/generations",
+            body={
+                "model": model,
+                "prompt": prompt,
+                "response_format": response_format,
+                "size": size,
+                "seed": seed,
+                "guidance_scale": guidance_scale,
+                "watermark": watermark,
+            },
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=ImagesResponse,
+        )

volcenginesdkarkruntime/types/images/__init__.py

@@ -0,0 +1 @@
+from .images import ImagesResponse

volcenginesdkarkruntime/types/images/images.py

@@ -0,0 +1,43 @@
+from typing import List
+
+from volcenginesdkarkruntime._models import BaseModel
+
+__all__ = ["ImagesResponse"]
+
+
+class Usage(BaseModel):
+    generated_images: int
+    """The number of images generated."""
+
+
+class Image(BaseModel):
+    url: str
+    """The URL of the generated image, if any."""
+
+    b64_json: str
+    """The Base 64 encoded string of the generated image, if any."""
+
+
+class Error(BaseModel):
+    message: str
+    """The reason for failed image generation"""
+
+    code: str
+    """The error code for failed image generation"""
+
+
+class ImagesResponse(BaseModel):
+    model: str
+    """The model used to generated the images."""
+
+    data: List[Image]
+    """The generated images."""
+
+    error: Error
+    """The error body, if applicable."""
+
+    usage: Usage
+    """The usage information for the generation of images."""
+
+    created_at: int
+    """The Unix timestamp when the image was generated."""

volcenginesdkexamples/volcenginesdkarkruntime/async_image_generations.py

@@ -0,0 +1,28 @@
+import asyncio
+
+from volcenginesdkarkruntime import AsyncArk
+
+# Authentication
+# 1.If you authorize your endpoint using an API key, you can set your api key to environment variable "ARK_API_KEY"
+# or specify api key by Ark(api_key="${YOUR_API_KEY}").
+# Note: If you use an API key, this API key will not be refreshed.
+# To prevent the API from expiring and failing after some time, choose an API key with no expiration date.
+client = AsyncArk()
+
+
+async def main():
+    print("----- async generate images -----")
+
+    result = client.images.generate(
+        model="${YOUR_ENDPOINT_ID}",
+        prompt="龙与地下城女骑士背景是起伏的平原，目光从镜头转向平原",
+        seed=1234567890,
+        watermark=True,
+        size="512x512",
+        guidance_scale=2.5,
+    )
+
+    print(await result)
+
+if __name__ == "__main__":
+    asyncio.run(main())

volcenginesdkexamples/volcenginesdkarkruntime/image_generations.py

@@ -0,0 +1,22 @@
+from volcenginesdkarkruntime import Ark
+
+# Authentication
+# 1.If you authorize your endpoint using an API key, you can set your api key to environment variable "ARK_API_KEY"
+# or specify api key by Ark(api_key="${YOUR_API_KEY}").
+# Note: If you use an API key, this API key will not be refreshed.
+# To prevent the API from expiring and failing after some time, choose an API key with no expiration date.
+client = Ark()
+
+if __name__ == "__main__":
+    print("----- generate images -----")
+
+    result = client.images.generate(
+        model="${YOUR_ENDPOINT_ID}",
+        prompt="龙与地下城女骑士背景是起伏的平原，目光从镜头转向平原",
+        seed=1234567890,
+        watermark=True,
+        size="512x512",
+        guidance_scale=2.5,
+    )
+
+    print(result)