Update app.py

app.py

@@ -1,43 +1,41 @@
 from fastapi import FastAPI
-from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-MODEL_NAME = "natalieparker/LumaAI-160M-v3"
+app = FastAPI()
+
+MODEL = "natalieparker/LumaAI-160M-v3"
 
-print("🔥 Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+print("🔄 Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
 
-print("🔥 Loading model...")
+print("🔄 Loading model on CPU...")
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    MODEL,
+    torch_dtype=torch.float32,   # CPU only
     low_cpu_mem_usage=True
 )
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-
-app = FastAPI()
-
-class GenerateRequest(BaseModel):
-    prompt: str
-    max_new_tokens: int = 150
-    temperature: float = 0.9
-    top_p: float = 0.9
-
-@app.post("/api/generate")
-def generate(req: GenerateRequest):
-    inputs = tokenizer(req.prompt, return_tensors="pt").to(device)
-
-    output = model.generate(
-        **inputs,
-        max_new_tokens=req.max_new_tokens,
-        temperature=req.temperature,
-        top_p=req.top_p,
-        do_sample=True,
-        repetition_penalty=1.05,
-    )
-
+model.to("cpu")
+
+@app.get("/")
+def root():
+    return {"status": "LumaAI API is live on CPU"}
+
+@app.post("/generate")
+def generate(prompt: str):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            temperature=0.9,
+            top_p=0.9,
+            repetition_penalty=1.05,
+            do_sample=True
+        )
     text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"response": text}
+    return {"response": text}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)