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()
|
||||
|
||||
@@ -64,6 +64,15 @@ class TestAiOutfitPanelDefaults(unittest.TestCase):
|
||||
|
||||
self.assertEqual(Path(panel._output_edit.text()), saved)
|
||||
|
||||
def test_generation_setting_label_is_image_concurrency(self):
|
||||
from PySide6.QtWidgets import QLabel
|
||||
|
||||
panel = self._panel()
|
||||
|
||||
labels = [label.text() for label in panel.findChildren(QLabel)]
|
||||
|
||||
self.assertIn("图片并发数", labels)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -91,15 +91,17 @@ class TestOutfitBatchRunner(unittest.TestCase):
|
||||
self.assertEqual(summary.success_count, 1)
|
||||
self.assertEqual(summary.results[0].attempts, 2)
|
||||
|
||||
def test_concurrency_runs_more_than_one_task_at_once(self):
|
||||
def test_concurrency_option_does_not_parallelize_excel_rows(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
active = [0]
|
||||
max_active = [0]
|
||||
started = []
|
||||
lock = threading.Lock()
|
||||
|
||||
def generate(task):
|
||||
with lock:
|
||||
started.append(task.row_index)
|
||||
active[0] += 1
|
||||
max_active[0] = max(max_active[0], active[0])
|
||||
time.sleep(0.05)
|
||||
@@ -116,7 +118,8 @@ class TestOutfitBatchRunner(unittest.TestCase):
|
||||
summary = runner.run()
|
||||
|
||||
self.assertEqual(summary.success_count, 4)
|
||||
self.assertGreaterEqual(max_active[0], 2)
|
||||
self.assertEqual(max_active[0], 1)
|
||||
self.assertEqual(started, [2, 3, 4, 5])
|
||||
|
||||
def test_task_cooldown_is_applied(self):
|
||||
from core.outfit_batch import OutfitBatchOptions, OutfitBatchRunner
|
||||
|
||||
Reference in New Issue
Block a user