self-media-james/scripts/jimeng_gen.py

"""即梦 AI 图片生成工具"""

import json
import os
import sys
import time
import base64
import datetime
import hashlib
import hmac
import requests

# ===== 配置 =====
ACCESS_KEY = 'AKLTOWVjYmE2NGVhZTFhNDQ2OThiYTNhZDdjZTZiYTc3ZTQ'
SECRET_KEY = 'WXpVd016azFPVEF3TURjNE5EbGxOV0psTlRNek1qaGpaalEyWkdKa1kyRQ=='
HOST = 'visual.volcengineapi.com'
ENDPOINT = 'https://visual.volcengineapi.com'
REGION = 'cn-north-1'
SERVICE = 'cv'


def _sign(key, msg):
    return hmac.new(key, msg.encode('utf-8'), hashlib.sha256).digest()


def _get_signature_key(key, date_stamp, region_name, service_name):
    k_date = _sign(key.encode('utf-8'), date_stamp)
    k_region = _sign(k_date, region_name)
    k_service = _sign(k_region, service_name)
    k_signing = _sign(k_service, 'request')
    return k_signing


def _format_query(parameters):
    return '&'.join(f'{k}={parameters[k]}' for k in sorted(parameters))


def jimeng_request(action, body_params):
    """发送即梦 API 请求"""
    t = datetime.datetime.now(datetime.UTC)
    current_date = t.strftime('%Y%m%dT%H%M%SZ')
    datestamp = t.strftime('%Y%m%d')

    query_params = _format_query({'Action': action, 'Version': '2022-08-31'})
    req_body = json.dumps(body_params)
    payload_hash = hashlib.sha256(req_body.encode('utf-8')).hexdigest()

    signed_headers = 'content-type;host;x-content-sha256;x-date'
    canonical_headers = (
        f'content-type:application/json\n'
        f'host:{HOST}\n'
        f'x-content-sha256:{payload_hash}\n'
        f'x-date:{current_date}\n'
    )
    canonical_request = f'POST\n/\n{query_params}\n{canonical_headers}\n{signed_headers}\n{payload_hash}'

    credential_scope = f'{datestamp}/{REGION}/{SERVICE}/request'
    string_to_sign = (
        f'HMAC-SHA256\n{current_date}\n{credential_scope}\n'
        + hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()
    )

    signing_key = _get_signature_key(SECRET_KEY, datestamp, REGION, SERVICE)
    signature = hmac.new(signing_key, string_to_sign.encode('utf-8'), hashlib.sha256).hexdigest()

    authorization = (
        f'HMAC-SHA256 Credential={ACCESS_KEY}/{credential_scope}, '
        f'SignedHeaders={signed_headers}, Signature={signature}'
    )

    headers = {
        'X-Date': current_date,
        'Authorization': authorization,
        'X-Content-Sha256': payload_hash,
        'Content-Type': 'application/json'
    }

    r = requests.post(f'{ENDPOINT}?{query_params}', headers=headers, data=req_body)
    return r.json()


def _ensure_dir(file_path):
    dir_name = os.path.dirname(file_path)
    if dir_name:
        os.makedirs(dir_name, exist_ok=True)


def generate_image(prompt, output_path, width=1024, height=1024):
    """文生图：提交任务并轮询获取结果，保存图片到本地"""
    print(f"[即梦] 生成图片: {prompt[:50]}...")

    # 1. 提交任务
    submit_resp = jimeng_request('CVSync2AsyncSubmitTask', {
        'req_key': 'jimeng_t2i_v40',
        'prompt': prompt,
        'width': width,
        'height': height
    })

    if submit_resp.get('code') != 10000:
        print(f"[即梦] 提交失败: {submit_resp}")
        return None

    task_id = submit_resp['data']['task_id']
    print(f"[即梦] 任务已提交, task_id: {task_id}")

    # 2. 轮询查询结果
    for i in range(30):
        time.sleep(5)
        result = jimeng_request('CVSync2AsyncGetResult', {
            'req_key': 'jimeng_t2i_v40',
            'task_id': task_id
        })

        resp_code = result.get('code')
        resp_data = result.get('data', {})

        # 检查任务是否完成
        if resp_code == 10000:
            # 尝试获取图片数据
            image_urls = resp_data.get('image_urls', [])
            binary_data = resp_data.get('binary_data_base64', [])

            if image_urls:
                # 下载图片
                img_url = image_urls[0]
                img_data = requests.get(img_url).content
                _ensure_dir(output_path)
                with open(output_path, 'wb') as f:
                    f.write(img_data)
                print(f"[即梦] 图片已保存: {output_path}")
                return output_path

            if binary_data:
                # base64 解码保存
                img_data = base64.b64decode(binary_data[0])
                _ensure_dir(output_path)
                with open(output_path, 'wb') as f:
                    f.write(img_data)
                print(f"[即梦] 图片已保存: {output_path}")
                return output_path

            # 任务可能还在处理中
            status = resp_data.get('status', 'unknown')
            if status == 'done':
                print(f"[即梦] 任务完成但未找到图片数据: {json.dumps(result, ensure_ascii=False)[:200]}")
                return None

        print(f"[即梦] 等待中... ({i+1}/30)")

    print("[即梦] 超时")
    return None


if __name__ == '__main__':
    prompt = sys.argv[1] if len(sys.argv) > 1 else "一只可爱的卡通龙虾站在发光的电路板上，科技感，扁平插画风格，蓝色调"
    output = sys.argv[2] if len(sys.argv) > 2 else "test_output.png"
    generate_image(prompt, output)