Set up nos bot to save training images

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

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:

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