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