services:

  # NVIDIA GB10 (Project DIGITS) — interactive shell
  cosmos3-gb10:
    build:
      context: ..
      dockerfile: docker/Dockerfile.gb10
    image: cosmos3-nano:gb10
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    ipc: host
    volumes:
      - ../Cosmos3-Nano:/workspace/cosmos3/Cosmos3-Nano:ro
      - ../Cosmos3-Nano-assets:/workspace/cosmos3/Cosmos3-Nano-assets:ro
      - ../output:/workspace/cosmos3/output
      - ../inputs:/workspace/cosmos3/inputs
      - ~/.cache/huggingface:/root/.cache/huggingface
    environment:
      - COSMOS3_DIR=/workspace/cosmos3
      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
      - TORCH_COMPILE_DISABLE=1
      - LD_LIBRARY_PATH=
      - HF_HUB_ENABLE_HF_TRANSFER=1
    stdin_open: true
    tty: true

  # NVIDIA Jetson Thor — interactive shell
  cosmos3-jetson-thor:
    build:
      context: ..
      dockerfile: docker/Dockerfile.jetson-thor
    image: cosmos3-nano:jetson-thor
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    ipc: host
    volumes:
      - ../Cosmos3-Nano:/workspace/cosmos3/Cosmos3-Nano:ro
      - ../Cosmos3-Nano-assets:/workspace/cosmos3/Cosmos3-Nano-assets:ro
      - ../output:/workspace/cosmos3/output
      - ../inputs:/workspace/cosmos3/inputs
      - ~/.cache/huggingface:/root/.cache/huggingface
    environment:
      - COSMOS3_DIR=/workspace/cosmos3
      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
      - TORCH_COMPILE_DISABLE=1
      - LD_LIBRARY_PATH=
      - HF_HUB_ENABLE_HF_TRANSFER=1
    stdin_open: true
    tty: true

  # Web UI (GB10) — change image/dockerfile for Jetson Thor
  cosmos3-webui:
    build:
      context: ..
      dockerfile: docker/Dockerfile.gb10
    image: cosmos3-nano:gb10
    runtime: nvidia
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    ipc: host
    ports:
      - "7860:7860"
    volumes:
      - ../Cosmos3-Nano:/workspace/cosmos3/Cosmos3-Nano:ro
      - ../Cosmos3-Nano-assets:/workspace/cosmos3/Cosmos3-Nano-assets:ro
      - ../output:/workspace/cosmos3/output
      - ../inputs:/workspace/cosmos3/inputs
      - ~/.cache/huggingface:/root/.cache/huggingface
    environment:
      - COSMOS3_DIR=/workspace/cosmos3
      - PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
      - TORCH_COMPILE_DISABLE=1
      - LD_LIBRARY_PATH=
      - HF_HUB_ENABLE_HF_TRANSFER=1
    entrypoint: ""
    command: /bin/bash -c "source .venv/bin/activate && python webui.py"