优化rmbg并发逻辑，实测并发生效

apps/rmbg/service.py

@@ -13,13 +13,18 @@ import asyncio
 import io
 import uuid
 import httpx
+import logging
+from concurrent.futures import ThreadPoolExecutor
 from settings import settings
 
-# 关闭不必要的警告
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
 warnings.filterwarnings("ignore", category=UserWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
-
-# 设置torch精度
 torch.set_float32_matmul_precision("high")
 
 class RmbgService:
@@ -30,25 +35,43 @@ class RmbgService:
         self.device = None
         self.save_dir = settings.save_dir
         self.download_url = settings.download_url
-        # 确保保存目录存在
         os.makedirs(self.save_dir, exist_ok=True)
-        # 创建异步HTTP客户端（复用连接，提高性能）
+        
-        self.http_client = httpx.AsyncClient(timeout=30.0, limits=httpx.Limits(max_keepalive_connections=20))
+        self.http_client = httpx.AsyncClient(
+            timeout=30.0, 
+            limits=httpx.Limits(
+                max_keepalive_connections=50,
+                max_connections=100
+            )
+        )
+        self.executor = ThreadPoolExecutor(max_workers=settings.max_workers)
+        self._gpu_semaphore = None
+        self._max_gpu_concurrent = settings.max_gpu_concurrent
         self._load_model()
     
+    @property
+    def gpu_semaphore(self):
+        """延迟初始化GPU信号量"""
+        if self._gpu_semaphore is None:
+            if self._max_gpu_concurrent == 0:
+                return None
+            try:
+                loop = asyncio.get_event_loop()
+                self._gpu_semaphore = asyncio.Semaphore(self._max_gpu_concurrent)
+            except RuntimeError:
+                self._gpu_semaphore = asyncio.Semaphore(self._max_gpu_concurrent)
+        return self._gpu_semaphore
+        
     def _load_model(self):
         """加载模型"""
-        # 设置设备
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        t0 = time.time()
         self.model = AutoModelForImageSegmentation.from_pretrained(self.model_path, trust_remote_code=True)
         self.model = self.model.to(self.device)
-        self.model.eval()  # 设置为评估模式
+        self.model.eval()
 
     def _process_image_sync(self, image):
-        """同步处理图像，移除背景（内部方法，在线程池中执行）"""
+        """同步处理图像，移除背景"""
         image_size = image.size
-        # 转换图像
         transform_image = transforms.Compose([
             transforms.Resize((1024, 1024)),
             transforms.ToTensor(),
@@ -56,30 +79,103 @@ class RmbgService:
         ])
         input_images = transform_image(image).unsqueeze(0).to(self.device)
         
-        # 推理
         with torch.no_grad():
             preds = self.model(input_images)[-1].sigmoid().cpu()
         
-        # 处理预测结果
         pred = preds[0].squeeze()
         pred_pil = transforms.ToPILImage()(pred)
         mask = pred_pil.resize(image_size)
-        
-        # 添加透明通道
         image.putalpha(mask)
         
-        # 清理显存
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
             gc.collect()
         
         return image
     
+    def _process_batch_images_sync(self, images_with_info):
+        """
+        批量处理图像（充分利用GPU并行能力）
+        
+        Args:
+            images_with_info: [(image, image_size, index), ...] 图像和元信息列表
+            
+        Returns:
+            [(processed_image, index), ...] 处理后的图像和索引
+        """
+        if not images_with_info:
+            return []
+        
+        try:
+            transform_image = transforms.Compose([
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ])
+            
+            batch_images = []
+            for image, image_size, index in images_with_info:
+                try:
+                    transformed = transform_image(image)
+                    batch_images.append(transformed)
+                except Exception as e:
+                    logger.error(f"图片转换失败 (index={index}): {str(e)}")
+                    raise
+            
+            if not batch_images:
+                raise Exception("没有有效的图片可以处理")
+            
+            input_batch = torch.stack(batch_images).to(self.device)
+            
+            with torch.no_grad():
+                model_output = self.model(input_batch)
+                if isinstance(model_output, (list, tuple)):
+                    preds = model_output[-1].sigmoid().cpu()
+                else:
+                    preds = model_output.sigmoid().cpu()
+            
+            results = []
+            for i, (image, image_size, index) in enumerate(images_with_info):
+                try:
+                    if len(preds.shape) == 4:
+                        pred = preds[i].squeeze()
+                    elif len(preds.shape) == 3:
+                        pred = preds[i]
+                    else:
+                        pred = preds[i].squeeze()
+                    
+                    pred_pil = transforms.ToPILImage()(pred)
+                    mask = pred_pil.resize(image_size)
+                    result_image = image.copy()
+                    result_image.putalpha(mask)
+                    results.append((result_image, index))
+                except Exception as e:
+                    logger.error(f"处理预测结果失败 (index={index}): {str(e)}")
+                    raise
+            
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                gc.collect()
+            
+            return results
+            
+        except Exception as e:
+            logger.error(f"批处理失败: {str(e)}")
+            raise
+    
     async def process_image(self, image):
-        """异步处理图像，移除背景（在线程池中执行同步操作）"""
+        """异步处理图像，移除背景"""
-        # 将同步的GPU操作放到线程池中执行，避免阻塞事件循环
         loop = asyncio.get_event_loop()
-        return await loop.run_in_executor(None, self._process_image_sync, image)
+        return await loop.run_in_executor(self.executor, self._process_image_sync, image)
+    
+    async def process_batch_images(self, images_with_info):
+        """异步批量处理图像"""
+        loop = asyncio.get_event_loop()
+        return await loop.run_in_executor(
+            self.executor, 
+            self._process_batch_images_sync, 
+            images_with_info
+        )
 
     def image_to_base64(self, image):
         """将PIL Image对象转换为base64字符串"""
@@ -88,23 +184,10 @@ class RmbgService:
         return base64.b64encode(buffered.getvalue()).decode('utf-8')
     
     def save_image_to_file(self, image):
-        """
+        """保存图片到文件并返回URL"""
-        保存图片到jfile/files目录并返回URL
-        
-        Args:
-            image: PIL Image对象
-            
-        Returns:
-            图片URL
-        """
-        # 生成唯一文件名
         filename = f"rmbg_{uuid.uuid4().hex[:10]}.png"
         file_path = os.path.join(self.save_dir, filename)
-        
-        # 保存图片
         image.save(file_path, format="PNG")
-        
-        # 构建URL
         image_url = f"{self.download_url}/{filename}"
         return image_url
 
@@ -120,32 +203,26 @@ class RmbgService:
         """
         temp_file = None
         try:
-            # 检查是否是URL
             if self.is_valid_url(image_path):
                 try:
-                    # 异步下载图片到临时文件
                     temp_file = await self.download_image(image_path)
                     image_path = temp_file
                 except Exception as e:
                     raise Exception(f"下载图片失败: {e}")
             
-            # 验证输入文件是否存在
             if not os.path.exists(image_path):
                 raise FileNotFoundError(f"输入图像不存在: {image_path}")
             
-            # 加载图像（IO操作，在线程池中执行）
             loop = asyncio.get_event_loop()
             image = await loop.run_in_executor(
-                None, 
+                self.executor, 
                 lambda: Image.open(image_path).convert("RGB")
             )
             
-            # 异步处理图像
             image_no_bg = await self.process_image(image)
             
-            # 保存图片到文件并获取URL（IO操作，在线程池中执行）
             image_url = await loop.run_in_executor(
-                None,
+                self.executor,
                 self.save_image_to_file,
                 image_no_bg
             )
@@ -156,7 +233,6 @@ class RmbgService:
             }
         
         finally:
-            # 清理临时文件
             if temp_file and os.path.exists(temp_file):
                 try:
                     os.unlink(temp_file)
@@ -174,19 +250,16 @@ class RmbgService:
             处理后的图像内容
         """
         try:
-            # 从文件内容创建PIL Image对象（IO操作，在线程池中执行）
             loop = asyncio.get_event_loop()
             image = await loop.run_in_executor(
-                None,
+                self.executor,
                 lambda: Image.open(io.BytesIO(file_content)).convert("RGB")
             )
             
-            # 异步处理图像
             image_no_bg = await self.process_image(image)
             
-            # 保存图片到文件并获取URL（IO操作，在线程池中执行）
             image_url = await loop.run_in_executor(
-                None,
+                self.executor,
                 self.save_image_to_file,
                 image_no_bg
             )
@@ -200,60 +273,77 @@ class RmbgService:
             raise Exception(f"处理图片失败: {e}")
 
     async def process_batch(self, urls):
-        """
+        """批量处理多个URL图像，流水线并发模式"""
-        批量处理多个URL图像，并发处理并流式返回结果
-        
-        Args:
-            urls: 图片URL列表
-            
-        Yields:
-            每个图片的处理结果（按完成顺序返回）
-        """
         total = len(urls)
         success_count = 0
         error_count = 0
+        batch_start_time = time.time()
         
-        # 创建并发任务
+        async def download_and_process(index, url):
-        async def process_single_url(index, url):
+            """下载并处理单张图片"""
-            """处理单个URL的包装函数"""
+            url_str = str(url)
             try:
-                url_str = str(url)
+                if self.is_valid_url(url_str):
-                result = await self.remove_background(url_str)
+                    temp_file = await self.download_image(url_str)
+                    image = await asyncio.get_event_loop().run_in_executor(
+                        self.executor,
+                        lambda: Image.open(temp_file).convert("RGB")
+                    )
+                    os.unlink(temp_file)
+                else:
+                    image = await asyncio.get_event_loop().run_in_executor(
+                        self.executor,
+                        lambda: Image.open(url_str).convert("RGB")
+                    )
+                
+                processed_image = await self.process_image(image)
+                
+                loop = asyncio.get_event_loop()
+                image_url = await loop.run_in_executor(
+                    self.executor,
+                    self.save_image_to_file,
+                    processed_image
+                )
+                
                 return {
                     "index": index,
                     "total": total,
                     "original_url": url_str,
                     "status": "success",
-                    "image_url": result["image_url"],
+                    "image_url": image_url,
                     "message": "处理成功"
                 }
+                
             except Exception as e:
+                logger.error(f"处理失败 (index={index}): {str(e)}")
                 return {
                     "index": index,
                     "total": total,
-                    "original_url": str(url),
+                    "original_url": url_str,
                     "status": "error",
                     "error": str(e),
                     "message": f"处理失败: {str(e)}"
                 }
         
-        # 创建所有任务
         tasks = [
-            process_single_url(i, url) 
+            download_and_process(i, url)
             for i, url in enumerate(urls, 1)
         ]
         
-        # 并发执行所有任务，使用as_completed按完成顺序返回
+        completed_order = 0
         for coro in asyncio.as_completed(tasks):
             result = await coro
+            completed_order += 1
+            
             if result["status"] == "success":
                 success_count += 1
             else:
                 error_count += 1
             
-            # 更新统计信息
             result["success_count"] = success_count
             result["error_count"] = error_count
+            result["completed_order"] = completed_order
+            result["batch_elapsed"] = round(time.time() - batch_start_time, 2)
             
             yield result
 
@@ -271,7 +361,6 @@ class RmbgService:
             response = await self.http_client.get(url)
             response.raise_for_status()
             
-            # 创建临时文件并写入内容（IO操作，在线程池中执行）
             def write_temp_file(content):
                 temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
                 temp_file.write(content)
@@ -280,7 +369,7 @@ class RmbgService:
             
             loop = asyncio.get_event_loop()
             temp_file_path = await loop.run_in_executor(
-                None,
+                self.executor,
                 write_temp_file,
                 response.content
             )
@@ -291,9 +380,8 @@ class RmbgService:
     
     async def cleanup(self):
         """清理资源"""
-        # 关闭HTTP客户端
         await self.http_client.aclose()
+        self.executor.shutdown(wait=True)
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
         gc.collect() 
-        print("资源已清理") 

apps/rmbg/settings.py

@@ -26,6 +26,10 @@ class Settings(BaseSettings):
     jingrow_api_key: Optional[str] = None
     jingrow_api_secret: Optional[str] = None
 
+    # 并发控制配置
+    max_workers: int = 30  # 线程池最大工作线程数（根据CPU核心数调整，22核44线程可设置20-30）
+    max_gpu_concurrent: int = 0  # GPU最大并发数（0表示不限制，根据显存大小设置，24GB显存建议10-15）
+
     class Config:
         env_file = ".env"