services:
  model-init:
    image: ${VOICE_IMAGE}
    build: &voice-build
      context: .
      dockerfile: Containerfile
      additional_contexts:
        cosyvoice_source: ${VOICE_SOURCE_DIR}
      args:
        BASE_IMAGE: ${BASE_IMAGE}
        SOURCE_REF: ${SOURCE_REF}
        REFERENCE_WAV_PATH: ${REFERENCE_WAV_PATH}
        REFERENCE_WAV_SHA256: ${REFERENCE_WAV_SHA256}
        PIP_VERSION: ${PIP_VERSION}
        PIP_INDEX_URL: ${PIP_INDEX_URL}
        SETUPTOOLS_INSTALL_SPEC: ${SETUPTOOLS_INSTALL_SPEC}
        TORCH_VERSION_IDENTITY: ${TORCH_VERSION_IDENTITY}
        TORCH_LIBRARY_PATH: ${TORCH_LIBRARY_PATH}
        TORCHAUDIO_VERSION_IDENTITY: ${TORCHAUDIO_VERSION_IDENTITY}
        OPENAI_WHISPER_INSTALL_SPEC: ${OPENAI_WHISPER_INSTALL_SPEC}
        ONNXRUNTIME_GPU_INSTALL_SPEC: ${ONNXRUNTIME_GPU_INSTALL_SPEC}
        ONNXRUNTIME_GPU_INDEX_URL: ${ONNXRUNTIME_GPU_INDEX_URL}
    restart: "no"
    environment:
      MODEL_REPOSITORY: ${MODEL_REPOSITORY}
      MODEL_BRANCH: ${MODEL_BRANCH}
      MODEL_REF: ${MODEL_REF}
      MODEL_REQUIRE_BRANCH_HEAD_MATCH: ${MODEL_REQUIRE_BRANCH_HEAD_MATCH}
    command: [/app/model-init.sh]
    volumes:
      - ${VOICE_CACHE_DIR}/model:/model
    logging: &voice-logs
      driver: json-file
      options:
        max-size: 20m
        max-file: "3"

  voice:
    image: ${VOICE_IMAGE}
    build: *voice-build
    restart: unless-stopped
    gpus: all
    shm_size: 4gb
    deploy:
      replicas: ${VOICE_REPLICAS}
    depends_on:
      model-init:
        condition: service_completed_successfully
    ports:
      - ${VOICE_BIND_ADDRESS}:${VOICE_PORT}:${VOICE_PORT}
    environment:
      MODEL_ID: ${MODEL_ID}
      MODEL_DIR: /model
      MODEL_REF: ${MODEL_REF}
      SOURCE_REF: ${SOURCE_REF}
      PROMPT_WAV: /opt/CosyVoice/${REFERENCE_WAV_PATH}
      PROMPT_TEXT: ${REFERENCE_PROMPT_TEXT}
      REFERENCE_VOICE_ID: ${REFERENCE_VOICE_ID}
      VOICE_DEVICE: ${VOICE_DEVICE}
      VOICE_FP16: ${VOICE_FP16}
      VOICE_LOAD_TRT: ${VOICE_LOAD_TRT}
      VOICE_LOAD_VLLM: ${VOICE_LOAD_VLLM}
      VOICE_WORKERS: ${VOICE_WORKERS}
      VOICE_CONCURRENCY: ${VOICE_CONCURRENCY}
      VOICE_SPEED_MIN: ${VOICE_SPEED_MIN}
      VOICE_SPEED_MAX: ${VOICE_SPEED_MAX}
      VOICE_PORT: ${VOICE_PORT}
      PYTORCH_CUDA_ALLOC_CONF: ${PYTORCH_CUDA_ALLOC_CONF}
      CUDA_MODULE_LOADING: ${CUDA_MODULE_LOADING}
      VOICE_SPEECH_TOKENIZER_PROVIDER: ${VOICE_SPEECH_TOKENIZER_PROVIDER}
      LD_LIBRARY_PATH: ${TORCH_LIBRARY_PATH}
    volumes:
      - ${VOICE_CACHE_DIR}/model:/model:ro
      - ${VOICE_DATA_DIR}:/data
    healthcheck:
      test: [CMD, curl, --fail, --silent, "http://127.0.0.1:${VOICE_PORT}/health"]
      interval: 30s
      timeout: 10s
      retries: 20
      start_period: 10m
    logging: *voice-logs

  frpc:
    image: ${FRPC_IMAGE}
    restart: unless-stopped
    depends_on:
      voice:
        condition: service_healthy
    env_file:
      - secret.env
    command: [-c, /etc/frp/frpc.toml]
    volumes:
      - ./frpc.toml:/etc/frp/frpc.toml:ro
    network_mode: host
    logging: *voice-logs
