Publish current ComfyUI O1Key code baseline
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
This commit is contained in:
Jony
@@ -0,0 +1,176 @@
|
||||
"""Offline tests for the image and video AI-writing helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
from io import BytesIO
|
||||
import json
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
sys.path.insert(0, str(ROOT))
|
||||
|
||||
from utils import chat_support as CHAT # noqa: E402
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, status, payload):
|
||||
self.status = status
|
||||
self.payload = payload
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, _exc_type, _exc, _tb):
|
||||
return False
|
||||
|
||||
async def json(self, **_kwargs):
|
||||
return self.payload
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, response):
|
||||
self.response = response
|
||||
self.calls = []
|
||||
|
||||
def post(self, url, **kwargs):
|
||||
self.calls.append((url, kwargs))
|
||||
return self.response
|
||||
|
||||
|
||||
class PromptOptimizerTests(unittest.TestCase):
|
||||
def test_payload_uses_fixed_model_high_reasoning_and_ordered_images(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
Image.new("RGB", (80, 120), "red").save(root / "first.png")
|
||||
Image.new("RGB", (120, 80), "blue").save(root / "second.png")
|
||||
|
||||
payload = CHAT.build_prompt_optimization_payload(
|
||||
"保留人物身份，改成雨夜街景",
|
||||
[
|
||||
{"name": "first.png", "subfolder": "", "type": "input"},
|
||||
{"name": "second.png", "subfolder": "", "type": "input"},
|
||||
],
|
||||
directory,
|
||||
)
|
||||
|
||||
self.assertEqual(payload["model"], "gpt-5.6-sol")
|
||||
self.assertEqual(payload["reasoning_effort"], "high")
|
||||
self.assertFalse(payload["stream"])
|
||||
self.assertIn("视觉元素绑定优先", payload["messages"][0]["content"])
|
||||
self.assertIn("必须保持不变", payload["messages"][0]["content"])
|
||||
|
||||
content = payload["messages"][1]["content"]
|
||||
self.assertEqual(content[1]["text"], "参考图 1（上传顺序第 1 张）")
|
||||
self.assertEqual(content[3]["text"], "参考图 2（上传顺序第 2 张）")
|
||||
image_items = [item for item in content if item["type"] == "image_url"]
|
||||
self.assertEqual(len(image_items), 2)
|
||||
for item in image_items:
|
||||
url = item["image_url"]["url"]
|
||||
self.assertTrue(url.startswith("data:image/jpeg;base64,"))
|
||||
with Image.open(BytesIO(base64.b64decode(url.split(",", 1)[1]))) as image:
|
||||
self.assertEqual(image.format, "JPEG")
|
||||
body_size = len(
|
||||
json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
)
|
||||
self.assertLessEqual(body_size, CHAT.PROMPT_OPTIMIZER_BODY_LIMIT_BYTES)
|
||||
|
||||
def test_rejects_reference_path_outside_input_directory(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
with self.assertRaisesRegex(ValueError, "路径不安全"):
|
||||
CHAT.build_prompt_optimization_payload(
|
||||
"测试提示词",
|
||||
[{"name": "outside.png", "subfolder": "..", "type": "input"}],
|
||||
directory,
|
||||
)
|
||||
|
||||
def test_non_streaming_request_extracts_only_the_final_prompt(self):
|
||||
session = _FakeSession(_FakeResponse(200, {
|
||||
"choices": [{"message": {"content": "```text\n最终视觉提示词\n```"}}],
|
||||
}))
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
result = asyncio.run(CHAT.optimize_image_prompt(
|
||||
session,
|
||||
"https://example.invalid",
|
||||
"offline-test-token",
|
||||
"原始提示词",
|
||||
[],
|
||||
directory,
|
||||
))
|
||||
|
||||
self.assertEqual(result, "最终视觉提示词")
|
||||
self.assertEqual(len(session.calls), 1)
|
||||
url, kwargs = session.calls[0]
|
||||
self.assertEqual(url, "https://example.invalid/v1/chat/completions")
|
||||
self.assertEqual(kwargs["json"]["model"], "gpt-5.6-sol")
|
||||
self.assertEqual(kwargs["json"]["reasoning_effort"], "high")
|
||||
self.assertFalse(kwargs["json"]["stream"])
|
||||
|
||||
def test_video_payload_uses_dedicated_preset_and_ordered_frame_roles(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
Image.new("RGB", (96, 128), "green").save(root / "start.png")
|
||||
Image.new("RGB", (128, 96), "yellow").save(root / "end.png")
|
||||
payload = CHAT.build_video_prompt_writing_payload(
|
||||
"人物从门口走到窗前",
|
||||
[
|
||||
{"name": "start.png", "subfolder": "", "type": "input"},
|
||||
{"name": "end.png", "subfolder": "", "type": "input"},
|
||||
],
|
||||
directory,
|
||||
{
|
||||
"generation_mode": "first_last_frame",
|
||||
"duration": "8",
|
||||
"aspect_ratio": "16:9",
|
||||
"generate_audio": True,
|
||||
"reference_video_count": 1,
|
||||
"reference_audio_count": 2,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(payload["model"], "gpt-5.6-sol")
|
||||
self.assertEqual(payload["reasoning_effort"], "high")
|
||||
self.assertFalse(payload["stream"])
|
||||
system_prompt = payload["messages"][0]["content"]
|
||||
self.assertIn("连续的时间过程", system_prompt)
|
||||
self.assertIn("首尾帧生视频", system_prompt)
|
||||
content = payload["messages"][1]["content"]
|
||||
self.assertIn("生成模式：首尾帧生视频", content[0]["text"])
|
||||
self.assertIn("目标时长：8", content[0]["text"])
|
||||
self.assertIn("生成音频：开启", content[0]["text"])
|
||||
self.assertIn("参考视频 1 个、参考音频 2 个", content[0]["text"])
|
||||
self.assertEqual(content[1]["text"], "首帧（第 1 张分析图）")
|
||||
self.assertEqual(content[3]["text"], "尾帧（第 2 张分析图）")
|
||||
self.assertEqual(len([item for item in content if item["type"] == "image_url"]), 2)
|
||||
|
||||
def test_video_request_uses_non_streaming_completion(self):
|
||||
session = _FakeSession(_FakeResponse(200, {
|
||||
"choices": [{"message": {"content": "连续镜头中的最终视频提示词"}}],
|
||||
}))
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
result = asyncio.run(CHAT.write_video_prompt(
|
||||
session,
|
||||
"https://example.invalid",
|
||||
"offline-test-token",
|
||||
"原始视频创意",
|
||||
[],
|
||||
directory,
|
||||
{"generation_mode": "text", "duration": "5"},
|
||||
))
|
||||
|
||||
self.assertEqual(result, "连续镜头中的最终视频提示词")
|
||||
self.assertEqual(session.calls[0][0], "https://example.invalid/v1/chat/completions")
|
||||
request = session.calls[0][1]["json"]
|
||||
self.assertIn("AI 视频生成提示词编写助手", request["messages"][0]["content"])
|
||||
self.assertFalse(request["stream"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=2)
|
||||
Reference in New Issue
Block a user