# !/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_key=LLM_API_KEY, base_url=LLM_BASE_URL, http_client=httpx.AsyncClient(timeout=LLM_TIMEOUT), ) try: resp = await client.chat.completions.create(model=LLM_MODEL, messages=[{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": text.strip()}, ], temperature=0.1, max_tokens=20, ) 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", type=Path, default=None) parser.add_argument("--seed", type=int, default=None) 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())
