# !/usr/bin/env python3 import argparse import asyncio import json import os import random import time from datetime import datetime from pathlib import Path import httpx from openai import AsyncOpenAI # 配置 # LLM_BASE_URL os.getenv(LLM_BASE_URL) # LLM_MODEL os.getenv(LLM_MODEL) LLM_API_KEY os.getenv(LLM_API_KEY) or empty LLM_TIMEOUT float(os.getenv(LLM_TIMEOUT, 30)) # 对象种类 CATEGORIES [徽章, 戒指, 动物, 人物, 其他] CATEGORIES_STR 、.join(CATEGORIES) SYSTEM_PROMPT f你是一个文本分类专家。根据给定的文生图提示词判断主要对象的种类。 可选种类{CATEGORIES_STR} 严格只输出一个词不要解释。 async def classify(text: str) - str: 分类提示词 if not text.strip(): return client AsyncOpenAI(api_keyLLM_API_KEY, base_urlLLM_BASE_URL, http_clienthttpx.AsyncClient(timeoutLLM_TIMEOUT), ) try: resp await client.chat.completions.create(modelLLM_MODEL, messages[{role: system, content: SYSTEM_PROMPT}, {role: user, content: text.strip()}, ], temperature0.1, max_tokens20, ) result resp.choices[0].message.content.strip() # 提取第一个词 for cat in CATEGORIES: if cat in result: return cat return result finally: await client.close() async def main(): parser argparse.ArgumentParser() parser.add_argument(--output, typePath, defaultNone) parser.add_argument(--seed, typeint, defaultNone) args parser.parse_args() test_prompts [ # 人物 A single 3D designer art toy figure, presented in a strict horizontal triptych layout with left side view, middle front view, and right rear view. Strictly limited to a palette of exactly 4 pure colors: 1. Pure white for the head, gloves, and boots; 2. Dark ultramarine blue for the torso, arms, and legs; 3. Crimson red for the flowing cape; 4. Silver-grey for the facial mask details and chest logo., # 徽章 A beautifully crafted enamel pin badge, featuring a detailed dragon design in gold and red colors, with subtle metallic sheen, professional product photography, white background, high resolution, studio lighting., # 戒指 A stunning platinum engagement ring with a 2-carat diamond center stone, surrounded by smaller pavé-set diamonds, intricate filigree details on the band, elegant jewelry photography, soft natural lighting, macro shot, 8K., # 动物 A majestic golden eagle soaring through the sky with wings fully extended, dramatic sunset background, sharp focus on feathers, hyper-realistic wildlife photography, National Geographic style, 8K, cinematic composition., # 其他 A futuristic skyscraper with a sleek glass facade, reflecting the surrounding cityscape, dramatic night lighting, long exposure photography, architectural marvel, modern minimalist design, high resolution, 8K., # 其他 A delicious gourmet burger on a wooden table, perfectly chargrilled beef patty, melted cheddar cheese, crispy bacon, fresh lettuce and tomato, steam rising, warm ambient lighting, food photography, shallow depth of field, 8K, mouth-watering., ] results [] # 逐个测试 for i, prompt in enumerate(test_prompts, 1): total_start time.time() predicted await classify(prompt) results.append({predicted: predicted}) print(f{i} {predicted} {time.time() - total_start:.2f}) if __name__ __main__: asyncio.run(main())