fastapi/japi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

rmbg同一批次内改为流式返回

Browse Source
This commit is contained in:
jingrow 2025-12-16 10:46:06 +00:00
parent 565115369d
commit 36f299403f
1 changed files with 79 additions and 68 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
apps/rmbg/service.py
View File

@ -607,14 +607,25 @@ class RmbgService:
local_batch_items = valid_items[local_batch_start:local_batch_end] local_batch_items = valid_items[local_batch_start:local_batch_end]
# 本地批次内并发调用 process_image，让全局队列有机会凑大 batch 并利用多 GPU # 本地批次内并发调用 process_image，让全局队列有机会凑大 batch 并利用多 GPU
tasks = [ # 改为按完成顺序流式返回，避免等待最慢的那张
self.process_image(image) tasks = []
for image, _, _, _ in local_batch_items task_meta = {}
] for image, _, index, url_str in local_batch_items:
batch_results = await asyncio.gather(*tasks, return_exceptions=True) t = asyncio.create_task(self.process_image(image))
tasks.append(t)
task_meta[t] = (index, url_str)
for (image, _, index, url_str), result_data in zip(local_batch_items, batch_results): # 使用 wait 循环而不是 as_completed，避免未等待的协程残留
if isinstance(result_data, Exception): pending_tasks = set(tasks)
while pending_tasks:
done, pending_tasks = await asyncio.wait(
pending_tasks, return_when=asyncio.FIRST_COMPLETED
)
for t in done:
index, url_str = task_meta[t]
try:
result_data = await t
except Exception as e:
error_count += 1 error_count += 1
completed_order += 1 completed_order += 1
result = { result = {
@ -622,8 +633,8 @@ class RmbgService:
"total": total, "total": total,
"original_url": url_str, "original_url": url_str,
"status": "error", "status": "error",
"error": str(result_data), "error": str(e),
"message": f"处理失败: {str(result_data)}", "message": f"处理失败: {str(e)}",
"success_count": success_count, "success_count": success_count,
"error_count": error_count, "error_count": error_count,
"completed_order": completed_order, "completed_order": completed_order,
@ -679,7 +690,7 @@ class RmbgService:
yield result yield result
except Exception as e: except Exception as e:
logger.error(f"批处理失败: {str(e)}") logger.error(f"批处理失败: {str(e)}")
for _, _, index, url_str, _ in valid_items: for _, _, index, url_str in valid_items:
error_count += 1 error_count += 1
completed_order += 1 completed_order += 1
result = { result = {