Update app.py

app.py

@@ -2,18 +2,19 @@ from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from modules.pmbl import PMBL
-from llama_cpp import Llama
+import torch
 
-# Check CUDA availability
-llm = Llama(model_path="./PMB-7b.Q6_K.gguf", n_gpu_layers=-1)
-print(f"CUDA available: {llm.model.context_params.n_gpu_layers > 0}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+print(f"CUDA device count: {torch.cuda.device_count()}")
+if torch.cuda.is_available():
+    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
 
 app = FastAPI(docs_url=None, redoc_url=None)
 
 app.mount("/static", StaticFiles(directory="static"), name="static")
 app.mount("/templates", StaticFiles(directory="templates"), name="templates")
 
-pmbl = PMBL("./PMB-7b.Q6_K.gguf")  # Path to your .gguf model
+pmbl = PMBL("./PMB-7b.Q6_K.gguf", gpu_layers=50)
 
 @app.head("/")
 @app.get("/")