Update app.py

app.py

@@ -1,22 +1,29 @@
 from fastapi import FastAPI, Form
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
+import os
 
 app = FastAPI()
 
-# Load NLLB model once on startup
 MODEL_NAME = "facebook/nllb-200-distilled-600M"
 SRC_LANG = "mar_Deva"  # Marathi
 TGT_LANG = "eng_Latn"  # English
 
-print("Loading model... This may take a minute.")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+tokenizer = None
+model = None
 
-tokenizer.src_lang = SRC_LANG
+def load_model():
+    global tokenizer, model
+    if tokenizer is None or model is None:
+        print("Loading NLLB model... This may take a minute.")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+        tokenizer.src_lang = SRC_LANG
+        print("Model loaded successfully!")
 
 @app.post("/translate")
 async def translate(text: str = Form(...)):
+    load_model()
     inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
         generated_tokens = model.generate(
@@ -29,4 +36,9 @@ async def translate(text: str = Form(...)):
 
 @app.get("/")
 def root():
-    return {"status": "NLLB API is running!"}
+    return {"status": "NLLB API is running (model loads on first request)!"}
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)