""" Seedream 图像生成 API 客户端 字节跳动 ARK - doubao-seedream-5-0-260128 模型 """ import os import time import logging from typing import Optional, List, Dict import httpx from openai import OpenAI # 模型名称映射表（旧名称 -> 新名称） MODEL_NAME_MAP: Dict[str, str] = { # doubao-seedream-5-0 系列 "doubao-seedream-5-0": "doubao-seedream-5-0-260128", # doubao-seedream-4-5 系列 "doubao-seedream-4-5": "doubao-seedream-4-5-251128", # doubao-seedream-4-0 系列 "doubao-seedream-4-0": "doubao-seedream-4-0-250828", } def normalize_model_name(model: str) -> str: """ 规范化模型名称 Args: model: 传入的模型名称 Returns: 规范化后的模型名称 """ return MODEL_NAME_MAP.get(model, model) class SeedreamClient: """ Seedream 图像生成客户端（字节跳动 ARK） 支持文生图功能 """ def __init__( self, api_key: Optional[str] = None, base_url: Optional[str] = None, local_proxy: Optional[str] = None, timeout: int = 120, ) -> None: """ 初始化 Seedream 客户端 Args: api_key: ARK API Key base_url: ARK API 基础 URL timeout: HTTP请求超时时间（秒） """ self.api_key = api_key or os.getenv("ARK_API_KEY") self.base_url = base_url or "https://ark.cn-beijing.volces.com/api/v3" self.local_proxy = local_proxy self.timeout = timeout if not self.api_key: logging.warning( "SeedreamClient missing api_key. Set ARK_API_KEY." ) client_kwargs = { "base_url": self.base_url, "api_key": self.api_key, "timeout": timeout, } if self.local_proxy: client_kwargs["http_client"] = httpx.Client(proxy=self.local_proxy, timeout=timeout) self.client = OpenAI(**client_kwargs) def generate_image( self, prompt: str, session_id: str, model: str = "doubao-seedream-4-5-251128", size: str = "1920*1080", image_paths: Optional[List[str]] = None, **kwargs ) -> List[str]: """ 生成图片 Args: prompt: 提示词 session_id: 任务或会话ID，用于构建存储路径 model: 模型名称 size: 生成图片的分辨率，如 "1920*1080", "1024*1024" image_paths: 参考图路径或URL列表 (图生图) **kwargs: 其他生成参数 Returns: 生成的图片路径列表 """ if not self.api_key: raise RuntimeError("ARK_API_KEY not set.") # 规范化模型名称（旧名称 -> 新名称） model = normalize_model_name(model) # 处理分辨率 (Seedream 要求至少 3686400 像素) # 常用 2K/4K 分辨率 size_map = { # 16:9 "1920*1080": (1920, 1080), "2048*1080": (2048, 1080), # 2K 电影 "2560*1440": (2560, 1440), # 2K QHD "3840*2160": (3840, 2160), # 4K UHD "4096*2160": (4096, 2160), # 4K 电影 # 9:16 "1080*1920": (1080, 1920), "1080*2048": (1080, 2048), "1440*2560": (1440, 2560), "2160*3840": (2160, 3840), "2160*4096": (2160, 4096), # 1:1 "1024*1024": (1024, 1024), "2048*2048": (2048, 2048), # 2K 正方 # 4:3 "1920*1440": (1920, 1440), "2560*1920": (2560, 1920), # 3:4 "1440*1920": (1440, 1920), "1920*2560": (1920, 2560), } width, height = 1920, 1080 # 默认 min_pixels = 3686400 if size: parts = size.split("*") if len(parts) == 2: w, h = int(parts[0]), int(parts[1]) width, height = w, h # 确保满足最小像素要求 if width * height < min_pixels: # 查找相同宽高比的常用分辨率 aspect_ratio = width / height for (w, h) in size_map.values(): if abs(w / h - aspect_ratio) < 0.01 and w * h >= min_pixels: width, height = w, h break else: # 没有找到合适的，按比例放大 scale = (min_pixels / (width * height)) ** 0.5 width = int(width * scale) height = int(height * scale) width = width if width % 2 == 0 else width + 1 height = height if height % 2 == 0 else height + 1 # 构建 extra_body extra_body = { "watermark": False, "sequential_image_generation": "disabled", } # 添加其他参数 if "seed" in kwargs: extra_body["seed"] = kwargs["seed"] if "quality" in kwargs: extra_body["quality"] = kwargs["quality"] if "style" in kwargs: extra_body["style"] = kwargs["style"] # 处理参考图 (图生图) image_urls = [] if image_paths and len(image_paths) > 0: # 处理参考图：支持 URL 和本地文件 ref_images = [] for p in image_paths: if p.startswith("http"): ref_images.append(p) elif os.path.exists(p): # 转换为 base64 URL import base64 with open(p, "rb") as f: img_data = base64.b64encode(f.read()).decode("utf-8") ext = os.path.splitext(p)[1].lower() mime = "image/png" if ext == ".png" else "image/jpeg" ref_images.append(f"data:{mime};base64,{img_data}") extra_body["image"] = ref_images # 调用 API if image_paths and len(image_paths) > 0: # 图生图 - image 放在 extra_body 中 response = self.client.images.generate( model=model, prompt=prompt, size=f"{width}x{height}", response_format="url", extra_body=extra_body, ) else: # 文生图 response = self.client.images.generate( model=model, prompt=prompt, size=f"{width}x{height}", response_format="url", extra_body=extra_body, ) # 下载图片到本地 generated_paths = [] if response.data: for idx, img_data in enumerate(response.data): if img_data.url: local_path = self._download_image( img_data.url, session_id, idx ) if local_path: generated_paths.append(local_path) return generated_paths def _download_image(self, url: str, session_id: str, idx: int) -> Optional[str]: """从URL下载图片到本地""" import requests # 构建存储路径 base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) result_dir = os.path.join(base_dir, "code", "result", "image", str(session_id)) os.makedirs(result_dir, exist_ok=True) file_name = f"seedream_{int(time.time())}_{idx}.png" file_path = os.path.join(result_dir, file_name) try: proxies = {"http": self.local_proxy, "https": self.local_proxy} if self.local_proxy else None response = requests.get(url, timeout=self.timeout, proxies=proxies) response.raise_for_status() with open(file_path, "wb") as f: f.write(response.content) return file_path except Exception as e: logging.error(f"Failed to download image from {url}: {e}") return None if __name__ == "__main__": import sys sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from config import Config # 加载 .env print("=== Seedream 可用性测试 ===") api_key = os.getenv("ARK_API_KEY", "") base_url = os.getenv("ARK_BASE_URL", "https://ark.cn-beijing.volces.com/api/v3") if not api_key: print("✗ ARK_API_KEY 未设置，跳过") sys.exit(1) print(f" API Key: {api_key[:6]}***{api_key[-4:]}") print(f" Base URL: {base_url}") client = SeedreamClient(api_key=api_key, base_url=base_url) # === 测试1: 文生图 === prompt = "星际穿越，黑洞，黑洞里冲出一辆支离破碎的复古列车，视觉冲击力，电影大片，末日既视感" print(f"\n[测试1: 文生图] Prompt: {prompt}") t0 = time.time() try: paths = client.generate_image( prompt=prompt, session_id="test_avail", model="doubao-seedream-5-0-260128", size="1920*1080", ) elapsed = time.time() - t0 if paths: print(f"✓ 生成 {len(paths)} 张图片 ({elapsed:.1f}s): {paths}") else: print(f"✗ 返回空列表 ({elapsed:.1f}s)") except Exception as e: print(f"✗ 图片生成失败: {e}") # === 测试2: 图生图 === # 需要一张已有的参考图路径 ref_image_path = "code/result/image/test_avail/test_input.png" if os.path.exists(ref_image_path): prompt_i2i = "将这只猫变成赛博朋克风格" print(f"\n[测试2: 图生图] Prompt: {prompt_i2i}") print(f" 参考图: {ref_image_path}") t0 = time.time() try: paths = client.generate_image( prompt=prompt_i2i, session_id="test_avail", model="doubao-seedream-5-0-260128", size="1920*1080", image_paths=[ref_image_path], ) elapsed = time.time() - t0 if paths: print(f"✓ 生成 {len(paths)} 张图片 ({elapsed:.1f}s): {paths}") else: print(f"✗ 返回空列表 ({elapsed:.1f}s)") except Exception as e: print(f"✗ 图生图失败: {e}") else: print(f"\n[测试2: 图生图] ✗ 参考图不存在: {ref_image_path}") print(" 跳过图生图测试，请先运行文生图测试生成参考图")