flow.yml

jtype: Flow
with:
  protocol: grpc
  port: 51005
  monitoring: true  # enable prometheus & grafana
  env:
    JINA_LOG_LEVEL: debug
executors:
  - name: dalle
    uses: executors/dalle/config.yml
    timeout_ready: -1  # slow download speed often leads to timeout
    env:
      CUDA_VISIBLE_DEVICES: 0  # change this if you have multiple GPU
      XLA_PYTHON_CLIENT_ALLOCATOR: platform  # https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    replicas: 1  # change this if you have larger VRAM
  - name: diffusion
    uses: GLID3Diffusion
    uses_with:
      glid3_path: ../glid-3-xl
      steps: 100
    py_modules:
      - executors/glid3/executor.py
    env:
      CUDA_VISIBLE_DEVICES: 0  # change this if you have multiple GPU
      XLA_PYTHON_CLIENT_ALLOCATOR: platform  # https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
    replicas: 1  # change this if you have larger VRAM
    needs: [gateway]
  - name: rerank
    uses: ReRank
    uses_with:
      clip_server: grpcs://demo-cas.jina.ai:2096
    py_modules:
      - executors/rerank/executor.py
    needs: [dalle, diffusion]
  - name: upscaler
    uses: SwinIRUpscaler
    py_modules:
      - executors/swinir/executor.py
    uses_with:
      swinir_path: ../SwinIR
    env:
      CUDA_VISIBLE_DEVICES: 0  # change this if you have multiple GPU
  - name: store
    uses: MyStore
    py_modules:
      - executors/store/executor.py