""" lcn_server.py — RKNN LCM Stable Diffusion FastAPI server (queued, multi-worker safe) Key goals: - One pipeline per worker thread (no shared RKNN objects across threads) - Determin guarantee: per-request seed -> np.RandomState - Deterministic input ordering handled in RKNN2Model (recommended) - Explicit data_format per model (UNet + VAE commonly NHWC on RKNN) - Queue backpressure (429 on overflow) - Clean startup/shutdown (FastAPI lifespan) - Returns PNG bytes + X-Seed header Env: MODEL_ROOT=/models/lcm_rknn PORT=4200 NUM_WORKERS=1..3 QUEUE_MAX=64 DEFAULT_SIZE=512x512 DEFAULT_STEPS=4 DEFAULT_GUIDANCE=1.0 DEFAULT_TIMEOUT=120 """ import io import os import json import time import queue import threading from dataclasses import dataclass from concurrent.futures import Future from typing import Optional, List, Dict, Tuple from contextlib import asynccontextmanager import numpy as np from fastapi import FastAPI, Response, HTTPException from pydantic import BaseModel, Field from diffusers import LCMScheduler from transformers import CLIPTokenizer from rknnlcm import RKNN2Model, RKNN2LatentConsistencyPipeline # ----------------------------- # Request schema (HTTP) # ----------------------------- class GenerateRequest(BaseModel): prompt: str size: str = Field(default=os.environ.get("DEFAULT_SIZE", "512x512"), pattern=r"^\d+x\d+$") num_inference_steps: int = Field(default=int(os.environ.get("DEFAULT_STEPS", "4")), ge=1, le=50) guidance_scale: float = Field(default=float(os.environ.get("DEFAULT_GUIDANCE", "1.0")), ge=0.0, le=20.0) seed: Optional[int] = Field(default=None, ge=0, le=2**31 - 1) @dataclass(frozen=True) class ModelPaths: root: str @property def scheduler_config(self) -> str: return os.path.join(self.root, "scheduler", "scheduler_config.json") @property def text_encoder(self) -> str: return os.path.join(self.root, "text_encoder") @property def unet(self) -> str: return os.path.join(self.root, "unet") @property def vae_decoder(self) -> str: return os.path.join(self.root, "vae_decoder") @dataclass class Job: req: GenerateRequest fut: Future submitted_at: float # ----------------------------- # RKNN multi-context configuration # ----------------------------- def build_rknn_context_cfgs_for_rk3588(num_workers: int) -> List[dict]: """ You must map these fields inside RKNN2Model if you actually support them. If your RKNN2Model does NOT accept these kwargs, set USE_RKNN_CONTEXT_CFGS=0. """ core_masks = ["NPU_CORE_0", "NPU_CORE_1", "NPU_CORE_2"] cfgs = [] for i in range(num_workers): cfgs.append( { "multi_context": True, # binding per-core is optional; if unstable, keep AUTO "core_mask": core_masks[i % len(core_masks)], # "core_mask": "NPU_CORE_AUTO", "context_name": f"w{i}", "worker_id": i, } ) return cfgs def parse_size(size_str: str) -> Tuple[int, int]: """ Parse 'WIDTHxHEIGHT' -> (width, height) """ w_str, h_str = size_str.lower().split("x") w, h = int(w_str), int(h_str) if w <= 0 or h <= 0: raise ValueError("size must be positive") return w, h def gen_seed_8_digits() -> int: # 0..99,999,999 inclusive return int(np.random.randint(0, 100_000_000)) # ----------------------------- # Pipeline Worker # ----------------------------- class PipelineWorker: """ Owns ONE pipeline instance. Execute jobs sequentially on this worker. """ def __init__( self, worker_id: int, paths: ModelPaths, scheduler_config: Dict, tokenizer: CLIPTokenizer, rknn_context_cfg: Optional[dict] = None, use_rknn_context_cfgs: bool = True, ): self.worker_id = worker_id self.paths = paths self.scheduler_config = scheduler_config self.tokenizer = tokenizer self.rknn_context_cfg = rknn_context_cfg or {} self.use_rknn_context_cfgs = use_rknn_context_cfgs self.pipe = None self._init_pipeline() def _mk_model(self, model_path: str, *, data_format: str) -> RKNN2Model: """ Create one RKNN2Model with explicit data_format. If your RKNN2Model supports multi_context/core_mask/etc, it will receive them. """ if self.use_rknn_context_cfgs: return RKNN2Model(model_path, data_format=data_format, **self.rknn_context_cfg) return RKNN2Model(model_path, data_format=data_format) def _init_pipeline(self): # IMPORTANT: per-worker scheduler instance (avoid shared mutable state) scheduler = LCMScheduler.from_config(self.scheduler_config) # Per-model explicit formats: # - text encoder is token/embedding, format mostly irrelevant; keep nchw # - unet + vae_decoder commonly require nhwc on RKNN self.pipe = RKNN2LatentConsistencyPipeline( text_encoder=self._mk_model(self.paths.text_encoder, data_format="nchw"), unet=self._mk_model(self.paths.unet, data_format="nhwc"), vae_decoder=self._mk_model(self.paths.vae_decoder, data_format="nhwc"), scheduler=scheduler, tokenizer=self.tokenizer, ) def run_job(self, job: Job) -> Tuple[bytes, int]: # Parse WIDTHxHEIGHT width, height = parse_size(job.req.size) # Deterministic per-request RNG seed = job.req.seed if job.req.seed is not None else gen_seed_8_digits() rng = np.random.RandomState(seed) result = self.pipe( prompt=job.req.prompt, height=height, width=width, num_inference_steps=job.req.num_inference_steps, guidance_scale=job.req.guidance_scale, generator=rng, ) pil_image = result["images"][0] buf = io.BytesIO() pil_image.save(buf, format="PNG") return buf.getvalue(), seed # ----------------------------- # Singleton Service # ----------------------------- class PipelineService: """ Singleton-ish service that: - loads scheduler_config + tokenizer once - starts N worker threads - queues requests and runs them on worker-owned pipelines """ _instance = None _instance_lock = threading.Lock() def __init__( self, paths: ModelPaths, num_workers: int, queue_max: int, rknn_context_cfgs: Optional[List[dict]] = None, use_rknn_context_cfgs: bool = True, ): self.paths = paths self.num_workers = max(1, int(num_workers)) self.q: "queue.Queue[Job]" = queue.Queue(maxsize=int(queue_max)) # Load scheduler config once (immutable dict) with open(self.paths.scheduler_config, "r") as f: self.scheduler_config = json.load(f) # Tokenizer is safe to share (read-only) self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch16") # Worker RKNN configs if rknn_context_cfgs is None: rknn_context_cfgs = build_rknn_context_cfgs_for_rk3588(self.num_workers) if len(rknn_context_cfgs) != self.num_workers: raise ValueError("rknn_context_cfgs must match num_workers length") self.workers: List[PipelineWorker] = [] self.threads: List[threading.Thread] = [] self._stop = threading.Event() # Create worker pipelines for i in range(self.num_workers): w = PipelineWorker( worker_id=i, paths=self.paths, scheduler_config=self.scheduler_config, tokenizer=self.tokenizer, rknn_context_cfg=rknn_context_cfgs[i], use_rknn_context_cfgs=use_rknn_context_cfgs, ) self.workers.append(w) # Start worker threads for i in range(self.num_workers): t = threading.Thread(target=self._worker_loop, args=(i,), daemon=True) t.start() self.threads.append(t) @classmethod def get_instance( cls, paths: ModelPaths, num_workers: int, queue_max: int, rknn_context_cfgs: Optional[List[dict]] = None, use_rknn_context_cfgs: bool = True, ) -> "PipelineService": with cls._instance_lock: if cls._instance is None: cls._instance = cls( paths=paths, num_workers=num_workers, queue_max=queue_max, rknn_context_cfgs=rknn_context_cfgs, use_rknn_context_cfgs=use_rknn_context_cfgs, ) return cls._instance def shutdown(self): self._stop.set() # Drain queue with errors while True: try: job = self.q.get_nowait() except queue.Empty: break if not job.fut.done(): job.fut.set_exception(RuntimeError("Service shutting down")) self.q.task_done() def submit(self, req: GenerateRequest, timeout_s: float = 0.25) -> Future: fut: Future = Future() job = Job(req=req, fut=fut, submitted_at=time.time()) try: self.q.put(job, timeout=timeout_s) except queue.Full: fut.set_exception(RuntimeError("Queue full")) return fut def _worker_loop(self, worker_idx: int): worker = self.workers[worker_idx] while not self._stop.is_set(): try: job = self.q.get(timeout=0.1) except queue.Empty: continue if job.fut.cancelled(): self.q.task_done() continue try: png, seed = worker.run_job(job) if not job.fut.done(): job.fut.set_result((png, seed)) except Exception as e: if not job.fut.done(): job.fut.set_exception(e) finally: self.q.task_done() # ----------------------------- # FastAPI server # ----------------------------- MODEL_ROOT = os.environ.get("MODEL_ROOT", "/models/lcm_rknn") NUM_WORKERS = int(os.environ.get("NUM_WORKERS", "1")) QUEUE_MAX = int(os.environ.get("QUEUE_MAX", "64")) PORT = int(os.environ.get("PORT", "4200")) REQUEST_TIMEOUT = float(os.environ.get("DEFAULT_TIMEOUT", "120")) # If your RKNN2Model does NOT accept multi_context/core_mask kwargs, set this to 0. USE_RKNN_CONTEXT_CFGS = os.environ.get("USE_RKNN_CONTEXT_CFGS", "1") not in ("0", "false", "False") paths = ModelPaths(root=MODEL_ROOT) @asynccontextmanager async def lifespan(app: FastAPI): # Create singleton service at startup app.state.service = PipelineService.get_instance( paths=paths, num_workers=NUM_WORKERS, queue_max=QUEUE_MAX, rknn_context_cfgs=build_rknn_context_cfgs_for_rk3588(NUM_WORKERS), use_rknn_context_cfgs=USE_RKNN_CONTEXT_CFGS, ) yield # Shutdown on app stop app.state.service.shutdown() app = FastAPI(lifespan=lifespan) @app.post("/generate", responses={200: {"content": {"image/png": {}}}}) def generate(req: GenerateRequest): service: PipelineService = app.state.service fut = service.submit(req, timeout_s=0.25) try: png_bytes, seed = fut.result(timeout=REQUEST_TIMEOUT) except Exception as e: msg = str(e) if "Queue full" in msg: raise HTTPException(status_code=429, detail="Too many requests (queue full). Try again.") raise HTTPException(status_code=500, detail=f"Generation failed: {msg}") return Response( content=png_bytes, media_type="image/png", headers={ "Cache-Control": "no-store", "X-Seed": str(seed), }, ) if __name__ == "__main__": import uvicorn uvicorn.run( app, host="0.0.0.0", port=PORT, log_config=None, # avoids logger dictConfig surprises )