# ROMA inference on 4x Quadro RTX 6000 (Turing, x86_64).
# Usage (run from the repo root on the workstation):
#   docker compose -f docker/docker-rtx6000/docker-compose.yml build
#   docker compose -f docker/docker-rtx6000/docker-compose.yml run --rm --service-ports roma bash
services:
  roma:
    build:
      dockerfile: ./docker/docker-rtx6000/Dockerfile
      context: ../..
      args:
        BASE_IMAGE: nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04
        TRANSFORMERS_REF: roma_patch
        PIP_INDEX: https://pypi.org/simple
        TORCH_INDEX: https://download.pytorch.org/whl/cu124
    image: roma-rtx6000:latest
    container_name: roma-rtx6000
    volumes:
      # ~16-22GB checkpoint, HF cache and demo media live on the host, not in the image.
      - ../../whole_model:/app/whole_model
      - ../../hf_cache:/root/.cache/huggingface
      - ../../demo_media:/app/demo_media
      - ../../output:/app/output
    ports:
      - "7860:7860"
      - "8000:8000"
    environment:
      - GRADIO_SERVER_NAME=0.0.0.0
      - HF_HUB_ENABLE_HF_TRANSFER=1
      # fp16 sharded across all 4 GPUs (no quantization). See README-RTX6000.md for fallbacks.
      - ROMA_DTYPE=float16
      - ROMA_ATTN=sdpa
      - ROMA_LOAD_8BIT=0
      # NOTE: intentionally NOT setting CUDA_VISIBLE_DEVICES, so device_map="auto" shards across all 4 GPUs.
      # - HF_TOKEN=${HF_TOKEN}   # uncomment if the checkpoint pull is rate-limited/gated
    ipc: host
    shm_size: "16gb"
    tty: true
    stdin_open: true
    command: bash
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all          # all 4 Quadro RTX 6000 cards
              capabilities: [gpu]
    restart: unless-stopped