Add GPU support

.dockerignore

@@ -0,0 +1,4 @@
+**.venv
+**.pytest_cache
+**__pycache__
+**.env

Dockerfile

@@ -18,7 +18,7 @@ RUN pip install --upgrade pip
 RUN pip install -U setuptools wheel
 
 # Install torch cpu version
-RUN pip install -U torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+RUN pip install -U torch torchvision torchaudio
 
 # Copy the requirements file into the container
 COPY --chown=user requirements.txt .

main.py

@@ -1,3 +1,4 @@
+import torch
 import zipfile
 from typing import Union
 from awq import AutoAWQForCausalLM
@@ -6,13 +7,20 @@ from tempfile import NamedTemporaryFile
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import RedirectResponse, FileResponse
-from dto import AWQConvertionRequest, GGUFConvertionRequest, GPTQConvertionRequest
+from .dto import AWQConvertionRequest, GGUFConvertionRequest, GPTQConvertionRequest
 
 ### FastAPI Initialization
 @asynccontextmanager
 async def lifespan(app:FastAPI):    
+    torch.cuda.empty_cache()
+    
+    print("Starting FastAPI server...")
+    print(f"Running on {"cuda" if torch.cuda.is_available() else "cpu"}")
+
     yield
 
+    torch.cuda.empty_cache()
+
 app = FastAPI(title="Huggingface Safetensor Model Converter to AWQ", version="0.1.0", lifespan=lifespan)
 ### -------
 
@@ -26,7 +34,7 @@ def redirect_to_docs():
 def convert(request: AWQConvertionRequest)->Union[FileResponse, dict]:
     
     try:
-        model = AutoAWQForCausalLM.from_pretrained(request.hf_model_name, device_map="cpu", trust_remote_code=True)
+        model = AutoAWQForCausalLM.from_pretrained(request.hf_model_name, trust_remote_code=True)
     except TypeError as e:
         raise HTTPException(status_code=400, detail=f"Is this model supported by AWQ Quantization? Check:https://github.com/mit-han-lab/llm-awq?tab=readme-ov-file {e}")
     

requirements.txt

@@ -8,7 +8,7 @@ pydantic
 fastapi[standard]
 transformers
 huggingface_hub[hf_xet]
-autoawq[cpu]>=0.2.8
+autoawq>=0.2.8
 pytest
 requests
 environs