Spaces:
Sleeping
Sleeping
File size: 5,438 Bytes
45769da 785e85b 45769da 785e85b 45769da 785e85b 45769da 785e85b 45769da 785e85b 45769da 0eff3a9 785e85b 0eff3a9 45769da 785e85b 45769da 785e85b 45769da 0eff3a9 45769da | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | import base64
import binascii
import json
import mimetypes
import secrets
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from io import BytesIO
from pathlib import Path
from urllib.parse import urlparse
from generate import CHECKPOINTS, MODEL_ID, MODEL_SIZES, generate_images, is_model_loaded, load_model
ROOT = Path(__file__).resolve().parent
def encode_frame(image):
buffer = BytesIO()
image.save(buffer, format="JPEG", quality=86)
return base64.b64encode(buffer.getvalue()).decode("ascii")
def generate_frames(seed, checkpoints=CHECKPOINTS, initial_image=None, model_id=MODEL_ID):
yield {"type": "status", "status": "loading", "model": model_id}
try:
for step, image in generate_images(seed, checkpoints=checkpoints, initial_image=initial_image, model_id=model_id):
yield {"type": "progress", "step": step, "total": 1000}
if image is not None:
yield {"type": "frame", "step": step, "image": encode_frame(image)}
_, _, device = load_model(model_id=model_id)
yield {"type": "complete", "seed": seed, "device": device, "model": model_id}
except Exception as error:
yield {"type": "error", "message": str(error)}
class DiffusionHandler(BaseHTTPRequestHandler):
def send_bytes(self, status, content_type, body):
self.send_response(status)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
self.wfile.write(body)
def do_GET(self):
path = urlparse(self.path).path
if path == "/api/health":
body = json.dumps({"model": MODEL_ID, "models": list(MODEL_SIZES), "loaded": is_model_loaded()}).encode()
self.send_bytes(200, "application/json", body)
return
relative_path = "index.html" if path == "/" else path.lstrip("/")
if relative_path not in {
"index.html",
"app.js",
"styles.css",
"assets/celeba_hq_validation_91_256.png",
"assets/cifar10_train_23_32.png",
}:
self.send_error(404)
return
file_path = ROOT / relative_path
content_type = mimetypes.guess_type(file_path.name)[0] or "application/octet-stream"
self.send_bytes(200, content_type, file_path.read_bytes())
def do_POST(self):
if urlparse(self.path).path != "/api/generate":
self.send_error(404)
return
try:
length = int(self.headers.get("Content-Length", "0"))
request = json.loads(self.rfile.read(length) or b"{}")
model_id = request.get("model", MODEL_ID)
if not isinstance(model_id, str) or model_id not in MODEL_SIZES:
raise ValueError("Unsupported model.")
seed = int(request.get("seed", secrets.randbelow(2**31)))
if seed < 0 or seed >= 2**63:
raise ValueError("Seed must be between 0 and 2^63 - 1.")
checkpoints = request.get("checkpoints", CHECKPOINTS)
if (
not isinstance(checkpoints, list)
or not 2 <= len(checkpoints) <= 12
or any(type(step) is not int or step < 0 or step > 1000 for step in checkpoints)
):
raise ValueError("Checkpoints must be a list of 2 to 12 integers from 0 to 1000.")
checkpoints = tuple(sorted(set(checkpoints)))
if checkpoints[0] != 0 or checkpoints[-1] != 1000:
raise ValueError("Checkpoints must include 0 and 1000.")
initial_image = None
encoded_image = request.get("initial_image")
if encoded_image is not None:
if not isinstance(encoded_image, str) or len(encoded_image) > 3_000_000:
raise ValueError("The forward-noise image is invalid or too large.")
from PIL import Image
encoded_image = encoded_image.partition(",")[2]
initial_image = Image.open(BytesIO(base64.b64decode(encoded_image, validate=True)))
initial_image.load()
except (ValueError, TypeError, json.JSONDecodeError, binascii.Error, OSError) as error:
self.send_bytes(400, "application/json", json.dumps({"error": str(error)}).encode())
return
self.send_response(200)
self.send_header("Content-Type", "application/x-ndjson; charset=utf-8")
self.send_header("Cache-Control", "no-store")
self.send_header("X-Accel-Buffering", "no")
self.end_headers()
try:
for event in generate_frames(seed, checkpoints, initial_image, model_id):
self.wfile.write((json.dumps(event) + "\n").encode())
self.wfile.flush()
except (BrokenPipeError, ConnectionResetError):
pass
def log_message(self, format_string, *args):
print(f"{self.address_string()} - {format_string % args}")
class DiffusionServer(ThreadingHTTPServer):
daemon_threads = True
if __name__ == "__main__":
server = DiffusionServer(("0.0.0.0", 8000), DiffusionHandler)
print(f"Diffusion Lab: http://{server.server_address[0]}:{server.server_address[1]}")
print(f"Model: {MODEL_ID} (loaded on first generation request)")
server.serve_forever()
|