Update app.py

app.py

@@ -21,6 +21,9 @@ from pydantic import BaseModel
 import asyncio
 import google.generativeai as genai
 
+# Set the TRANSFORMERS_CACHE environment variable to a writable directory
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface_cache"
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -70,7 +73,10 @@ def load_model(task: str, model_name: str = None):
         model_to_load = model_name or MODELS.get(task)
         
         if task in ["chatbot", "translation"]:
-            return genai.GenerativeModel(model_to_load)
+            logger.info(f"Initializing Gemini model: {model_to_load}")
+            model = genai.GenerativeModel(model_to_load)
+            logger.info(f"Gemini model loaded in {time.time() - start_time:.2f}s")
+            return model
             
         if task == "visual-qa":
             processor = ViltProcessor.from_pretrained(model_to_load)
@@ -91,9 +97,12 @@ def load_model(task: str, model_name: str = None):
                 logger.info(f"VQA raw output: {answer}")
                 return answer
             
+            logger.info(f"Visual QA model loaded in {time.time() - start_time:.2f}s")
             return vqa_function
         
-        return pipeline(task, model=model_to_load)
+        model = pipeline(task, model=model_to_load)
+        logger.info(f"Pipeline model loaded in {time.time() - start_time:.2f}s")
+        return model
         
     except Exception as e:
         logger.error(f"Model load failed: {str(e)}")
@@ -676,7 +685,16 @@ async def list_models():
 async def startup_event():
     """Pre-load models at startup with timeout"""
     logger.info("Starting model pre-loading...")
-    
+
+    # Load Gemini models synchronously
+    for task in ["chatbot", "translation"]:
+        try:
+            load_model(task)  # Synchronous call
+            logger.info(f"Successfully loaded {task} model")
+        except Exception as e:
+            logger.error(f"Error pre-loading {task}: {str(e)}")
+
+    # Load Hugging Face models asynchronously
     async def load_model_with_timeout(task):
         try:
             await asyncio.wait_for(load_model(task), timeout=60.0)
@@ -685,15 +703,17 @@ async def startup_event():
             logger.warning(f"Timeout loading {task} model - will load on demand")
         except Exception as e:
             logger.error(f"Error pre-loading {task}: {str(e)}")
-    
+
     await asyncio.gather(
         load_model_with_timeout("summarization"),
         load_model_with_timeout("image-to-text"),
-        load_model_with_timeout("visual-qa"),
-        load_model_with_timeout("chatbot"),
-        load_model_with_timeout("translation")
+        load_model_with_timeout("visual-qa")
     )
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
+    # Ensure the upload_dir is writable
+    logger.info(f"Checking write permissions for {upload_dir}")
+    if not os.access(upload_dir, os.W_OK):
+        logger.error(f"No write permissions for {upload_dir}")
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)