Update app.py

app.py

@@ -11,9 +11,10 @@ import torch
 import whisper
 from moviepy.editor import VideoFileClip
 from pydub import AudioSegment
-import fitz  # PyMuPDF for handling PDFs
-import docx  # For handling .docx files
+import fitz
+import docx
 import yt_dlp
+from functools import lru_cache
 
 # Configure logging
 logging.basicConfig(
@@ -22,87 +23,120 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-# Get HuggingFace token from environment variable
-HUGGINGFACE_TOKEN = os.environ.get('HUGGINGFACE_TOKEN')
-if not HUGGINGFACE_TOKEN:
-    logger.error("HUGGINGFACE_TOKEN environment variable not set")
-    raise ValueError("Please set the HUGGINGFACE_TOKEN environment variable")
+class ModelManager:
+    _instance = None
+    
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(ModelManager, cls).__new__(cls)
+            cls._instance._initialized = False
+        return cls._instance
+    
+    def __init__(self):
+        if not self._initialized:
+            self.tokenizer = None
+            self.model = None
+            self.news_generator = None
+            self.whisper_model = None
+            self._initialized = True
+    
+    @spaces.GPU(duration=60)
+    def initialize_models(self):
+        """Initialize models with Zero GPU optimizations"""
+        try:
+            # Get HuggingFace token
+            HUGGINGFACE_TOKEN = os.environ.get('HUGGINGFACE_TOKEN')
+            if not HUGGINGFACE_TOKEN:
+                raise ValueError("HUGGINGFACE_TOKEN environment variable not set")
+
+            logger.info("Starting model initialization...")
+            model_name = "meta-llama/Llama-2-7b-chat-hf"
+
+            # Load tokenizer
+            logger.info("Loading tokenizer...")
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                model_name,
+                token=HUGGINGFACE_TOKEN,
+                use_fast=False
+            )
+            if self.tokenizer is None:
+                raise RuntimeError("Failed to initialize tokenizer")
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+            # Load model with specific GPU memory settings
+            logger.info("Loading model...")
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                token=HUGGINGFACE_TOKEN,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                low_cpu_mem_usage=True,
+                max_memory={0: "8GiB"}
+            )
+            if self.model is None:
+                raise RuntimeError("Failed to initialize model")
+
+            # Create pipeline
+            logger.info("Creating pipeline...")
+            self.news_generator = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                device_map="auto",
+                torch_dtype=torch.float16,
+                max_length=2048,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                repetition_penalty=1.2
+            )
+            if self.news_generator is None:
+                raise RuntimeError("Failed to initialize news generator pipeline")
 
-# Global variables for models
-tokenizer = None
-model = None
-news_generator = None
-whisper_model = None
+            # Load Whisper model
+            logger.info("Loading Whisper model...")
+            self.whisper_model = whisper.load_model("base", device="cuda")
+            if self.whisper_model is None:
+                raise RuntimeError("Failed to initialize Whisper model")
 
-def check_models_initialized():
-    """Check if all models are properly initialized"""
-    if None in (tokenizer, model, news_generator, whisper_model):
-        raise RuntimeError("Models not properly initialized. Please ensure initialization was successful.")
+            logger.info("All models initialized successfully")
+            return True
 
-@spaces.GPU(duration=60)
-def initialize_models():
-    """Initialize models with Zero GPU optimizations"""
-    global tokenizer, model, news_generator, whisper_model
-    
-    try:
-        logger.info("Starting model initialization...")
-        model_name = "meta-llama/Llama-2-7b-chat-hf"
-        
-        # Load tokenizer
-        logger.info("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(
-            model_name,
-            token=HUGGINGFACE_TOKEN
-        )
-        if tokenizer is None:
-            raise RuntimeError("Failed to initialize tokenizer")
-        tokenizer.pad_token = tokenizer.eos_token
-        
-        # Load model
-        logger.info("Loading model...")
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            token=HUGGINGFACE_TOKEN,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            low_cpu_mem_usage=True
-        )
-        if model is None:
-            raise RuntimeError("Failed to initialize model")
-        
-        # Create pipeline
-        logger.info("Creating pipeline...")
-        news_generator = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            device_map="auto",
-            torch_dtype=torch.bfloat16,
-            max_length=2048,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.95,
-            repetition_penalty=1.2
-        )
-        if news_generator is None:
-            raise RuntimeError("Failed to initialize news generator pipeline")
-        
-        # Load Whisper model
-        logger.info("Loading Whisper model...")
-        whisper_model = whisper.load_model("base")
-        if whisper_model is None:
-            raise RuntimeError("Failed to initialize Whisper model")
+        except Exception as e:
+            logger.error(f"Error during model initialization: {str(e)}")
+            self.reset_models()
+            raise
+
+    def reset_models(self):
+        """Reset all models to None"""
+        self.tokenizer = None
+        self.model = None
+        self.news_generator = None
+        self.whisper_model = None
         
-        logger.info("All models initialized successfully")
-        return True
-    except Exception as e:
-        logger.error(f"Error during model initialization: {str(e)}")
-        # Reset all models to None if initialization fails
-        tokenizer = None
-        model = None
-        news_generator = None
-        whisper_model = None
-        raise
+        # Clear CUDA cache
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def check_models_initialized(self):
+        """Check if all models are properly initialized"""
+        if None in (self.tokenizer, self.model, self.news_generator, self.whisper_model):
+            logger.warning("Models not initialized, attempting to initialize...")
+            self.initialize_models()
+
+    def get_models(self):
+        """Get initialized models, initializing if necessary"""
+        self.check_models_initialized()
+        return self.tokenizer, self.model, self.news_generator, self.whisper_model
+
+# Create global model manager instance
+model_manager = ModelManager()
+
+# Initialize models at startup
+try:
+    model_manager.initialize_models()
+except Exception as e:
+    logger.error(f"Initial model initialization failed: {str(e)}")
 
 def download_social_media_video(url):
     """Download a video from social media."""
@@ -154,8 +188,8 @@ def preprocess_audio(audio_file):
 def transcribe_audio(file):
     """Transcribe an audio or video file."""
     try:
-        # Check if models are initialized
-        check_models_initialized()
+        # Get initialized models
+        _, _, _, whisper_model = model_manager.get_models()
         
         if isinstance(file, str) and file.startswith('http'):
             file_path = download_social_media_video(file)
@@ -231,8 +265,8 @@ def process_social_content(url):
 @spaces.GPU(duration=60)
 def generate_news(instructions, facts, size, tone, *args):
     try:
-        # Check if models are initialized
-        check_models_initialized()
+        # Get initialized models
+        tokenizer, _, news_generator, _ = model_manager.get_models()
         
         # Initialize knowledge base
         knowledge_base = {
@@ -358,14 +392,14 @@ Follow these requirements:
 
     except Exception as e:
         logger.error(f"Error generating news: {str(e)}")
-        # Try to reinitialize models if they're not working
         try:
-            initialize_models()
+            model_manager.reset_models()
+            model_manager.initialize_models()
             logger.info("Models reinitialized successfully")
         except Exception as reinit_error:
             logger.error(f"Failed to reinitialize models: {str(reinit_error)}")
         return f"Error generating the news article: {str(e)}", ""
-
+        
 def create_demo():
     with gr.Blocks() as demo:
         gr.Markdown("## Generador de noticias todo en uno")