| import base64 | |
| import hashlib | |
| import json | |
| import time | |
| import gradio as gr | |
| 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)) | |
| def dequeue(): | |
| if not q.empty(): | |
| pr = json.loads(q.get_nowait()) | |
| return json.dumps({ | |
| "status": "ok", | |
| "prompt": pr["prompt"], | |
| "id": pr["id"] | |
| }) | |
| return json.dumps({ | |
| "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 json.dumps({"status": "ok"}) | |
| class ValueClass(): | |
| def __init__(self, value): | |
| self.value = value | |
| 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["status"] = "busy" | |
| w["prompt"] = pr["prompt"] | |
| w["id"] = pr["id"] | |
| break | |
| else: | |
| q.put(json.dumps(pr)) | |
| time.sleep(1) | |
| with gr.Blocks() as bl: | |
| with gr.Row("none"): | |
| shad_out = gr.Textbox(visible=False) | |
| shad_dequeue = gr.Button(value="Dequeue", visible=False) | |
| shad_in = gr.Textbox(visible=False) | |
| shad_complete = gr.Button(value="Complete", visible=False) | |
| shad_submitted = gr.Checkbox(visible=False) | |
| with gr.Row("panel"): | |
| gr.Label("Enter a prompt to generate an image. This is a work-in-progress. Release 1.") | |
| with gr.Row("panel"): | |
| prompt_input = gr.Textbox(placeholder="Enter a prompt here", visible=shad_complete.value) | |
| with gr.Row("panel"): | |
| submit_button = gr.Button(value="Generate", visible=True) | |
| with gr.Row("panel"): | |
| image_output = gr.Image(label="Generated Image") | |
| submit_button.click(enqueue, inputs=[prompt_input], api_name="enqueue") | |
| shad_dequeue.click(dequeue, outputs=[shad_out], api_name="dequeue") | |
| shad_complete.click(complete, inputs=[shad_in], outputs=[shad_out], api_name="complete") | |
| def kill_the_noise_skrillex(btn): | |
| prompt_input.visible = False | |
| prompt_input.value = "" | |
| submit_button.visible = False | |
| return btn | |
| prompt_input = gr.Textbox(label="Prompt") | |
| submit_button = gr.Button("Submit") | |
| demo = gr.Interface(fn=kill_the_noise_skrillex, | |
| inputs=[submit_button], | |
| outputs=[submit_button]) | |
| demo.launch() | |