DX: Reliable imports (PYTHONPATH) + robust CUDA build hints in setup.sh

README.md

@@ -70,6 +70,38 @@ cd DiffSplat
 bash settings/setup.sh
 ```
 
+### ▶️ Optional: Containerized quickstart (Docker/RunPod)
+
+If you prefer a batteries-included container with SSH, JupyterLab, and FileBrowser pre-wired, you can use the optional files under `docker/runpod/`.
+
+Build the image (from repo root):
+
+```bash
+docker build -f docker/runpod/Dockerfile -t diffsplat-runpod .
+```
+
+Run locally with GPU (Docker Desktop + WSL2 on Windows/macOS, or native Linux):
+
+```bash
+docker run --gpus all \
+    -p 22:22 -p 80:80 -p 9090:9090 -p 9999:9999 \
+    -e JUPYTER_PASSWORD=changeme \
+    diffsplat-runpod
+```
+
+Services inside the container:
+
+- SSH: port 22 (key-based auth; set `PUBLIC_KEY` to inject your key)
+- JupyterLab: port 9999 (`--ip=0.0.0.0`, `--no-browser`; use token from logs or `JUPYTER_PASSWORD`)
+- FileBrowser: port 9090
+- Nginx: port 80 (optional welcome)
+
+On RunPod, map these ports in the UI. HTTP ports (80/9090/9999) will be proxied to `https://<pod>-<port>.proxy.runpod.net/`. SSH is a raw TCP port; connect with:
+
+```bash
+ssh -i /path/to/key -p <mapped_port> root@<runpod_host>
+```
+
 
 ## 📊 Dataset
 

scripts/infer.sh

@@ -8,6 +8,9 @@ export HF_HOME=~/.cache/huggingface
 export TORCH_HOME=~/.cache/torch
 export NCCL_DEBUG=VERSION
 
+# Ensure local imports like `import src.*` work when invoked from repo root
+export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$(pwd):$(pwd)/src"
+
 python3 ${FILE} \
     --config_file ${CONFIG_FILE} \
     --tag ${TAG} \

scripts/train.sh

@@ -14,6 +14,9 @@ export HF_HOME=~/.cache/huggingface
 export TORCH_HOME=~/.cache/torch
 export NCCL_DEBUG=VERSION
 
+# Ensure local imports like `import src.*` work when invoked from repo root
+export PYTHONPATH="${PYTHONPATH:+$PYTHONPATH:}$(pwd):$(pwd)/src"
+
 accelerate launch \
     --num_machines $NUM_MACHINES \
     --num_processes $(( $NUM_MACHINES * $NUM_LOCAL_GPUS )) \

settings/requirements.txt

@@ -1,24 +1,24 @@
-plyfile
-ipython
+plyfile==0.9
+ipython==8.26.0
 numpy==1.26.4
-matplotlib
-Pillow
-imageio
-imageio-ffmpeg
-rembg[gpu]
-lpips
-einops
-safetensors
-sentencepiece
-accelerate
-transformers
+matplotlib==3.8.4
+Pillow==10.3.0
+imageio==2.34.0
+imageio-ffmpeg==0.4.9
+rembg[gpu]==2.0.57
+lpips==0.1.4
+einops==0.7.0
+safetensors==0.4.3
+sentencepiece==0.2.0
+accelerate==0.30.1
+transformers==4.45.0
 diffusers==0.32
 kiui
-omegaconf
-pyarrow
-redis
-loguru
-pandas
+omegaconf==2.3.0
+pyarrow==15.0.2
+redis==5.0.1
+loguru==0.7.2
+pandas==2.2.2
 image-reward
 git+https://github.com/openai/CLIP.git
 deepspeed

settings/setup.sh

@@ -1,15 +1,36 @@
 PROJECT_DIR=$(pwd)
 
+# Build hints for CUDA extensions and Torch CMake paths (best-effort, safe if CUDA not present)
+export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
+export CPATH=${CPATH:-$CUDA_HOME/include:$CPATH}
+TORCH_CMAKE_PREFIX=$(python3 -c "import torch; import os; print(getattr(torch.utils, 'cmake_prefix_path', ''))" 2>/dev/null || true)
+if [ -n "$TORCH_CMAKE_PREFIX" ]; then
+	export CMAKE_PREFIX_PATH="$TORCH_CMAKE_PREFIX${CMAKE_PREFIX_PATH:+:$CMAKE_PREFIX_PATH}"
+fi
+
 # Pytorch
 pip3 install -i https://download.pytorch.org/whl/cu121 -U torch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1
 pip3 install -i https://download.pytorch.org/whl/cu121 -U xformers==0.0.27
 
 # A modified gaussian splatting (+ alpha, depth, normal rendering)
-cd extensions && git clone https://github.com/BaowenZ/RaDe-GS.git --recursive && cd RaDe-GS/submodules
-pip3 install ./diff-gaussian-rasterization
+cd extensions && git clone https://github.com/BaowenZ/RaDe-GS.git --recursive --depth 1 && cd RaDe-GS/submodules
+# Prefer current env for CUDA extensions to avoid resolver swapping
+if [ -d diff-gaussian-rasterization ]; then
+	pip3 install --no-build-isolation ./diff-gaussian-rasterization || pip3 install ./diff-gaussian-rasterization
+fi
+if [ -d simple-knn ]; then
+	pip3 install --no-build-isolation ./simple-knn || true
+fi
 cd ${PROJECT_DIR}
 
 # Others
 pip3 install -U gpustat
 pip3 install -U -r settings/requirements.txt
-sudo apt-get install -y ffmpeg
+# Install ffmpeg if available and permitted (do not fail on environments without sudo)
+if command -v apt-get >/dev/null 2>&1; then
+	if command -v sudo >/dev/null 2>&1; then
+		sudo apt-get update -y && sudo apt-get install -y ffmpeg || true
+	else
+		apt-get update -y && apt-get install -y ffmpeg || true
+	fi
+fi