Update app.py

app.py

@@ -11,6 +11,8 @@ import rembg
 import onnxruntime as ort
 from concurrent.futures import ThreadPoolExecutor
 import asyncio
+import gc
+from asyncio import Semaphore
 
 app = FastAPI()
 
@@ -19,9 +21,15 @@ options = ort.SessionOptions()
 options.intra_op_num_threads = 2  # Limita o número de threads para evitar sobrecarga
 options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL  # Execução sequencial para melhor desempenho em CPU
 
+# Cria uma sessão global do rembg para reutilização
+session = rembg.new_session(model_name="u2net", session_options=options)
+
 # Pool de threads para executar tarefas bloqueantes
 executor = ThreadPoolExecutor(max_workers=4)
 
+# Semáforo para limitar requisições simultâneas
+semaphore = Semaphore(4)  # Ajuste conforme a capacidade da sua CPU/RAM
+
 def resize_image(image, max_size=512):
     """Redimensiona a imagem para uma largura máxima de 512px, mantendo a proporção."""
     width, height = image.size
@@ -32,38 +40,46 @@ def resize_image(image, max_size=512):
     return image
 
 def process_image(image_url):
-    # Baixa a imagem da URL fornecida
-    response = requests.get(image_url)
-    response.raise_for_status()
-    
-    # Abre a imagem usando Pillow
-    image = Image.open(BytesIO(response.content))
-    
-    # Pré-processamento: apenas redimensiona para 512px
-    image = resize_image(image, max_size=512)
-    
-    # Remove o fundo da imagem usando rembg
-    output = rembg.remove(image, session_options=options)
-    
-    # Pós-processamento: suaviza as bordas
-    output = output.filter(ImageFilter.SMOOTH_MORE)
-    
-    # Converte a imagem de volta para bytes
-    img_byte_arr = BytesIO()
-    output.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
-    
-    return img_byte_arr
+    try:
+        # Baixa a imagem da URL fornecida
+        response = requests.get(image_url)
+        response.raise_for_status()
+        
+        # Abre a imagem usando Pillow
+        image = Image.open(BytesIO(response.content))
+        
+        # Pré-processamento: apenas redimensiona para 512px
+        image = resize_image(image, max_size=512)
+        
+        # Remove o fundo da imagem usando rembg (reutiliza a sessão global)
+        output = rembg.remove(image, session=session)
+        
+        # Pós-processamento: suaviza as bordas
+        output = output.filter(ImageFilter.SMOOTH_MORE)
+        
+        # Converte a imagem de volta para bytes
+        img_byte_arr = BytesIO()
+        output.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        # Libera memória explicitamente
+        del image, output
+        gc.collect()  # Força a liberação de memória
+        
+        return img_byte_arr
+    except Exception as e:
+        raise e
 
 @app.get("/remove-background")
 async def remove_background(image_url: str):
-    try:
-        # Executa o processamento da imagem em um thread separado
-        loop = asyncio.get_event_loop()
-        img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
+    async with semaphore:  # Limita o número de requisições simultâneas
+        try:
+            # Executa o processamento da imagem em um thread separado
+            loop = asyncio.get_event_loop()
+            img_byte_arr = await loop.run_in_executor(executor, process_image, image_url)
+            
+            # Retorna a imagem processada diretamente no navegador
+            return StreamingResponse(img_byte_arr, media_type="image/png")
         
-        # Retorna a imagem processada diretamente no navegador
-        return StreamingResponse(img_byte_arr, media_type="image/png")
-    
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
+        except Exception as e:
+            raise HTTPException(status_code=400, detail=str(e))