Spaces:
Sleeping
Sleeping
| 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() | |