Update app.py

app.py

@@ -2,19 +2,18 @@ from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from modules.pmbl import PMBL
-import torch
+from llama_cpp import Llama
 
-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)}")
+# 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}")
 
 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")  # Replace with the path to your model
+pmbl = PMBL("./PMB-7b.Q6_K.gguf")  # Path to your .gguf model
 
 @app.head("/")
 @app.get("/")