| 1 | """ |
| 2 | Seedream 图像生成 API 客户端 |
| 3 | 字节跳动 ARK - doubao-seedream-5-0-260128 模型 |
| 4 | """ |
| 5 | |
| 6 | import os |
| 7 | import time |
| 8 | import logging |
| 9 | from typing import Optional, List, Dict |
| 10 | import httpx |
| 11 | from openai import OpenAI |
| 12 | |
| 13 | # 模型名称映射表（旧名称 -> 新名称） |
| 14 | MODEL_NAME_MAP: Dict[str, str] = { |
| 15 | # doubao-seedream-5-0 系列 |
| 16 | "doubao-seedream-5-0": "doubao-seedream-5-0-260128", |
| 17 | # doubao-seedream-4-5 系列 |
| 18 | "doubao-seedream-4-5": "doubao-seedream-4-5-251128", |
| 19 | # doubao-seedream-4-0 系列 |
| 20 | "doubao-seedream-4-0": "doubao-seedream-4-0-250828", |
| 21 | } |
| 22 | |
| 23 | |
| 24 | def normalize_model_name(model: str) -> str: |
| 25 | """ |
| 26 | 规范化模型名称 |
| 27 | |
| 28 | Args: |
| 29 | model: 传入的模型名称 |
| 30 | |
| 31 | Returns: |
| 32 | 规范化后的模型名称 |
| 33 | """ |
| 34 | return MODEL_NAME_MAP.get(model, model) |
| 35 | |
| 36 | |
| 37 | class SeedreamClient: |
| 38 | """ |
| 39 | Seedream 图像生成客户端（字节跳动 ARK） |
| 40 | 支持文生图功能 |
| 41 | """ |
| 42 | |
| 43 | def __init__( |
| 44 | self, |
| 45 | api_key: Optional[str] = None, |
| 46 | base_url: Optional[str] = None, |
| 47 | local_proxy: Optional[str] = None, |
| 48 | timeout: int = 120, |
| 49 | ) -> None: |
| 50 | """ |
| 51 | 初始化 Seedream 客户端 |
| 52 | |
| 53 | Args: |
| 54 | api_key: ARK API Key |
| 55 | base_url: ARK API 基础 URL |
| 56 | timeout: HTTP请求超时时间（秒） |
| 57 | """ |
| 58 | self.api_key = api_key or os.getenv("ARK_API_KEY") |
| 59 | self.base_url = base_url or "https://ark.cn-beijing.volces.com/api/v3" |
| 60 | self.local_proxy = local_proxy |
| 61 | self.timeout = timeout |
| 62 | |
| 63 | if not self.api_key: |
| 64 | logging.warning( |
| 65 | "SeedreamClient missing api_key. Set ARK_API_KEY." |
| 66 | ) |
| 67 | |
| 68 | client_kwargs = { |
| 69 | "base_url": self.base_url, |
| 70 | "api_key": self.api_key, |
| 71 | "timeout": timeout, |
| 72 | } |
| 73 | if self.local_proxy: |
| 74 | client_kwargs["http_client"] = httpx.Client(proxy=self.local_proxy, timeout=timeout) |
| 75 | |
| 76 | self.client = OpenAI(**client_kwargs) |
| 77 | |
| 78 | def generate_image( |
| 79 | self, |
| 80 | prompt: str, |
| 81 | session_id: str, |
| 82 | model: str = "doubao-seedream-4-5-251128", |
| 83 | size: str = "1920*1080", |
| 84 | image_paths: Optional[List[str]] = None, |
| 85 | **kwargs |
| 86 | ) -> List[str]: |
| 87 | """ |
| 88 | 生成图片 |
| 89 | |
| 90 | Args: |
| 91 | prompt: 提示词 |
| 92 | session_id: 任务或会话ID，用于构建存储路径 |
| 93 | model: 模型名称 |
| 94 | size: 生成图片的分辨率，如 "1920*1080", "1024*1024" |
| 95 | image_paths: 参考图路径或URL列表 (图生图) |
| 96 | **kwargs: 其他生成参数 |
| 97 | |
| 98 | Returns: |
| 99 | 生成的图片路径列表 |
| 100 | """ |
| 101 | if not self.api_key: |
| 102 | raise RuntimeError("ARK_API_KEY not set.") |
| 103 | |
| 104 | # 规范化模型名称（旧名称 -> 新名称） |
| 105 | model = normalize_model_name(model) |
| 106 | |
| 107 | # 处理分辨率 (Seedream 要求至少 3686400 像素) |
| 108 | # 常用 2K/4K 分辨率 |
| 109 | size_map = { |
| 110 | # 16:9 |
| 111 | "1920*1080": (1920, 1080), |
| 112 | "2048*1080": (2048, 1080), # 2K 电影 |
| 113 | "2560*1440": (2560, 1440), # 2K QHD |
| 114 | "3840*2160": (3840, 2160), # 4K UHD |
| 115 | "4096*2160": (4096, 2160), # 4K 电影 |
| 116 | # 9:16 |
| 117 | "1080*1920": (1080, 1920), |
| 118 | "1080*2048": (1080, 2048), |
| 119 | "1440*2560": (1440, 2560), |
| 120 | "2160*3840": (2160, 3840), |
| 121 | "2160*4096": (2160, 4096), |
| 122 | # 1:1 |
| 123 | "1024*1024": (1024, 1024), |
| 124 | "2048*2048": (2048, 2048), # 2K 正方 |
| 125 | # 4:3 |
| 126 | "1920*1440": (1920, 1440), |
| 127 | "2560*1920": (2560, 1920), |
| 128 | # 3:4 |
| 129 | "1440*1920": (1440, 1920), |
| 130 | "1920*2560": (1920, 2560), |
| 131 | } |
| 132 | |
| 133 | width, height = 1920, 1080 # 默认 |
| 134 | min_pixels = 3686400 |
| 135 | |
| 136 | if size: |
| 137 | parts = size.split("*") |
| 138 | if len(parts) == 2: |
| 139 | w, h = int(parts[0]), int(parts[1]) |
| 140 | width, height = w, h |
| 141 | |
| 142 | # 确保满足最小像素要求 |
| 143 | if width * height < min_pixels: |
| 144 | # 查找相同宽高比的常用分辨率 |
| 145 | aspect_ratio = width / height |
| 146 | for (w, h) in size_map.values(): |
| 147 | if abs(w / h - aspect_ratio) < 0.01 and w * h >= min_pixels: |
| 148 | width, height = w, h |
| 149 | break |
| 150 | else: |
| 151 | # 没有找到合适的，按比例放大 |
| 152 | scale = (min_pixels / (width * height)) ** 0.5 |
| 153 | width = int(width * scale) |
| 154 | height = int(height * scale) |
| 155 | width = width if width % 2 == 0 else width + 1 |
| 156 | height = height if height % 2 == 0 else height + 1 |
| 157 | |
| 158 | # 构建 extra_body |
| 159 | extra_body = { |
| 160 | "watermark": False, |
| 161 | "sequential_image_generation": "disabled", |
| 162 | } |
| 163 | |
| 164 | # 添加其他参数 |
| 165 | if "seed" in kwargs: |
| 166 | extra_body["seed"] = kwargs["seed"] |
| 167 | if "quality" in kwargs: |
| 168 | extra_body["quality"] = kwargs["quality"] |
| 169 | if "style" in kwargs: |
| 170 | extra_body["style"] = kwargs["style"] |
| 171 | |
| 172 | # 处理参考图 (图生图) |
| 173 | image_urls = [] |
| 174 | if image_paths and len(image_paths) > 0: |
| 175 | # 处理参考图：支持 URL 和本地文件 |
| 176 | ref_images = [] |
| 177 | for p in image_paths: |
| 178 | if p.startswith("http"): |
| 179 | ref_images.append(p) |
| 180 | elif os.path.exists(p): |
| 181 | # 转换为 base64 URL |
| 182 | import base64 |
| 183 | with open(p, "rb") as f: |
| 184 | img_data = base64.b64encode(f.read()).decode("utf-8") |
| 185 | ext = os.path.splitext(p)[1].lower() |
| 186 | mime = "image/png" if ext == ".png" else "image/jpeg" |
| 187 | ref_images.append(f"data:{mime};base64,{img_data}") |
| 188 | extra_body["image"] = ref_images |
| 189 | |
| 190 | # 调用 API |
| 191 | if image_paths and len(image_paths) > 0: |
| 192 | # 图生图 - image 放在 extra_body 中 |
| 193 | response = self.client.images.generate( |
| 194 | model=model, |
| 195 | prompt=prompt, |
| 196 | size=f"{width}x{height}", |
| 197 | response_format="url", |
| 198 | extra_body=extra_body, |
| 199 | ) |
| 200 | else: |
| 201 | # 文生图 |
| 202 | response = self.client.images.generate( |
| 203 | model=model, |
| 204 | prompt=prompt, |
| 205 | size=f"{width}x{height}", |
| 206 | response_format="url", |
| 207 | extra_body=extra_body, |
| 208 | ) |
| 209 | |
| 210 | # 下载图片到本地 |
| 211 | generated_paths = [] |
| 212 | if response.data: |
| 213 | for idx, img_data in enumerate(response.data): |
| 214 | if img_data.url: |
| 215 | local_path = self._download_image( |
| 216 | img_data.url, session_id, idx |
| 217 | ) |
| 218 | if local_path: |
| 219 | generated_paths.append(local_path) |
| 220 | |
| 221 | return generated_paths |
| 222 | |
| 223 | def _download_image(self, url: str, session_id: str, idx: int) -> Optional[str]: |
| 224 | """从URL下载图片到本地""" |
| 225 | import requests |
| 226 | |
| 227 | # 构建存储路径 |
| 228 | base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| 229 | result_dir = os.path.join(base_dir, "code", "result", "image", str(session_id)) |
| 230 | os.makedirs(result_dir, exist_ok=True) |
| 231 | |
| 232 | file_name = f"seedream_{int(time.time())}_{idx}.png" |
| 233 | file_path = os.path.join(result_dir, file_name) |
| 234 | |
| 235 | try: |
| 236 | proxies = {"http": self.local_proxy, "https": self.local_proxy} if self.local_proxy else None |
| 237 | response = requests.get(url, timeout=self.timeout, proxies=proxies) |
| 238 | response.raise_for_status() |
| 239 | with open(file_path, "wb") as f: |
| 240 | f.write(response.content) |
| 241 | return file_path |
| 242 | except Exception as e: |
| 243 | logging.error(f"Failed to download image from {url}: {e}") |
| 244 | return None |
| 245 | |
| 246 | |
| 247 | if __name__ == "__main__": |
| 248 | import sys |
| 249 | sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| 250 | from config import Config # 加载 .env |
| 251 | |
| 252 | print("=== Seedream 可用性测试 ===") |
| 253 | api_key = os.getenv("ARK_API_KEY", "") |
| 254 | base_url = os.getenv("ARK_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3") |
| 255 | if not api_key: |
| 256 | print("✗ ARK_API_KEY 未设置，跳过") |
| 257 | sys.exit(1) |
| 258 | print(f" API Key: {api_key[:6]}***{api_key[-4:]}") |
| 259 | print(f" Base URL: {base_url}") |
| 260 | |
| 261 | client = SeedreamClient(api_key=api_key, base_url=base_url) |
| 262 | |
| 263 | # === 测试1: 文生图 === |
| 264 | prompt = "星际穿越，黑洞，黑洞里冲出一辆支离破碎的复古列车，视觉冲击力，电影大片，末日既视感" |
| 265 | print(f"\n[测试1: 文生图] Prompt: {prompt}") |
| 266 | t0 = time.time() |
| 267 | try: |
| 268 | paths = client.generate_image( |
| 269 | prompt=prompt, |
| 270 | session_id="test_avail", |
| 271 | model="doubao-seedream-5-0-260128", |
| 272 | size="1920*1080", |
| 273 | ) |
| 274 | elapsed = time.time() - t0 |
| 275 | if paths: |
| 276 | print(f"✓ 生成 {len(paths)} 张图片 ({elapsed:.1f}s): {paths}") |
| 277 | else: |
| 278 | print(f"✗ 返回空列表 ({elapsed:.1f}s)") |
| 279 | except Exception as e: |
| 280 | print(f"✗ 图片生成失败: {e}") |
| 281 | |
| 282 | # === 测试2: 图生图 === |
| 283 | # 需要一张已有的参考图路径 |
| 284 | ref_image_path = "code/result/image/test_avail/test_input.png" |
| 285 | if os.path.exists(ref_image_path): |
| 286 | prompt_i2i = "将这只猫变成赛博朋克风格" |
| 287 | print(f"\n[测试2: 图生图] Prompt: {prompt_i2i}") |
| 288 | print(f" 参考图: {ref_image_path}") |
| 289 | t0 = time.time() |
| 290 | try: |
| 291 | paths = client.generate_image( |
| 292 | prompt=prompt_i2i, |
| 293 | session_id="test_avail", |
| 294 | model="doubao-seedream-5-0-260128", |
| 295 | size="1920*1080", |
| 296 | image_paths=[ref_image_path], |
| 297 | ) |
| 298 | elapsed = time.time() - t0 |
| 299 | if paths: |
| 300 | print(f"✓ 生成 {len(paths)} 张图片 ({elapsed:.1f}s): {paths}") |
| 301 | else: |
| 302 | print(f"✗ 返回空列表 ({elapsed:.1f}s)") |
| 303 | except Exception as e: |
| 304 | print(f"✗ 图生图失败: {e}") |
| 305 | else: |
| 306 | print(f"\n[测试2: 图生图] ✗ 参考图不存在: {ref_image_path}") |
| 307 | print(" 跳过图生图测试，请先运行文生图测试生成参考图") |
| 308 |