fix: make torch import optional and add PaddlePaddle GPU memory management
Problem: - Backend failed to start with ModuleNotFoundError for torch module - torch was imported as hard dependency but not in requirements.txt - Project uses PaddlePaddle which has its own CUDA implementation Changes: - Make torch import optional with try/except in ocr_service.py - Make torch import optional in pp_structure_enhanced.py - Add cleanup_gpu_memory() method using PaddlePaddle's memory management - Add check_gpu_memory() method to monitor available GPU memory - Use paddle.device.cuda.empty_cache() for GPU cleanup - Use torch.cuda only if TORCH_AVAILABLE flag is True - Add cleanup calls after OCR processing to prevent OOM errors - Add memory checks before GPU-intensive operations Benefits: - Backend can start without torch installed - GPU memory is properly managed using PaddlePaddle - Optional torch support provides additional memory monitoring - Prevents GPU OOM errors during document processing 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -9,7 +9,16 @@ import logging
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
import json
|
||||
import gc
|
||||
|
||||
# Optional torch import for additional GPU memory management
|
||||
try:
|
||||
import torch
|
||||
TORCH_AVAILABLE = True
|
||||
except ImportError:
|
||||
TORCH_AVAILABLE = False
|
||||
|
||||
import paddle
|
||||
from paddleocr import PPStructureV3
|
||||
from app.models.unified_document import ElementType
|
||||
|
||||
@@ -155,6 +164,18 @@ class PPStructureEnhanced:
|
||||
logger.error(f"Enhanced PP-StructureV3 analysis error: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
|
||||
# Clean up GPU memory on error
|
||||
try:
|
||||
if TORCH_AVAILABLE and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
if paddle.device.is_compiled_with_cuda():
|
||||
paddle.device.cuda.empty_cache()
|
||||
gc.collect()
|
||||
except:
|
||||
pass # Ignore cleanup errors
|
||||
|
||||
return {
|
||||
'elements': [],
|
||||
'total_elements': 0,
|
||||
|
||||
Reference in New Issue
Block a user