Tool_Agent/tools/local/seedream/generate_seedream.py

"""
SeeDream 4.0 图片生成 - Python逻辑封装
支持标准 OpenAI Image API 格式的调用
"""

import os
import requests
from dotenv import load_dotenv

# 向上加载环境参数
env_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), ".env")
if os.path.exists(env_path):
    load_dotenv(env_path)
else:
    load_dotenv()

APIYI_KEY = os.getenv("SEEDREAM_API_KEY") or os.getenv("APIYI_KEY")
if APIYI_KEY and not APIYI_KEY.startswith("sk-"):
    APIYI_KEY = f"sk-{APIYI_KEY}"
APIYI_BASE_URL = os.getenv("APIYI_BASE_URL", "https://api.apiyi.com/v1")

if not APIYI_KEY:
    print("错误: 请在 .env 文件中设置 APIYI_KEY 或 SEEDREAM_API_KEY")
    exit(1)

def generate_seedream_image(
        prompt: str,
        model: str = "seedream-4-5-251128",
        size: str = "2048x2048",
        quality: str = "hd",
        n: int = 1,
        response_format: str = "url",
        image_url: str = None
    ):
    """
    生成图片核心函数
    
    Args:
        prompt: 图片描述提示词
        model: 模型选择 (seedream-4-5-251128)
        size: 分辨率/尺寸 (如 "1K", "2K", "4K", "2048x2048")
        quality: 图片质量 (standard, hd)
        n: 生成数量
        response_format: 响应格式 (url 推荐)
        image_url: 可选的参考图外链，用于图生图
    """
    
    url = f"{APIYI_BASE_URL}/images/generations"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {APIYI_KEY}"
    }
    
    payload = {
        "model": model,
        "prompt": prompt,
        "size": size,
        "quality": quality,
        "n": n,
        "response_format": response_format
    }
    
    if image_url:
        payload["image"] = image_url
    
    print(f"📡 发送请求...")
    print(f" - API: {url}")
    print(f" - Prompt: {prompt}")
    
    try:
        response = requests.post(url, headers=headers, json=payload, timeout=120)
        
        if response.status_code != 200:
            error_msg = f"API请求失败，状态码: {response.status_code}"
            try:
                error_detail = response.json()
                error_msg += f", 错误详情: {error_detail}"
            except Exception:
                error_msg += f", 响应内容: {response.text[:500]}"
            print(f"❌ {error_msg}")
            return {"error": error_msg}
            
        print("✅ API请求成功")
        return response.json()
    except Exception as e:
        error_msg = f"发生异常: {str(e)}"
        print(f"❌ {error_msg}")
        return {"error": error_msg}

if __name__ == "__main__":
    result = generate_seedream_image("A futuristic utopia, flying cars, beautiful 4K setup")
    if "data" in result:
        print(f"结果: {result['data'][0]['url']}")
    else:
        print(f"完整输出: {result}")