From bf71e8b389a360a72b6e6c8acf6edb4b8ad78d3f Mon Sep 17 00:00:00 2001 From: Your Name Date: Wed, 24 Jun 2026 22:38:09 +0530 Subject: [PATCH] Update coreml_server.py --- scripts/workers/coreml_server.py | 127 ++++++++++++++++++++----------- 1 file changed, 83 insertions(+), 44 deletions(-) diff --git a/scripts/workers/coreml_server.py b/scripts/workers/coreml_server.py index 60be525..2589f3d 100644 --- a/scripts/workers/coreml_server.py +++ b/scripts/workers/coreml_server.py @@ -8,9 +8,20 @@ import subprocess import sys import tempfile import time +import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path +import numpy as np +import torch +from python_coreml_stable_diffusion.pipeline import ( + StableDiffusionPipeline, + StableDiffusionXLPipeline, + get_coreml_pipe, + SCHEDULER_MAP +) + + MODEL_VERSION_BY_NAME = { "coreml-stable-diffusion-v1-5": "runwayml/stable-diffusion-v1-5", @@ -20,6 +31,24 @@ MODEL_VERSION_BY_NAME = { } +def get_scheduler_class(name: str): + name = str(name).lower().replace(" ", "").replace("_", "").replace("-", "") + if "eulerancestral" in name or "eulera" in name: + return SCHEDULER_MAP.get("EulerAncestralDiscrete") + if "euler" in name: + return SCHEDULER_MAP.get("EulerDiscrete") + if "dpm" in name: + return SCHEDULER_MAP.get("DPMSolverMultistep") + if "ddim" in name: + return SCHEDULER_MAP.get("DDIM") + if "lms" in name: + return SCHEDULER_MAP.get("LMSDiscrete") + if "pndm" in name: + return SCHEDULER_MAP.get("PNDM") + return None + + + def json_response(handler, code, payload): body = json.dumps(payload).encode("utf-8") handler.send_response(code) @@ -107,6 +136,29 @@ class CoreMLServerState: self.steps = steps self.cfg_scale = cfg_scale self.started_at = time.time() + self.lock = threading.Lock() + + # Load CoreML pipeline components on startup + print(f"[coreml-npu] Loading PyTorch reference configuration for model version: {self.model_version}", flush=True) + SDP = StableDiffusionXLPipeline if 'xl' in self.model_version else StableDiffusionPipeline + pytorch_pipe = SDP.from_pretrained( + self.model_version, + use_auth_token=True, + ) + + compute_unit = os.environ.get("COREML_COMPUTE_UNIT", "CPU_AND_NE") + sources = infer_model_sources(self.resources) + + print(f"[coreml-npu] Loading Core ML models from: {self.resources} (compute unit: {compute_unit}, sources: {sources})", flush=True) + self.pipe = get_coreml_pipe( + pytorch_pipe=pytorch_pipe, + mlpackages_dir=self.resources, + model_version=self.model_version, + compute_unit=compute_unit, + sources=sources + ) + print("[coreml-npu] Core ML models loaded successfully.", flush=True) + def make_handler(state: CoreMLServerState): @@ -163,53 +215,37 @@ def make_handler(state: CoreMLServerState): negative_prompt = str(payload.get("negative_prompt") or "").strip() started = time.time() - out_dir = Path(tempfile.mkdtemp(prefix="portable-diffusion-coreml-")) - cmd = [ - sys.executable, - "-m", - "python_coreml_stable_diffusion.pipeline", - "--prompt", - prompt, - "-i", - str(state.resources), - "-o", - str(out_dir), - "--compute-unit", - os.environ.get("COREML_COMPUTE_UNIT", "CPU_AND_NE"), - "--seed", - str(seed), - "--model-version", - state.model_version, - "--num-inference-steps", - str(steps), - "--guidance-scale", - str(guidance), - "--model-sources", - infer_model_sources(state.resources), - ] - if negative_prompt: - cmd.extend(["--negative-prompt", negative_prompt]) + + # Dynamic sampler / scheduler swap + sampler = payload.get("sample_method") or payload.get("sampler") or payload.get("sampler_name") + if sampler and hasattr(state, "pipe") and state.pipe: + scheduler_cls = get_scheduler_class(sampler) + if scheduler_cls is not None and state.pipe.scheduler.__class__ != scheduler_cls: + print(f"[coreml-npu] Switching scheduler to {scheduler_cls.__name__}", flush=True) + state.pipe.scheduler = scheduler_cls.from_config(state.pipe.scheduler.config) print("[coreml-npu] generating image", flush=True) - proc = subprocess.Popen( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1, - ) - output_lines = [] - assert proc.stdout is not None - for line in proc.stdout: - output_lines.append(line.rstrip()) - print(line, end="", flush=True) - code = proc.wait() - if code != 0: - raise RuntimeError("CoreML generation failed with code %s:\n%s" % (code, "\n".join(output_lines[-30:]))) + np.random.seed(seed) + + # Thread safe execution of the in-memory pipeline + with state.lock: + output = state.pipe( + prompt=prompt, + height=state.pipe.height, + width=state.pipe.width, + num_inference_steps=steps, + guidance_scale=guidance, + negative_prompt=negative_prompt if negative_prompt else None, + ) + + image = output.images[0] + + # Convert to base64 PNG + import io + buffered = io.BytesIO() + image.save(buffered, format="PNG") + encoded = base64.b64encode(buffered.getvalue()).decode("ascii") - image_path = latest_png(out_dir) - encoded = base64.b64encode(image_path.read_bytes()).decode("ascii") - shutil.rmtree(out_dir, ignore_errors=True) print("[coreml-npu] decoding complete", flush=True) json_response(self, 200, { "created": int(time.time()), @@ -220,8 +256,11 @@ def make_handler(state: CoreMLServerState): "duration_sec": round(time.time() - started, 2), }) except Exception as exc: + import traceback + traceback.print_exc() json_response(self, 500, {"ok": False, "error": str(exc)}) + def log_message(self, fmt, *args): print("[coreml-npu-http] " + fmt % args, flush=True)