feat: use image concurrency for outfit directories

This commit is contained in:
2026-06-22 18:10:54 +08:00
parent 00ed41db14
commit e83dd07dc0
7 changed files with 131 additions and 45 deletions
+39
View File
@@ -3,6 +3,8 @@ import os
import shutil
import sys
import tempfile
import threading
import time
import unittest
from io import BytesIO
from pathlib import Path
@@ -262,6 +264,22 @@ class TestAiOutfitCore(unittest.TestCase):
self.assertEqual(len(result.output_paths), 2)
self.assertFalse((out / "a" / "img2.jpg").exists())
def test_generate_directory_uses_image_concurrency(self):
from core.ai_outfit import generate_outfit_image
d = self._make_dir_with_images("a", files=("img1.png", "img2.png", "img3.png", "img4.png"))
out = self.tmp / "out"
client = _ConcurrentRecordingClient(self._image_bytes())
result = generate_outfit_image(
self._dir_task(d), "x", out, model_config={}, api_client=client,
image_concurrency=2)
self.assertTrue(result.success, result.error)
self.assertEqual(client.calls, 4)
self.assertGreaterEqual(client.max_active, 2)
self.assertEqual(len(result.output_paths), 4)
class _FakeClient:
def __init__(self, image_bytes):
@@ -309,5 +327,26 @@ class _FailOnClient:
return self._image_bytes
class _ConcurrentRecordingClient:
"""Thread-safe fake client that records concurrent generate() calls."""
def __init__(self, image_bytes):
self._image_bytes = image_bytes
self.calls = 0
self.active = 0
self.max_active = 0
self._lock = threading.Lock()
def generate(self, prompt, image_path, resolution="1K"):
with self._lock:
self.calls += 1
self.active += 1
self.max_active = max(self.max_active, self.active)
time.sleep(0.05)
with self._lock:
self.active -= 1
return self._image_bytes
if __name__ == "__main__":
unittest.main()