# =============================================================================
# ACE-Step Forge — React workspace + API Compose stack
# =============================================================================
#
# Primary UI:  http://localhost:3000
# API:         http://localhost:8001 (useful for scripts and integrations)
# Legacy UI:   docker compose --profile legacy up acestep-gradio
#
# The Forge service proxies /api/* to the API container, so the browser never
# needs a second origin or a CORS exception.
# =============================================================================

services:
  acestep:
    build:
      context: .
      dockerfile: Dockerfile
    image: ghcr.io/sunwood-ai-labs/ace-step-forge:local
    container_name: acestep-api
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids:
                - "${ACESTEP_GPU_DEVICE_ID:-0}"
              capabilities: [gpu]
    env_file:
      - path: .env
        required: false
    environment:
      NVIDIA_VISIBLE_DEVICES: "${ACESTEP_GPU_DEVICE_ID:-0}"
      CUDA_VISIBLE_DEVICES: "${ACESTEP_GPU_DEVICE_ID:-0}"
      ACESTEP_MODE: api
      ACESTEP_LLM_BACKEND: ${ACESTEP_LLM_BACKEND:-pt}
      # The API reads ACESTEP_NO_INIT. Leave it true for a fast Forge-ready
      # server; set it false in .env to eagerly load models at startup.
      ACESTEP_NO_INIT: ${ACESTEP_NO_INIT:-true}
      ACESTEP_CONFIG_PATH: ${ACESTEP_CONFIG_PATH:-acestep-v15-turbo}
      ACESTEP_LM_MODEL_PATH: ${ACESTEP_LM_MODEL_PATH:-acestep-5Hz-lm-4B}
      ACESTEP_EXTRA_ARGS: ${ACESTEP_EXTRA_ARGS:-}
      TOKENIZERS_PARALLELISM: "false"
    ports:
      - "${API_PORT:-8001}:8001"
    volumes:
      - ./checkpoints:/app/checkpoints
      - hf_cache:/root/.cache/huggingface
      - ./gradio_outputs:/app/gradio_outputs
      - ./output:/app/output
    shm_size: "2gb"
    restart: unless-stopped
    healthcheck:
      test: ["CMD-SHELL", "curl -sf http://localhost:8001/health > /dev/null"]
      interval: 30s
      timeout: 10s
      start_period: 300s
      retries: 3

  forge:
    build:
      context: ./frontend
      dockerfile: Dockerfile
    image: ghcr.io/sunwood-ai-labs/ace-step-forge-ui:local
    container_name: acestep-forge
    depends_on:
      acestep:
        condition: service_healthy
    ports:
      - "${FORGE_PORT:-3000}:8080"
    restart: unless-stopped
    healthcheck:
      test: ["CMD-SHELL", "wget -q -O /dev/null http://127.0.0.1:8080/healthz"]
      interval: 30s
      timeout: 5s
      start_period: 10s
      retries: 3

  # A local-only Streamable HTTP MCP gateway. It delegates generation to the
  # existing API service, so completed tracks use the same shared Library.
  acestep-mcp:
    build:
      context: .
      dockerfile: Dockerfile.mcp
    image: ghcr.io/sunwood-ai-labs/ace-step-forge-mcp:local
    container_name: acestep-mcp
    depends_on:
      acestep:
        condition: service_healthy
    env_file:
      - path: .env
        required: false
    environment:
      ACESTEP_MCP_BIND_HOST: 0.0.0.0
      ACESTEP_MCP_PORT: 8002
      ACESTEP_MCP_API_BASE_URL: http://acestep:8001
      # Override this in .env when the API has a non-default public URL/port.
      ACESTEP_MCP_PUBLIC_API_BASE_URL: ${ACESTEP_MCP_PUBLIC_API_BASE_URL:-http://localhost:8001}
      ACESTEP_MCP_ALLOWED_HOSTS: ${ACESTEP_MCP_ALLOWED_HOSTS:-localhost:*,127.0.0.1:*,[::1]:*}
      ACESTEP_MCP_ALLOWED_ORIGINS: ${ACESTEP_MCP_ALLOWED_ORIGINS:-http://localhost:*,http://127.0.0.1:*,http://[::1]:*}
    # MCP is deliberately local-only. Use Tailscale Serve/ACLs if it needs to
    # be consumed by another Tailnet device.
    ports:
      - "127.0.0.1:${MCP_PORT:-8002}:8002"
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8002/health', timeout=2).close()"]
      interval: 30s
      timeout: 5s
      start_period: 10s
      retries: 3

  # The original bundled Gradio UI is preserved, but is opt-in so it does not
  # compete with the API model worker used by Forge.
  acestep-gradio:
    profiles: ["legacy"]
    build:
      context: .
      dockerfile: Dockerfile
    image: ghcr.io/sunwood-ai-labs/ace-step-forge:local
    container_name: acestep-gradio
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids:
                - "${ACESTEP_GPU_DEVICE_ID:-0}"
              capabilities: [gpu]
    env_file:
      - path: .env
        required: false
    environment:
      NVIDIA_VISIBLE_DEVICES: "${ACESTEP_GPU_DEVICE_ID:-0}"
      CUDA_VISIBLE_DEVICES: "${ACESTEP_GPU_DEVICE_ID:-0}"
      ACESTEP_MODE: gradio
      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-4B}
      ACESTEP_EXTRA_ARGS: ${ACESTEP_EXTRA_ARGS:-}
      TOKENIZERS_PARALLELISM: "false"
    ports:
      - "${GRADIO_PORT:-7860}:7860"
    volumes:
      - ./checkpoints:/app/checkpoints
      - hf_cache:/root/.cache/huggingface
      - ./gradio_outputs:/app/gradio_outputs
      - ./output:/app/output
    shm_size: "2gb"
    restart: unless-stopped
    healthcheck:
      test: ["CMD-SHELL", "curl -sf http://localhost:7860/ > /dev/null"]
      interval: 30s
      timeout: 10s
      start_period: 300s
      retries: 3

volumes:
  hf_cache:

# Docker Desktop can exhaust its automatic bridge-network pool when many local
# projects coexist. Keep this stack deterministic while allowing an override in
# .env when 172.16.43.0/24 is already routed in a particular environment.
networks:
  default:
    name: ace-step-forge-network
    ipam:
      config:
        - subnet: ${ACESTEP_DOCKER_SUBNET:-172.16.43.0/24}