Update app.py

app.py

@@ -1,15 +1,14 @@
 import os
 import uuid
 import httpx
+import torch  # <-- MISSING IMPORT ADDED
+import logging
+from typing import Dict
 from fastapi import FastAPI, Request, BackgroundTasks, HTTPException
 from fastapi.responses import JSONResponse
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import logging
 import uvicorn
-from typing import Dict
-
-# Initialize FastAPI app FIRST
-app = FastAPI()
+from contextlib import asynccontextmanager
 
 # Configuration
 MODEL_ID = "google/gemma-1.1-2b-it"
@@ -39,9 +38,10 @@ class ScriptGenerator:
             self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
-                torch_dtype=torch.float32,
+                torch_dtype=torch.float32,  # Now torch is defined
                 device_map=None,
-                token=HF_TOKEN
+                token=HF_TOKEN,
+                low_cpu_mem_usage=True
             ).to(DEVICE)
             self.loaded = True
             logger.info("Model loaded successfully")
@@ -49,6 +49,15 @@ class ScriptGenerator:
             logger.error(f"Model loading failed: {str(e)}")
             raise
 
+# Modern lifespan handler (replaces @app.on_event)
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    generator = ScriptGenerator()
+    generator.load_model()
+    yield
+    # Cleanup if needed
+
+app = FastAPI(lifespan=lifespan)
 generator = ScriptGenerator()
 
 def generate_script(topic: str) -> str:
@@ -149,10 +158,6 @@ async def get_status(job_id: str):
         raise HTTPException(status_code=404, detail="Job not found")
     return jobs[job_id]
 
-@app.on_event("startup")
-async def startup():
-    generator.load_model()
-
 if __name__ == "__main__":
     uvicorn.run(
         app,