File size: 2,475 Bytes
54ca13d d1d2aa6 54ca13d d1d2aa6 54ca13d d1d2aa6 54ca13d d1d2aa6 54ca13d 1de5e48 54ca13d 1de5e48 54ca13d 1de5e48 54ca13d 1de5e48 54ca13d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from flask import Flask, render_template, request, redirect, url_for, jsonify, Request, Response
import json
import base64
import hashlib
import json
import time
import os
import queue
import zstandard as zstd
q = queue.Queue()
"""
Format:
{
"status": "queue" | "progress" | "done"
"prompt": "Some ducks..."
"id": "abc123"
}
"""
models = []
WORKERS = []
def enqueue(prompt: str):
tm = time.time()
hsh = hashlib.sha256(prompt.encode("utf-8")).hexdigest()
md = {
"status": "queue",
"prompt": prompt,
"id": f"{hsh}.{tm}"
}
models.append(md)
q.put(json.dumps(md))
return jsonify({
"status": "ok"
})
def dequeue():
if not q.empty():
pr = json.loads(q.get_nowait())
return jsonify({
"status": "ok",
"prompt": pr["prompt"],
"id": pr["id"]
})
return jsonify({
"status": "empty"
})
def complete(data):
jsn = json.loads(data)
for i in range(len(models)):
if models[i]["id"] == jsn["_id"]:
models[i]["status"] = "done"
for fl in jsn["files"]:
rd = zstd.decompress(base64.b64decode(fl["data"]))
os.makedirs(f"files/{fl['path']}", exist_ok=True)
os.rmdir(f"files/{fl['path']}")
with open(f"files/{fl['path']}", "wb") as f:
f.write(rd)
f.flush()
f.close()
break
return jsonify({"status": "ok"})
def worker():
while True:
if not q.empty():
pr = json.loads(q.get_nowait())
pr["status"] = "progress"
for w in WORKERS:
if w["status"] == "idle":
w["prompt"] = pr["prompt"]
w["id"] = pr["id"]
break
else:
q.put(jsonify(pr))
time.sleep(1)
app = Flask(__name__)
if __name__ == "__main__":
app.add_url_rule("/enqueue/<prompt>", "enqueue", enqueue, methods=["POST"])
app.add_url_rule("/dequeue", "dequeue", dequeue, methods=["GET"])
app.add_url_rule("/complete", "complete", complete, methods=["POST"])
app.add_url_rule("/", "index", lambda: """
<html>
<head>
<title>Snail</title>
</head>
<body>
<h1>Snail</h1>
</html>
""", methods=["GET"])
app.static_folder = "public"
app.run(port=7860, host="0.0.0.0") |