feat: use image concurrency for outfit directories
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user