# =============================================================================
# ACE-Step 1.5 + ace-step-ui — Full stack for Intel XPU / Arc GPUs
# =============================================================================
#
# Dual-GPU hosts (e.g. A750 + A770):
#   ZE_AFFINITY_MASK selects which Level Zero device(s) are visible.
#   On a host where [0]=A750 and [1]=A770, set ZE_AFFINITY_MASK=1 so only the
#   A770 is seen (renumbered as xpu:0).
#   Do NOT also set ONEAPI_DEVICE_SELECTOR=level_zero:1 — after affinity the
#   surviving GPU is index 0; selecting :1 yields "No XPU devices".
#
# Usage:
#   mkdir -p datasets/uploads lora_output checkpoints
#   cp .env.xpu.example .env
#   docker compose -f docker-compose.xpu.yml up -d --build
#   open http://YOUR_SERVER_IP:3003
#
# =============================================================================

services:
  acestep-xpu:
    build:
      context: .
      dockerfile: Dockerfile.xpu
    image: acestep-xpu:latest
    container_name: acestep-xpu

    devices:
      - /dev/dri:/dev/dri

    env_file:
      - path: .env
        required: false
    environment:
      - ACESTEP_MODE=${ACESTEP_MODE:-gradio-api}
      - ACESTEP_LLM_BACKEND=${ACESTEP_LLM_BACKEND:-pt}
      - ACESTEP_INIT_SERVICE=${ACESTEP_INIT_SERVICE:-true}
      - ACESTEP_CONFIG_PATH=${ACESTEP_CONFIG_PATH:-acestep-v15-turbo}
      - ACESTEP_LM_MODEL_PATH=${ACESTEP_LM_MODEL_PATH:-acestep-5Hz-lm-1.7B}
      - ACESTEP_OFFLOAD_TO_CPU=${ACESTEP_OFFLOAD_TO_CPU:-true}
      - ACESTEP_OFFLOAD_DIT_TO_CPU=${ACESTEP_OFFLOAD_DIT_TO_CPU:-true}
      - ACESTEP_LM_OFFLOAD_TO_CPU=${ACESTEP_LM_OFFLOAD_TO_CPU:-true}
      - ACESTEP_ALLOW_4B_LM=${ACESTEP_ALLOW_4B_LM:-true}
      - ACESTEP_TURBO_MAX_INFER_STEPS=${ACESTEP_TURBO_MAX_INFER_STEPS:-0}
      - ACESTEP_DTYPE=${ACESTEP_DTYPE:-}
      - ACESTEP_EXTRA_ARGS=${ACESTEP_EXTRA_ARGS:---quantization none}
      - TOKENIZERS_PARALLELISM=false
      - PYTORCH_DEVICE=xpu
      - TORCH_COMPILE_BACKEND=eager
      - SYCL_CACHE_PERSISTENT=1
      - SYCL_PI_LEVEL_ZERO_USE_IMMEDIATE_COMMANDLISTS=1
      - ZE_FLAT_DEVICE_HIERARCHY=COMPOSITE
      # Only expose A770 when host is [0]=A750 [1]=A770. Override in .env if needed.
      - ZE_AFFINITY_MASK=${ZE_AFFINITY_MASK:-1}
      - DATASETS_DIR=/app/datasets
      - LORA_OUTPUT_DIR=/app/lora_output

    ports:
      - "${API_PORT:-8001}:8001"
      - "${GRADIO_PORT:-7860}:7860"

    volumes:
      - ./checkpoints:/app/checkpoints
      - hf_cache:/root/.cache/huggingface
      - ./gradio_outputs:/app/gradio_outputs
      - ./output:/app/output
      - ./datasets:/app/datasets
      - ./lora_output:/app/lora_output

    shm_size: "2gb"

    restart: unless-stopped

    labels:
      - autoheal=true

    healthcheck:
      test: >-
        curl -sf http://localhost:8001/health > /dev/null 2>&1 ||
        curl -sf http://localhost:8001/ > /dev/null 2>&1 ||
        curl -sf http://localhost:7860/ > /dev/null 2>&1 ||
        exit 1
      interval: 30s
      timeout: 15s
      start_period: 360s
      retries: 5

  acestep-ui:
    build:
      context: .
      dockerfile: Dockerfile.ui
    image: acestep-ui:latest
    container_name: acestep-ui

    depends_on:
      acestep-xpu:
        condition: service_started

    environment:
      - ACESTEP_API_URL=http://acestep-xpu:8001
      - PORT=3001
      - FRONTEND_PORT=3003
      - DATABASE_PATH=/app/data/acestep.db
      - JWT_SECRET=${JWT_SECRET:-ace-step-ui-local-secret-change-me}
      - FRONTEND_URL=${FRONTEND_URL:-http://localhost:3003}
      - PEXELS_API_KEY=${PEXELS_API_KEY:-}
      - ENABLE_DOCKER_RESTART=${ENABLE_DOCKER_RESTART:-true}
      - ENABLE_LOG_STREAM=${ENABLE_LOG_STREAM:-true}
      - XPU_CONTAINER_NAME=${XPU_CONTAINER_NAME:-acestep-xpu}
      - UI_CONTAINER_NAME=${UI_CONTAINER_NAME:-acestep-ui}
      - DOCKER_HOST=unix:///var/run/docker.sock
      - DATASETS_DIR=/app/datasets
      - DATASETS_UPLOADS_DIR=/app/datasets/uploads
      - ACESTEP_PATH=/app
      - LORA_OUTPUT_DIR=/app/lora_output

    ports:
      - "${UI_PORT:-3003}:3003"
      - "${UI_API_PORT:-3004}:3001"

    volumes:
      - ui_data:/app/data
      - ui_audio:/app/server/public/audio
      - /var/run/docker.sock:/var/run/docker.sock
      - ./datasets:/app/datasets
      - ./lora_output:/app/lora_output
      - ./checkpoints:/app/checkpoints:ro

    restart: unless-stopped

    labels:
      - autoheal=true

    healthcheck:
      test: >-
        curl -sf http://localhost:3003/ >/dev/null 2>&1 ||
        curl -sf http://localhost:3001/ >/dev/null 2>&1 ||
        exit 1
      interval: 30s
      timeout: 10s
      start_period: 90s
      retries: 5

  autoheal:
    image: willfarrell/autoheal:latest
    container_name: acestep-autoheal
    restart: unless-stopped
    environment:
      - AUTOHEAL_CONTAINER_LABEL=autoheal
      - AUTOHEAL_INTERVAL=30
      - AUTOHEAL_START_PERIOD=360
      - AUTOHEAL_DEFAULT_STOP_TIMEOUT=20
    volumes:
      - /var/run/docker.sock:/var/run/docker.sock

volumes:
  hf_cache:
  ui_data:
  ui_audio: