add stable diffusion slackbot blog

examples/stable_diffusion_slackbot/code/bot.py

@@ -0,0 +1,74 @@
+import io 
+import os
+from typing import Optional
+
+from fastapi import Request
+from modal import Image, Secret, SharedVolume, Stub, web_endpoint
+
+stub = Stub("oxen-stable-diffusion-bot")
+volume = SharedVolume().persist("stable-diff-model-vol")
+
+CACHE_PATH = "/root/model_cache"
+
+@stub.function(
+    gpu="T4",
+    image=(Image.debian_slim()
+        .pip_install("diffusers", "transformers", "scipy", "ftfy", "accelerate", "torch", "slack-sdk")),
+    shared_volumes={CACHE_PATH: volume},
+    secret=Secret.from_name("huggingface-token"),
+)
+
+async def run_stable_diffusion(prompt: str, channel_name: Optional[str] = None):
+    from diffusers import StableDiffusionPipeline
+    from torch import float16
+
+    pipe = StableDiffusionPipeline.from_pretrained(
+        "bartuso/ox_2",
+        use_auth_token=os.environ["HUGGINGFACE-TOKEN"],
+        torch_dtype=float16,
+        cache_dir=CACHE_PATH,
+        device_map="auto"
+    )
+
+    image = pipe(prompt, num_inference_steps=50).images[0]
+
+    # Convert PIL image to PNG byte array 
+    with io.BytesIO() as buf:
+        image.save(buf, format="PNG")
+        img_bytes = buf.getvalue()
+
+    if channel_name: 
+        post_image_to_slack.call(prompt, channel_name, img_bytes)
+
+    return img_bytes
+
+@stub.function()
+@web_endpoint(method="POST")
+async def entrypoint(request: Request):
+    body = await request.form()
+    prompt = body["text"]
+    run_stable_diffusion.spawn(prompt, body["channel_name"])
+    return f"Running stable diffusion for {prompt}."
+
+@stub.function(
+    image=Image.debian_slim().pip_install("slack-sdk"),
+    secret=Secret.from_name("slack-secret"),
+)
+def post_image_to_slack(title: str, channel_name: str, image_bytes: bytes):
+    import slack_sdk
+
+    client = slack_sdk.WebClient(token=os.environ["SLACK_BOT_TOKEN"])
+    client.files_upload(channels=channel_name, title=title, content=image_bytes)
+
+# Testing 
+@stub.local_entrypoint()
+def run(
+    prompt: str = "an image of the oxenai ox eating cereal",
+    output_dir: str= "/tmp/stable-diffusion"
+):
+    os.makedirs(output_dir, exist_ok=True)
+    img_bytes = run_stable_diffusion.call(prompt)
+    output_path = os.path.join(output_dir, "output.png")
+    with open(output_path, "wb") as f:
+        f.write(img_bytes)
+    print(f"Wrote data to {output_path}")

examples/stable_diffusion_slackbot/code/requirements.txt

@@ -0,0 +1,57 @@
+aiohttp==3.8.4
+aiosignal==1.3.1
+aiostream==0.4.5
+anyio==3.7.0
+asgiref==3.7.2
+async-timeout==4.0.2
+attrs==23.1.0
+certifi==2023.5.7
+charset-normalizer==3.1.0
+click==8.1.3
+cloudpickle==2.0.0
+diffusers==0.17.1
+exceptiongroup==1.1.1
+fastapi==0.97.0
+filelock==3.12.2
+frozenlist==1.3.3
+fsspec==2023.6.0
+grpclib==0.4.3
+h2==4.1.0
+hpack==4.0.0
+huggingface-hub==0.15.1
+hyperframe==6.0.1
+idna==3.4
+importlib-metadata==6.7.0
+markdown-it-py==3.0.0
+mdurl==0.1.2
+modal==0.49.2437
+modal-client==0.49.2437
+multidict==6.0.4
+numpy==1.25.0
+packaging==23.1
+Pillow==9.5.0
+protobuf==4.23.3
+pydantic==1.10.9
+Pygments==2.15.1
+PyYAML==6.0
+regex==2023.6.3
+requests==2.31.0
+rich==13.4.2
+safetensors==0.3.1
+sigtools==4.0.1
+sniffio==1.3.0
+starlette==0.27.0
+synchronicity==0.5.3
+tblib==1.7.0
+tokenizers==0.13.3
+toml==0.10.2
+tqdm==4.65.0
+transformers==4.30.2
+typer==0.9.0
+types-certifi==2021.10.8.3
+types-toml==0.10.8.6
+typing_extensions==4.6.3
+urllib3==2.0.3
+watchfiles==0.19.0
+yarl==1.9.2
+zipp==3.15.0

examples/stable_diffusion_slackbot/code/server.py

@@ -0,0 +1,189 @@
+from flask import Flask, request 
+from slack_sdk import WebClient
+from oxen import RemoteRepo
+from threading import Thread
+from slack_sdk.errors import SlackApiError
+import dotenv
+import os
+import requests
+import codecs
+import shortuuid
+import time
+import json
+from unidecode import unidecode
+import hashlib
+
+# Configure things
+dotenv.load_dotenv()
+client = WebClient(token=os.environ['SLACK_BOT_TOKEN'])
+app = Flask(__name__)
+
+# Configure Oxen repo
+repo = RemoteRepo("ba/slackbot-oxen")
+repo.checkout("dev")
+
+# Constants
+VOTES = {
+    "+1": "Approve",
+    "-1": "Disapprove",
+}
+IMAGE_DIR = "images"
+DF_PATH = "annotations/train.csv"
+MODEL_VERSION = "0-1"
+HASH_LENGTH = 12
+USER_HASH_LENGTH = 6
+
+if not os.path.exists(IMAGE_DIR):
+    os.makedirs(IMAGE_DIR)
+
+def hash_file_id(file_id):
+    """Hash slack file id into a file name"""
+    return hashlib.sha256(file_id.encode('utf-8')).hexdigest()[0:HASH_LENGTH]
+
+def commit_df_to_oxen(row):
+    start_commit = time.time()
+    try:
+        repo.add_df_row(DF_PATH, row)
+        print("time to add row", time.time() - start_commit)
+        repo.commit(f"Remote commit - {row['rater']} voting on image {row['path'].split('/')[-1]}")
+    except Exception as e:
+        print('Error adding df row to Oxen', e)    
+    print(f"Time to commit: {time.time() - start_commit}")
+
+def commit_image_to_oxen(filepath):
+    start_time_image = time.time()
+    try:
+        repo.add(filepath, IMAGE_DIR)
+        repo.commit(f"Adding image image {filepath.split('/')[-1]}")
+    except Exception as e:
+        print('Error adding image to Oxen', e)
+    print(f"Time to commit image: {time.time() - start_time_image}")
+
+def download_image(url, image_id):
+    headers = {"Authorization": f"Bearer {os.environ['SLACK_BOT_TOKEN']}"}
+    img_data = requests.get(url, headers=headers).content
+    file_hash = hash_file_id(image_id)
+    file_path = f'{IMAGE_DIR}/{file_hash}.png'
+
+    with open(file_path, 'wb') as handler:
+        handler.write(img_data)
+    return file_path
+
+def fetch_message(conversation_id, message_ts):
+    try:
+        result = client.conversations_history(
+            channel=conversation_id,
+            inclusive=True,
+            oldest = message_ts,
+            limit=1
+        )
+    except SlackApiError as e:
+        print(f"Error: {e}")
+    return result
+
+def fetch_file(file_id):
+    try:
+        file = client.files_info(
+            file=file_id,
+            count=1
+        )
+    except SlackApiError as e:
+        print(f"Error: {e}")
+    return file
+
+
+def is_valid_reaction(reaction_data):
+    if reaction_data['event']['reaction'] not in list(VOTES.keys()):
+        print("Invalid reaction, skipping")
+    return "Success"
+
+def is_valid_reaction_message(message_data):
+    if len(message_data["messages"]) < 1:
+        print("No message found, aborting")
+        return False
+    message = message_data["messages"][0]
+    if message["user"] != os.environ["SLACK_BOT_USER_ID"]:
+        print("Message not authored by bot user, aborting")
+        return False
+    # Check if the message has files 
+    if "files" not in message:
+        print("No files found, aborting")
+        return False
+    return True
+
+def is_valid_file_upload(file_data: str) -> bool:
+    if file_data['event']['user_id'] != os.environ["SLACK_BOT_USER_ID"]:
+        print("File upload not by bot user, aborting")
+        return False
+    return True
+
+
+def handle_reaction(reaction_data, message_data):
+    start_reaction = time.time()
+    file_id = message_data["files"][0]["id"]
+    file_hash = hash_file_id(file_id)
+    prompt = unidecode(message_data["files"][0]["title"])
+    oxen_row = {
+        "prompt": prompt,
+        "path": f"{IMAGE_DIR}/{file_hash}.png",
+        "rating": VOTES[reaction_data['event']['reaction']],
+        "rater": hashlib.sha256(reaction_data['event']['user'].encode("utf-8")).hexdigest()[0:USER_HASH_LENGTH],
+        "model_version": MODEL_VERSION,
+    }
+    print("Time elapsed before upload df call", time.time() - start_reaction)
+    commit_df_to_oxen(oxen_row)
+    print("Total execution time for reaction: ", time.time() - start_reaction)
+
+def handle_file_upload(file_id):
+    file = fetch_file(file_id)
+    image_url = file["file"]["url_private_download"]
+    filepath = download_image(image_url, file_id)
+    # Commit to oxen
+    commit_image_to_oxen(filepath)
+
+
+@app.route("/")
+def hello():
+    return ""
+
+# Post route to accept incoming data 
+@app.route("/post", methods=["POST"])
+def post():
+    start = time.time()
+    data = request.get_json()
+    if data['event']['type'] == 'file_created':
+        if not is_valid_file_upload(data):
+            return "Skipping"
+
+        thr = Thread(target=handle_file_upload, args=[data['event']['file_id']])
+        thr.start()
+
+        print("Total execution time for file upload: ", time.time() - start)
+        return "Success"
+
+    if data['event']['type'] == 'reaction_added':
+        if not is_valid_reaction(data):
+            return "Skipping"
+
+        conversation_id = data['event']['item']['channel']
+        message_ts = data['event']['item']['ts']
+
+        message_data = fetch_message(conversation_id, message_ts)
+
+        if not is_valid_reaction_message(message_data):
+            return "Skipping"
+
+        message = message_data["messages"][0]
+
+        thr = Thread(target=handle_reaction, args=[data, message])
+        thr.start() 
+
+        return "Success"
+
+    print("total execution time", time.time() - start)
+    return "Success"
+
+if __name__ == '__main__':
+    app.run(host="0.0.0.0", port=8000, debug=True)
+
+