ai-swap/backend/modules/comfyui/ai_copywriter.py

import io
import base64
from PIL import Image
from volcenginesdkarkruntime import Ark


client = Ark(
    api_key= "817dff39-5586-4f9b-acba-55004167c0b1",
    base_url="https://ark.cn-beijing.volces.com/api/v3",
)


default_prompt = """
## 角色：小红书博主，文案能力优秀
## 任务：基于图像内容，生成小红书高赞文案，文案风格为小红书风格，需要有标题，具体文案内容和话题标签（话题标签至少5个）。
## 要求：
- 字数控制在400字以内，多使用合适的emoji表情
- 除了输出标题、具体文案内容和话题标签，不要输出任何额外的内容
## 输出格式：
- 标题：<标题文本>
- 文案：<文案文本>
- 话题标签：<话题标签文本>
"""
    
def encode_image(pil_image):
    """将PIL Image对象转换为base64编码"""
    buffered = io.BytesIO()
    pil_image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode('utf-8')

def gen_copywriter(image_path, sys_prompt=default_prompt):
    base64_image = encode_image(image_path)
    response = client.chat.completions.create(
        model="doubao-1-5-vision-pro-32k-250115",
        temperature=1,
        max_tokens=200,
        messages=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": sys_prompt,
                    },
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:image/jpg;base64,{base64_image}"
                        },
                    },
                ],
            }
        ],
    )

    return response.choices[0].message.content


if __name__ == "__main__":

    with Image.open("./backend/data/cloth.png") as img:
        print(gen_copywriter(img))