Develop

Dockerfile

@@ -28,7 +28,7 @@ COPY ./src/ /app/
 
 EXPOSE 8000
 
-COPY ./start.sh /app/start.sh
+COPY ./scripts/start.sh /app/start.sh
 RUN chmod +x /app/start.sh
 
 CMD ./start.sh

README.md

@@ -1,29 +1,55 @@
 # TTI-FastAPI
 
-Serving Text to Image Model Using FastAPI and Celery
+Serving Text to Image Model Using FastAPI, RabbitMQ and Celery worker
 
-## Installation
-1. build docker image
+## How to start
+### Using docker-compose(recommended)
+1. Clone repository
 ```shell
 git clone https://github.com/ainize-team/TTI-FastAPI
 cd TTI-FastAPI
-docker build -t tti-fast-api .
 ```
 
-2. Run Docker Image
+2. Edit [docker-compose.yml](./docker-compose.yml), [rabbitmq.env](./envs/rabbimq.env.sample) and [fastapi.env](./envs/fastapi.env.sample) for your project.
+
+3. Run containers
+```shell
+docker-compose up -d
+```
+
+4. (Optional) config rabbimq user setting
+```shell
+docker exec tti-rabbitmq -it /bin/bash
+cd scripts
+./init_rabbitmq
+```
+
+### Using docker
+1. Clone repository
+```shell
+git clone https://github.com/ainize-team/TTI-FastAPI
+cd TTI-FastAPI
+```
+
+2. Build docker image
+```shell
+docker build -t tti-fastapi .
+```
+
+3. Create docker container
 ```
 docker run -d --name tti-fastapi -p 8000:8000 \
     -e BROKER_URI=<broker_uri> \
     -e FIREBASE_DATABASE_URL=<firebase_realtime_database_url>  \
     -v <firebase_credential_dir_path>:/app/key tti-fastapi
 ```
 
-Or, you can use the [.env file](./.env.sample) to run as follows.
+Or, you can use the [.env file](./envs/fastapi.env.sample) to run as follows.
 
 ```shell
-docker build -t tti-fast-api .
+docker build -t tti-fastapi .
 docker run -d --name tti-fastapi -p 8000:8000 \
-     --env-file {.env_file_path} \
+    --env-file .env \
     -v <firebase_credential_dir_path>:/app/key tti-fastapi
 ```
 

docker-compose.yml

@@ -0,0 +1,37 @@
+version: '3'
+
+services:
+  rabbitmq:
+    image: rabbitmq:3.11.2-management
+    container_name: tti-rabbitmq
+    ports:
+      - 5672:5672
+      - 15672:15672
+    env_file:
+      - ./envs/rabbitmq.env
+    volumes:
+      - ./scripts/init_rabbitmq.sh:/scripts/init_rabbitmq.sh
+    networks:
+      - dev
+    restart: unless-stopped
+
+  api-server:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    container_name: tti-api
+    ports:
+      - 8000:8000
+    env_file:
+      - ./envs/fastapi.env
+    volumes:
+      - <firebase_credential_dir_path>:/app/key/
+    networks:
+      - dev
+    restart: unless-stopped
+    depends_on:
+     - rabbitmq
+
+networks:
+  dev:
+    driver: bridge

envs/fastapi.env.sample

@@ -0,0 +1,4 @@
+APP_NAME="Stable Diffusion FastAPI Dev Server" 
+BROKER_BASE_URI=amqp://<user_name>:<password>@<rabbitmq_container_name>:5672
+FIREBASE_DATABASE_URL=<firebase_realtime_database_url>
+FIREBASE_APP_NAME=text-to-art

envs/rabbitmq.env.sample

@@ -0,0 +1,3 @@
+RABBITMQ_DEFAULT_USER=<rabbitmq_default_user>
+RABBITMQ_DEFAULT_PASS=<rabbitmq_default_password>
+RABBITMQ_VHOSTS=vhost1,vhost2

scripts/init_rabbitmq.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+echo "Init Rabbitmq config..."
+
+USER=$RABBITMQ_DEFAULT_USER
+VHOSTS=$RABBITMQ_VHOSTS
+
+# Split the vhosts into an array using a comma (,) as the delimiter.
+IFS=',' read -ra VHOST_ARRAY <<< "$VHOSTS"
+
+echo ${VHOST_ARRAY}
+
+for vhost in "${VHOST_ARRAY[@]}"
+do
+    rabbitmqctl add_vhost $vhost
+    rabbitmqctl set_permissions -p $vhost $USER ".*" ".*" ".*"
+done

src/enums.py

@@ -22,13 +22,8 @@ class ResponseStatusEnum(StrEnum):
 
 
 class ModelEnum(StrEnum):
-    STABLE_DIFFUSION_V1_4 = "stable-diffusion-v1-4"
-    STABLE_DIFFUSION_V1_5 = "stable-diffusion-v1-5"
-    STABLE_DIFFUSION_V2 = "stable-diffusion-v2"
     STABLE_DIFFUSION_V2_1 = "stable-diffusion-v2-1"
-    STABLE_DIFFUSION_V2_768 = "stable-diffusion-v2-768"
     STABLE_DIFFUSION_V2_1_768 = "stable-diffusion-v2-1-768"
-    OPENJOURNEY_V2 = "openjourney-v2"
 
 
 class SchedulerType(StrEnum):

src/event_handlers.py

@@ -3,10 +3,11 @@
 import firebase_admin
 from celery import Celery
 from fastapi import FastAPI
-from firebase_admin import credentials, db
+from firebase_admin import credentials
 from loguru import logger
 
 from config import celery_settings, firebase_settings
+from enums import ModelEnum
 
 
 def _setup_firebase(app: FastAPI) -> None:
@@ -19,12 +20,9 @@ def _setup_firebase(app: FastAPI) -> None:
 
 def _setup_celery(app: FastAPI) -> None:
     logger.info("Setup Celery")
-    ref = db.reference(f"{firebase_settings.firebase_app_name}/models")
-    data = ref.get()
-    if data is None:
-        exit()
     app.state.celery = {}
-    for model_id in data.keys():
+    for model in ModelEnum:
+        model_id = model.value
         app.state.celery[model_id] = Celery(broker=f"{celery_settings.broker_base_uri}/{model_id}")