diff --git a/docker-compose.discord.yml b/docker-compose.discord.yml new file mode 100644 index 00000000..5085bc6c --- /dev/null +++ b/docker-compose.discord.yml @@ -0,0 +1,29 @@ +version: "3.8" + +services: + nos-server: + image: autonomi/nos:latest-discord + 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 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/examples/discord/nos_bot.py b/examples/discord/nos_bot.py new file mode 100755 index 00000000..6c08f2bf --- /dev/null +++ b/examples/discord/nos_bot.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python + +import io +import os + +import discord +from discord.ext import commands + +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_client = InferenceClient() +nos_client.WaitForServer() +if not nos_client.IsHealthy(): + raise RuntimeError("NOS server is not healthy") + +# Set permissions for our bot to allow it to read messages: +intents = discord.Intents.default() +intents.message_content = True + +# Create our bot: +bot = commands.Bot(command_prefix="$", intents=intents) + +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 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, + model_from_job_id, + prompts=[prompt], + width=512, + height=512, + num_images=1, + ) + image = response["images"][0] + + image_bytes = io.BytesIO() + image.save(image_bytes, format="PNG") + image_bytes.seek(0) + + 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}") + job_to_thread[job_id] = thread + + +# Pull API token out of environment and run the bot: +bot_token = os.environ.get("BOT_TOKEN") +if bot_token is None: + raise Exception("BOT_TOKEN environment variable not set") + +bot.run(bot_token) 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 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 deleted file mode 100755 index a2b535b2..00000000 --- a/nos/experimental/discord/nos_bot.py +++ /dev/null @@ -1,52 +0,0 @@ -#!/usr/bin/env python - -import io -import os - -import discord -from discord.ext import commands - -import nos -from nos.client import InferenceClient, TaskType - - -# 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(): - raise RuntimeError("NOS server is not healthy") - -# Set permissions for our bot to allow it to read messages: -intents = discord.Intents.default() -intents.message_content = True - -# Create our bot: -bot = commands.Bot(command_prefix="$", intents=intents) - -# Create a callback to read messages and generate images from prompt: -@bot.command() -async def generate(ctx, *, prompt): - response = nos_client.Run( - TaskType.IMAGE_GENERATION, - "stabilityai/stable-diffusion-2", - prompts=[prompt], - width=512, - height=512, - num_images=1, - ) - image = response["images"][0] - - image_bytes = io.BytesIO() - image.save(image_bytes, format="PNG") - image_bytes.seek(0) - - await ctx.send(file=discord.File(image_bytes, filename="image.png")) - - -# Pull API token out of environment and run the bot: -bot_token = os.environ.get("BOT_TOKEN") -if bot_token is None: - raise Exception("BOT_TOKEN environment variable not set") - -bot.run(bot_token) diff --git a/scripts/entrypoint.sh b/scripts/entrypoint.sh index e70c162a..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 +nos-grpc-server & +echo "Starting NOS bot..." +python ./nos_bot.py