| 1 | # https://yunwu.apifox.cn/api-347960869 |
| 2 | |
| 3 | import asyncio |
| 4 | import logging |
| 5 | import aiohttp |
| 6 | from typing import List, Optional |
| 7 | from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential |
| 8 | from utils.retry import after_func |
| 9 | from utils.image import image_path_to_b64 |
| 10 | from interfaces.image_output import ImageOutput |
| 11 | |
| 12 | |
| 13 | class ImageGeneratorDoubaoSeedreamYunwuAPI: |
| 14 | def __init__( |
| 15 | self, |
| 16 | api_key: str, |
| 17 | model: str = "doubao-seedream-4-0-250828", |
| 18 | |
| 19 | ): |
| 20 | self.api_key = api_key |
| 21 | self.base_url = "https://yunwu.ai/v1/images/generations" |
| 22 | self.model = model |
| 23 | |
| 24 | |
| 25 | @retry( |
| 26 | stop=stop_after_attempt(3), |
| 27 | wait=wait_exponential(multiplier=1, max=30), |
| 28 | retry=retry_if_exception_type((aiohttp.ClientError, asyncio.TimeoutError)), |
| 29 | reraise=True, |
| 30 | after=after_func, |
| 31 | ) |
| 32 | async def generate_single_image( |
| 33 | self, |
| 34 | prompt: str, |
| 35 | reference_image_paths: List[str] = [], |
| 36 | size: Optional[str] = None, |
| 37 | **kwargs, |
| 38 | ) -> ImageOutput: |
| 39 | """ |
| 40 | size: [1024x1024, 4096x4096] |
| 41 | """ |
| 42 | |
| 43 | logging.info(f"Calling {self.model} to generate image...") |
| 44 | |
| 45 | image = [ |
| 46 | image_path_to_b64(path, mime=True) for path in reference_image_paths |
| 47 | ] |
| 48 | |
| 49 | payload = { |
| 50 | "model": self.model, |
| 51 | "prompt": prompt, |
| 52 | "sequential_image_generation": "disabled", # "auto" or "disabled" |
| 53 | # "sequential_image_generation_options": { |
| 54 | # "max_images": 1 |
| 55 | # }, |
| 56 | "response_format": "url", |
| 57 | "size": size if size is not None else "1024x1024", |
| 58 | } |
| 59 | if len(image) > 0: |
| 60 | payload["image"] = image |
| 61 | |
| 62 | headers = { |
| 63 | "Authorization": f"Bearer {self.api_key}", |
| 64 | "Content-Type": "application/json", |
| 65 | } |
| 66 | |
| 67 | async with aiohttp.ClientSession() as session: |
| 68 | async with session.post(self.base_url, json=payload, headers=headers) as response: |
| 69 | response_json = await response.json() |
| 70 | if response.status >= 400: |
| 71 | raise RuntimeError(f"Image generation failed with HTTP {response.status}: {response_json}") |
| 72 | |
| 73 | data = response_json['data'][0]['url'] |
| 74 | return ImageOutput(fmt="url", ext="png", data=data) |
| 75 |