67 lines
1.8 KiB
Python
67 lines
1.8 KiB
Python
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
|
|
import requests
|
|
from rich import traceback
|
|
|
|
from bots.data import config, db, marseygen_queue
|
|
from bots.clients.drama import DramaClient
|
|
from bots.clients.runpod import RunpodClient
|
|
|
|
traceback.install()
|
|
|
|
# Set up logging
|
|
logging.basicConfig(
|
|
filename=f"{config['data_dir']}/logs/marseygen.log",
|
|
filemode="a",
|
|
format="%(name)s - %(levelname)s - %(message)s",
|
|
level=logging.INFO,
|
|
)
|
|
logger = logging.getLogger("marseygen")
|
|
|
|
|
|
def main():
|
|
if "marseygen_id" not in db:
|
|
db["marseygen_id"] = 0
|
|
if "marseygen_last_processed" not in db:
|
|
db["marseygen_last_processed"] = 0
|
|
|
|
runpod_client = RunpodClient(logger=logger)
|
|
|
|
# If requests are still in the queue and
|
|
# an instance is running, return and let it complete.
|
|
if marseygen_queue and runpod_client.is_running("marseygen"):
|
|
return
|
|
|
|
drama_client = DramaClient(config["marseygen"]["token"], logger=logger)
|
|
|
|
# Fetch new requests and add each to the queue.
|
|
comments, newest_id = drama_client.fetch_new_comments(
|
|
after=db["marseygen_last_processed"]
|
|
)
|
|
db["marseygen_last_processed"] = newest_id
|
|
|
|
# Add new requests to queue.
|
|
for comment in comments:
|
|
prompts = re.findall(
|
|
r"^!sd (.*)$", comment["body"], re.MULTILINE | re.IGNORECASE
|
|
)
|
|
prompts = prompts[:5]
|
|
for prompt in prompts:
|
|
prompt = prompt.replace("`", "")
|
|
prompt = prompt.replace("marsey", "Marsey")
|
|
db["marseygen_id"] += 1
|
|
marseygen_queue[str(db["marseygen_id"])] = (comment, prompt)
|
|
|
|
# Create instance if there are requests to be fulfilled.
|
|
if marseygen_queue:
|
|
runpod_client.create_instance(
|
|
f"{config['data_dir']}/marseygen/create_runpod.json"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|