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/", "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: """ Snail

Snail

""", methods=["GET"]) app.static_folder = "public" app.run(port=7860, host="0.0.0.0")