fix cache and load model

app.py

@@ -18,9 +18,17 @@ from pydub import AudioSegment, effects
 import pandas as pd
 from moviepy.editor import VideoFileClip
 from together import Together
+import asyncio
+import logging
 
-os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/hf_cache"
-os.makedirs("/tmp/hf_cache", exist_ok=True)
+HF_CACHE_DIR = "/tmp/hf_cache"
+WHISPER_CACHE_DIR = "/tmp/whisper_cache"
+
+os.makedirs(HF_CACHE_DIR, exist_ok=True)
+os.makedirs(WHISPER_CACHE_DIR, exist_ok=True)
+
+os.environ["HUGGINGFACE_HUB_CACHE"] = HF_CACHE_DIR
+os.environ["TORCH_HOME"] = WHISPER_CACHE_DIR
 
 token = os.environ.get('HF_TOKEN')
 together_api_key = os.environ.get('TOGETHER_API_KEY')
@@ -28,25 +36,40 @@ ngrok_auth_token = os.environ.get('NGROK_AUTH_TOKEN')
 
 pipelines, models, others = [], [], []
 
-n = torch.cuda.device_count()
-
-if n == 0:
-    device = "cpu"
-    pipelines.append(Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=token).to(device))
-    models.append(whisper.load_model("large").to(device))
-elif n == 1:
-    device = "cuda:0"
-    pipelines.append(Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=token).to(device))
-    models.append(whisper.load_model("large").to(device))
-else:
-    device_pyannote = torch.device("cuda:0")
-    device_whisper = torch.device("cuda:1")
-    pipeline = Pipeline.from_pretrained(
-        "pyannote/speaker-diarization-3.1",
-        use_auth_token=token
-    )
-    pipeline.to(device_pyannote)
-    model = whisper.load_model("large").to(device_whisper)
+def load_model_bundle():
+    """โหลด pyannote + whisper โดยใช้ device ตาม GPU ที่มี"""
+    n = torch.cuda.device_count()
+    logger.info(f"🖥️ Found {n} CUDA device(s)")
+    
+    if n == 0:
+        device = "cpu"
+        pipeline = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization-3.1",
+            use_auth_token=token,
+            cache_dir=HF_CACHE_DIR
+        ).to(device)
+        model = whisper.load_model("large", download_root=WHISPER_CACHE_DIR).to(device)
+    elif n == 1:
+        device = "cuda:0"
+        pipeline = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization-3.1",
+            use_auth_token=token,
+            cache_dir=HF_CACHE_DIR
+        ).to(device)
+        model = whisper.load_model("large", download_root=WHISPER_CACHE_DIR).to(device)
+    else:
+        device_pyannote = torch.device("cuda:0")
+        device_whisper = torch.device("cuda:1")
+        pipeline = Pipeline.from_pretrained(
+            "pyannote/speaker-diarization-3.1",
+            use_auth_token=token,
+            cache_dir=HF_CACHE_DIR
+        ).to(device_pyannote)
+        model = whisper.load_model("large", download_root=WHISPER_CACHE_DIR).to(device_whisper)
+    
+    pipelines.append(pipeline)
+    models.append(model)
+
 
 nest_asyncio.apply()
 together = Together(api_key=together_api_key)
@@ -73,19 +96,16 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-@app.on_event("startup")
-def on_startup():
-    global pipeline, model, device
-    pipeline, model, device = setup_models()
-    # ... any other startup logic
-
 @app.get("/health")
-def health_check():
-    return {
-        "status": "ok",
-        "model_loaded": model is not None,
-        "diarization_pipeline_loaded": pipeline is not None,
-        "device": device
+async def health_check():
+   return {
+        "status": "running",
+        "models_loaded": {
+            "pipelines": len(pipelines),
+            "whisper_models": len(models)
+        },
+        "cuda_available": torch.cuda.is_available(),
+        "cuda_devices": torch.cuda.device_count() if torch.cuda.is_available() else 0
     }
 
 @app.get("/")
@@ -93,7 +113,7 @@ def check_api():
     return {"message": "API is up and running"}
 
 @app.get("/key")
-def check_env():
+async def check_env():
     return {
         "env": os.environ.get("ENV", "dev"),
         "openai_key_exists": bool(os.environ.get("OPENAI_API_KEY")),