docker-compose.yml

version: '3.9'

x-gpu-base-service: &gpu_service
  #runtime: nvidia
  privileged: true
  devices:
    - /dev/nvidia0:/dev/nvidia0
    - /dev/nvidiactl:/dev/nvidiactl
    - /dev/nvidia-caps:/dev/nvidia-caps
    - /dev/nvidia-modeset:/dev/nvidia-modeset
    - /dev/nvidia-uvm:/dev/nvidia-uvm
    - /dev/nvidia-uvm-tools:/dev/nvidia-uvm-tools
  deploy:
    resources:
      reservations:
        devices:
          - driver: nvidia
            device_ids: ['0']
            capabilities: [gpu]

x-base_service: &base_service
  user: "${UID:-0}:${GID:-0}"
  #network_mode: "host"
  ports:
    - "7860:7680"
  build:
    context: ./service
    args:
      # Compile time args
      pyver: "3.10"
      pyimage: python:3.10-slim
      #XFORMERS_COMMAND: /bin/bash /docker/install-container-dep.sh --upgrade-strategy only-if-needed  /docker/xformers-*.whl
      XFORMERS_COMMAND: /bin/bash /docker/install-container-dep.sh xformers==0.0.20
      #TORCH_COMMAND: /bin/bash /docker/install-container-dep.sh /docker/tensorflow-*.whl /docker/torch-*.whl /docker/torchvision-*.whl /docker/torchaudio-*.whl
      TORCH_COMMAND: /bin/bash /docker/install-container-dep.sh torch==2.0.1+cu118 torchvision==0.15.2+cu118 torchaudio==2.0.1+cu118
      #TENSORFLOW_COMMAND: /bin/bash /docker/install-container-dep.sh /docker/tensorflow-*.whl
      TENSORFLOW_COMMAND: /bin/bash /docker/install-container-dep.sh tensorflow==2.12.0
      # General configuration
      PIP_REPOSITORY: "https://download.pytorch.org/whl/cu118"
      PYTORCH_CUDA_ALLOC_CONF: "garbage_collection_threshold:0.9,max_split_size_mb:256"
      TRITON_VERSION: "2.0.0"
      DEEPSPEED_VERSION: "0.9.2"
      CUDNN_VERSION: "8.6.0.163"
      CUDA_VERSION: cuda-11-8
      CUDA_DRIVERS: cuda-drivers-535
      CUDA_KEYRING: https://developer.download.nvidia.com/compute/cuda/repos/debian11/x86_64/cuda-keyring_1.1-1_all.deb
      DS_BUILD_OPS: 1
      TORCH_CUDA_ARCH_LIST: 7.5+PTX
      NVCC_FLAGS: --use_fast_math
      JAX: False
      TPU: False
      DEEPSPEED: False
  volumes:
    - &v1 ./data:/data
    - &v2 /tmp/.X11-unix:/tmp/.X11-unix
  deploy:
    restart_policy:
      delay: 5s
      max_attempts: 10
      window: 120s
      
name: kohya-docker

services:

  kohya: &kohya_service
    <<: [*base_service, *gpu_service]
    profiles: ["kohya"]
    environment: 
      - TF_ENABLE_ONEDNN_OPTS=1
      - XDG_CACHE_HOME=/data/.cache
      - ACCELERATE=False
      - DISPLAY=unix$DISPLAY
      - SAFETENSORS_FAST_GPU=1
#      - RUN_ARGS=/koyah_ss/kohya_gui.py --listen 0.0.0.0 --server_port 7680
      - RUN_ARGS=/koyah_ss/gui.sh --listen 0.0.0.0 --server_port 7680
      - RUNNER=/docker/run.sh
      
  kohya_debug:
    <<: [*base_service]
    profiles: ["kohya_debug"]
    stdin_open: true
    tty: true
    environment:
      - RUNNER=/docker/debug.sh