gpu work

Dockerfile

@@ -1,5 +1,4 @@
-FROM nvidia/cuda:latest
-#FROM python:3.9
+FROM python:3.9
 
 # Using the Ubuntu image (our OS)
 # Update package manager (apt-get) 

main.py

@@ -13,6 +13,8 @@ import torch
 from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 
 
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+
 app = FastAPI()
 
 @app.get("/generate")

requirements.txt

@@ -3,4 +3,6 @@ requests==2.27.*
 sentencepiece==0.1.*
 torch==1.11.*
 transformers==4.*
-uvicorn[standard]==0.17.*
+uvicorn[standard]==0.17.*
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch