Update Dockerfile

Dockerfile

@@ -1,13 +1,6 @@
-FROM nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04
-#CMD nvidia-smi
+# Use an official PyTorch image with CUDA support as the base image
+FROM pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime
 
-ENV DEBIAN_FRONTEND=noninteractive 
-
-# Install Python 3.9 (or any other version you need)
-RUN apt-get install -y python3.9
-
-RUN apt-get update -y && apt-get upgrade -y && apt-get install -y sudo && apt-get install -y python3-pip && pip3 install --upgrade pip
-RUN apt-get install -y curl gnupg wget htop sudo git git-lfs software-properties-common build-essential libgl1 zip unzip
 
 # Set up a new user named "user" with user ID 1000
 RUN useradd -m -u 1000 user
@@ -24,9 +17,6 @@ WORKDIR $HOME/app
 # Clone your repository or add your code to the container
 RUN git clone -b dev https://github.com/camenduru/DiffBIR $HOME/app
 
-# Install PyTorch and torchvision with a specific CUDA toolkit version
-#RUN pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 -f https://download.pytorch.org/whl/torch_stable.html
-
 # Install Python dependencies
 RUN pip3 install -q einops pytorch_lightning gradio omegaconf xformers==0.0.20 transformers lpips opencv-python
 
@@ -58,5 +48,9 @@ RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co
 
 RUN find $HOME/app
 
+# Set the environment variable to specify the GPU device
+ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
+ENV CUDA_VISIBLE_DEVICES=0
+
 # Define the command to run your application
 CMD ["python", "gradio_diffbir.py", "--ckpt", "/home/user/app/models/general_full_v1.ckpt", "--config", "/home/user/app/configs/model/cldm.yaml", "--reload_swinir", "--swinir_ckpt", "/home/user/app/models/general_swinir_v1.ckpt"]