Update app.py

app.py

@@ -2,14 +2,16 @@ import os
 import tempfile
 import time
 import json
+import threading
+import gc
 from pathlib import Path
 import uuid
 import logging
+from datetime import datetime, timedelta
 
 import torch
 import yt_dlp as youtube_dl
 from flask import Flask, request, jsonify
-# from flask_cors import CORS
 from transformers import pipeline
 from transformers.pipelines.audio_utils import ffmpeg_read
 import ffmpeg
@@ -19,7 +21,6 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = Flask(__name__)
-# CORS(app)
 
 # Configuration
 MODEL_NAME = "openai/whisper-large-v3"
@@ -27,23 +28,112 @@ BATCH_SIZE = 8
 FILE_LIMIT_MB = 1000
 YT_LENGTH_LIMIT_S = 3600  # 1 hour limit for YouTube
 MAX_FILE_SIZE = FILE_LIMIT_MB * 1024 * 1024  # Convert to bytes
+MODEL_TIMEOUT_MINUTES = 60  # مدت زمان نگهداری مدل در حافظه (به دقیقه)
 
 # Device configuration
 device = 0 if torch.cuda.is_available() else "cpu"
 logger.info(f"Using device: {device}")
 
-# Initialize Whisper pipeline
-try:
-    pipe = pipeline(
-        task="automatic-speech-recognition",
-        model=MODEL_NAME,
-        chunk_length_s=30,
-        device=device,
-    )
-    logger.info("Whisper model loaded successfully")
-except Exception as e:
-    logger.error(f"Error loading Whisper model: {e}")
-    raise
+# Global model management
+class ModelManager:
+    def __init__(self):
+        self.pipe = None
+        self.last_used = None
+        self.model_lock = threading.Lock()
+        self.cleanup_timer = None
+        self.is_loading = False
+        
+    def load_model(self):
+        """بارگذاری مدل در صورت عدم وجود"""
+        with self.model_lock:
+            if self.pipe is not None:
+                self.last_used = datetime.now()
+                return self.pipe
+            
+            if self.is_loading:
+                # اگر مدل در حال بارگذاری است، منتظر بمانید
+                while self.is_loading:
+                    time.sleep(0.5)
+                return self.pipe
+            
+            try:
+                self.is_loading = True
+                logger.info("Loading Whisper model...")
+                
+                self.pipe = pipeline(
+                    task="automatic-speech-recognition",
+                    model=MODEL_NAME,
+                    chunk_length_s=30,
+                    device=device,
+                )
+                
+                self.last_used = datetime.now()
+                self.start_cleanup_timer()
+                logger.info("Whisper model loaded successfully")
+                
+            except Exception as e:
+                logger.error(f"Error loading Whisper model: {e}")
+                self.pipe = None
+                raise
+            finally:
+                self.is_loading = False
+                
+            return self.pipe
+    
+    def get_model(self):
+        """دریافت مدل (با بارگذاری در صورت نیاز)"""
+        if self.pipe is None:
+            return self.load_model()
+        
+        self.last_used = datetime.now()
+        return self.pipe
+    
+    def cleanup_model(self):
+        """پاکسازی مدل از حافظه"""
+        with self.model_lock:
+            if self.pipe is not None:
+                logger.info("Cleaning up model from memory...")
+                del self.pipe
+                self.pipe = None
+                
+                # پاکسازی کش CUDA در صورت استفاده از GPU
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+                
+                # فراخوانی garbage collector
+                gc.collect()
+                logger.info("Model cleanup completed")
+                
+            if self.cleanup_timer:
+                self.cleanup_timer.cancel()
+                self.cleanup_timer = None
+    
+    def start_cleanup_timer(self):
+        """شروع تایمر پاکسازی"""
+        if self.cleanup_timer:
+            self.cleanup_timer.cancel()
+        
+        self.cleanup_timer = threading.Timer(
+            MODEL_TIMEOUT_MINUTES * 60, 
+            self.check_and_cleanup
+        )
+        self.cleanup_timer.start()
+    
+    def check_and_cleanup(self):
+        """بررسی و پاکسازی مدل در صورت عدم استفاده"""
+        with self.model_lock:
+            if self.last_used and self.pipe:
+                time_diff = datetime.now() - self.last_used
+                if time_diff > timedelta(minutes=MODEL_TIMEOUT_MINUTES):
+                    self.cleanup_model()
+                else:
+                    # اگر هنوز زمان پاکسازی نرسیده، دوباره تایمر را تنظیم کنید
+                    remaining_time = MODEL_TIMEOUT_MINUTES * 60 - time_diff.total_seconds()
+                    self.cleanup_timer = threading.Timer(remaining_time, self.check_and_cleanup)
+                    self.cleanup_timer.start()
+
+# Global model manager instance
+model_manager = ModelManager()
 
 # Supported languages for Whisper (99 languages)
 SUPPORTED_LANGUAGES = {
@@ -141,6 +231,9 @@ def download_youtube_audio(yt_url, output_path):
 def process_audio_file(file_path, task="transcribe", language="auto", return_timestamps=False):
     """Process audio file with Whisper"""
     try:
+        # دریافت مدل (با lazy loading)
+        pipe = model_manager.get_model()
+        
         # Read audio file
         with open(file_path, "rb") as f:
             inputs = f.read()
@@ -173,13 +266,53 @@ def process_audio_file(file_path, task="transcribe", language="auto", return_tim
 @app.route('/health', methods=['GET'])
 def health_check():
     """Health check endpoint"""
+    model_status = "loaded" if model_manager.pipe is not None else "not_loaded"
     return jsonify({
         "status": "healthy",
         "model": MODEL_NAME,
         "device": str(device),
+        "model_status": model_status,
+        "model_timeout_minutes": MODEL_TIMEOUT_MINUTES,
         "supported_languages": list(SUPPORTED_LANGUAGES.keys())
     })
 
+@app.route('/model/status', methods=['GET'])
+def model_status():
+    """وضعیت مدل را بررسی کنید"""
+    is_loaded = model_manager.pipe is not None
+    last_used = model_manager.last_used.isoformat() if model_manager.last_used else None
+    
+    return jsonify({
+        "model_loaded": is_loaded,
+        "last_used": last_used,
+        "timeout_minutes": MODEL_TIMEOUT_MINUTES,
+        "is_loading": model_manager.is_loading
+    })
+
+@app.route('/model/preload', methods=['POST'])
+def preload_model():
+    """پیش‌بارگذاری مدل"""
+    try:
+        model_manager.get_model()
+        return jsonify({
+            "success": True,
+            "message": "Model preloaded successfully"
+        })
+    except Exception as e:
+        return jsonify({
+            "success": False,
+            "error": str(e)
+        }), 500
+
+@app.route('/model/unload', methods=['POST'])
+def unload_model():
+    """پاکسازی دستی مدل"""
+    model_manager.cleanup_model()
+    return jsonify({
+        "success": True,
+        "message": "Model unloaded from memory"
+    })
+
 @app.route('/languages', methods=['GET'])
 def get_supported_languages():
     """Get list of supported languages"""
@@ -404,8 +537,15 @@ def get_extension_hooks():
         "description": "Extension hooks for plugins like CSS customization, myCred integration, etc."
     })
 
+# Cleanup on app shutdown
+@app.teardown_appcontext
+def cleanup_model_on_shutdown(exception):
+    """پاکسازی مدل هنگام خروج از برنامه"""
+    model_manager.cleanup_model()
+
 if __name__ == '__main__':
-    # Try to initialize model on startup
-    logger.info("Starting Flask application...")
-    initialize_model()
-    app.run(host='0.0.0.0', port=7860, debug=False)
+    try:
+        app.run(host='0.0.0.0', port=7860, debug=False)
+    finally:
+        # پاکسازی نهایی
+        model_manager.cleanup_model()