From 0d415894c7413042d3643b7155a142b57a4d9060 Mon Sep 17 00:00:00 2001 From: outtanames Date: Mon, 28 Aug 2023 15:09:27 -0700 Subject: [PATCH 1/5] Set up nos bot to save training images run bot inside a discord server alongside nos Kick off a training job on saved images --- makefiles/Makefile.base.mk | 3 ++ nos/experimental/discord/nos_bot.py | 53 ++++++++++++++++++++++++++++- scripts/entrypoint.sh | 2 +- 3 files changed, 56 insertions(+), 2 deletions(-) diff --git a/makefiles/Makefile.base.mk b/makefiles/Makefile.base.mk index 5bba0b71..31019e0f 100644 --- a/makefiles/Makefile.base.mk +++ b/makefiles/Makefile.base.mk @@ -113,3 +113,6 @@ docker-compose-upd-cpu: docker-build-cpu docker-compose-upd-gpu: docker-build-gpu docker compose -f docker-compose.gpu.yml up + +docker-compose-upd-discord-bot: docker-build-gpu + docker compose -f docker-compose.discord.yml up diff --git a/nos/experimental/discord/nos_bot.py b/nos/experimental/discord/nos_bot.py index a2b535b2..462fa0da 100755 --- a/nos/experimental/discord/nos_bot.py +++ b/nos/experimental/discord/nos_bot.py @@ -8,10 +8,10 @@ import nos from nos.client import InferenceClient, TaskType +from nos.constants import NOS_TMP_DIR # Init nos server, wait for it to spin up then confirm its healthy: -nos.init(runtime="gpu") nos_client = InferenceClient() nos_client.WaitForServer() if not nos_client.IsHealthy(): @@ -24,9 +24,13 @@ # Create our bot: bot = commands.Bot(command_prefix="$", intents=intents) +TRAINING_CHANNEL_NAME = "training" +NOS_TRAINING_DIR = NOS_TMP_DIR / "train" + # Create a callback to read messages and generate images from prompt: @bot.command() async def generate(ctx, *, prompt): + # pull the channel id so we know which model to run: response = nos_client.Run( TaskType.IMAGE_GENERATION, "stabilityai/stable-diffusion-2", @@ -44,6 +48,53 @@ async def generate(ctx, *, prompt): await ctx.send(file=discord.File(image_bytes, filename="image.png")) +@bot.command() +async def train(ctx): + # check that its in the training channel + if ctx.channel.name != TRAINING_CHANNEL_NAME: + print("not in training channel, returning!") + return + + if not ctx.message.attachments: + print("no attachments to train on, returning!") + return + + # create a thread for this training job: + thread_name = str(ctx.message.id) + thread = await ctx.channel.create_thread(name=thread_name, type=discord.ChannelType.public_thread) + + await thread.send(f"Created a new thread: {thread.name}") + + dirname = NOS_TRAINING_DIR / thread_name + dirname.mkdir(parents=True, exist_ok=True) + + await thread.send("saving at dir: " + str(dirname)) + + # save the attachments + for attachment in ctx.message.attachments: + print(f"got attachement: {attachment.filename}") + await attachment.save(os.path.join(dirname, attachment.filename)) + await thread.send(f"Image {attachment.filename} saved!") + + # Kick off a nos training run + from nos.server._service import TrainingService + + svc = TrainingService() + job_id = svc.train( + method="stable-diffusion-dreambooth-lora", + training_inputs={ + "model_name": "stabilityai/stable-diffusion-2-1", + "instance_directory": dirname, + }, + metadata={ + "name": "sdv21-dreambooth-lora-test-bench", + }, + ) + assert job_id is not None + + thread.send(f"Started training job: {job_id}") + + # Pull API token out of environment and run the bot: bot_token = os.environ.get("BOT_TOKEN") if bot_token is None: diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh index e70c162a..38090a2b 100755 --- a/scripts/entrypoint.sh +++ b/scripts/entrypoint.sh @@ -7,4 +7,4 @@ echo "Starting Ray server with OMP_NUM_THREADS=${OMP_NUM_THREADS}..." OMP_NUM_THREADS=${OMP_NUM_THREADS} ray start --head echo "Starting NOS server..." -nos-grpc-server +nos-grpc-server && python ./nos_bot.py From 1992adb0d8bcf0faf5e3432227d2d7337492c9bb Mon Sep 17 00:00:00 2001 From: outtanames Date: Tue, 29 Aug 2023 15:32:15 -0700 Subject: [PATCH 2/5] Add discord docker compose --- docker-compose.discord.yml | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) create mode 100644 docker-compose.discord.yml diff --git a/docker-compose.discord.yml b/docker-compose.discord.yml new file mode 100644 index 00000000..73936698 --- /dev/null +++ b/docker-compose.discord.yml @@ -0,0 +1,29 @@ +version: "3.8" + +services: + nos-server: + image: autonomi/nos:latest-gpu + build: + context: . + dockerfile: docker/Dockerfile.discord + args: + - TARGET=gpu + - BASE_IMAGE=nvidia/cuda:11.8.0-base-ubuntu22.04 + ports: + - 50051:50051 + - 8265:8265 + environment: + - NOS_HOME=/app/.nos + - NOS_LOGGING_LEVEL=DEBUG + volumes: + - ~/.nosd:/app/.nos + - /dev/shm:/dev/shm + ipc: host + deploy: + resources: + reservations: + devices: + - capabilities: [gpu] + limits: + cpus: "6" + memory: 6G From 2068e954ca49421bf63e1dfe11f719d8f2dc5491 Mon Sep 17 00:00:00 2001 From: outtanames Date: Tue, 29 Aug 2023 15:35:32 -0700 Subject: [PATCH 3/5] Move everything to examples --- examples/discord/Dockerfile | 16 ++++++++++++++++ .../experimental => examples}/discord/nos_bot.py | 0 examples/discord/requirements.txt | 4 ++++ 3 files changed, 20 insertions(+) create mode 100644 examples/discord/Dockerfile rename {nos/experimental => examples}/discord/nos_bot.py (100%) create mode 100644 examples/discord/requirements.txt diff --git a/examples/discord/Dockerfile b/examples/discord/Dockerfile new file mode 100644 index 00000000..28387057 --- /dev/null +++ b/examples/discord/Dockerfile @@ -0,0 +1,16 @@ +FROM python:3.8-slim + +ENV PROJECT nos-bot + +WORKDIR /tmp/$PROJECT +ADD requirements.txt . + +# Install nos client and discord dependencies +RUN pip install -r requirements.txt + +WORKDIR /app +COPY requirements.txt . +RUN pip install -r requirements.txt +COPY . . + +CMD ["python", "nos_bot.py"] \ No newline at end of file diff --git a/nos/experimental/discord/nos_bot.py b/examples/discord/nos_bot.py similarity index 100% rename from nos/experimental/discord/nos_bot.py rename to examples/discord/nos_bot.py diff --git a/examples/discord/requirements.txt b/examples/discord/requirements.txt new file mode 100644 index 00000000..dd9b5a0c --- /dev/null +++ b/examples/discord/requirements.txt @@ -0,0 +1,4 @@ +autonomi-nos +discord==2.3.2 +discord.py==2.3.2 +docker From 02f8dbc2b07a921a63d19ccd57730ea72be2da57 Mon Sep 17 00:00:00 2001 From: outtanames Date: Wed, 30 Aug 2023 11:46:09 -0700 Subject: [PATCH 4/5] Update docker compose and entrypoint to run both nos bot and server --- docker-compose.discord.yml | 2 +- scripts/entrypoint.sh | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docker-compose.discord.yml b/docker-compose.discord.yml index 73936698..5085bc6c 100644 --- a/docker-compose.discord.yml +++ b/docker-compose.discord.yml @@ -2,7 +2,7 @@ version: "3.8" services: nos-server: - image: autonomi/nos:latest-gpu + image: autonomi/nos:latest-discord build: context: . dockerfile: docker/Dockerfile.discord diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh index 38090a2b..c7dd69e1 100755 --- a/scripts/entrypoint.sh +++ b/scripts/entrypoint.sh @@ -7,4 +7,6 @@ echo "Starting Ray server with OMP_NUM_THREADS=${OMP_NUM_THREADS}..." OMP_NUM_THREADS=${OMP_NUM_THREADS} ray start --head echo "Starting NOS server..." -nos-grpc-server && python ./nos_bot.py +nos-grpc-server & +echo "Starting NOS bot..." +python ./nos_bot.py From 679e0afee9abab3e973837f252ac11741c449517 Mon Sep 17 00:00:00 2001 From: outtanames Date: Wed, 30 Aug 2023 12:09:50 -0700 Subject: [PATCH 5/5] job to thread --- examples/discord/nos_bot.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/examples/discord/nos_bot.py b/examples/discord/nos_bot.py index 462fa0da..6c08f2bf 100755 --- a/examples/discord/nos_bot.py +++ b/examples/discord/nos_bot.py @@ -27,13 +27,24 @@ TRAINING_CHANNEL_NAME = "training" NOS_TRAINING_DIR = NOS_TMP_DIR / "train" +# Init a dictionary to map job ids to thread ids: +thread_to_job = {} + # Create a callback to read messages and generate images from prompt: @bot.command() async def generate(ctx, *, prompt): - # pull the channel id so we know which model to run: + # Pull the thread id so we know which model to run generation on: + if ctx.channel.id not in thread_to_job: + await ctx.send("No thread found for this channel, please train a model first!") + return + + job_id = thread_to_job.get(ctx.channel.id) + + # TODO (Sudeep/Scott): What is the setup for training given the job id? + model_from_job_id = "custom/" + job_id response = nos_client.Run( TaskType.IMAGE_GENERATION, - "stabilityai/stable-diffusion-2", + model_from_job_id, prompts=[prompt], width=512, height=512, @@ -47,7 +58,7 @@ async def generate(ctx, *, prompt): await ctx.send(file=discord.File(image_bytes, filename="image.png")) - + @bot.command() async def train(ctx): # check that its in the training channel @@ -93,6 +104,7 @@ async def train(ctx): assert job_id is not None thread.send(f"Started training job: {job_id}") + job_to_thread[job_id] = thread # Pull API token out of environment and run the bot: