add cuda settings

Dockerfile

@@ -5,9 +5,11 @@ FROM python:3.9
 WORKDIR /code
 
 # Copy the current directory contents into the container at /code
+COPY ./requirements_cuda.txt /code/requirements_cuda.txt
 COPY ./requirements.txt /code/requirements.txt
 
 # Install requirements.txt 
+RUN pip install --no-cache-dir --upgrade -r /code/requirements_cuda.txt
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 # Set up a new user named "user" with user ID 1000

main.py

@@ -37,9 +37,9 @@ def draw(data: Data):
 
 		model_id = 'stabilityai/stable-diffusion-2'
 
-		#pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
-		pipe = StableDiffusionPipeline.from_pretrained(model_id)
-		#pipe = pipe.to('cuda')
+		#pipe = StableDiffusionPipeline.from_pretrained(model_id)
+		pipe = StableDiffusionPipeline.from_pretrained(model_id, revision='fp16', torch_dtype=torch.float16)
+		pipe = pipe.to('cuda')
 
 		image = pipe(prompt, negative_prompt=n_prompt).images[0]
 

requirements.txt

@@ -1,7 +1,6 @@
 fastapi==0.74.*
 requests==2.27.*
 sentencepiece==0.1.*
-torch==1.11.*
 transformers==4.*
 uvicorn[standard]==0.17.*
 diffusers[torch]==0.13.*

requirements_cuda.txt

@@ -0,0 +1,4 @@
+--extra-index-url https://download.pytorch.org/whl/cu117
+torch
+torchvision
+torchaudio