made more determinant in unet parsed params

lcm_server.py

@@ -1,3 +1,26 @@
+"""
+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
@@ -6,7 +29,8 @@ import queue
 import threading
 from dataclasses import dataclass
 from concurrent.futures import Future
-from typing import Optional, List, Tuple
+from typing import Optional, List, Dict, Tuple
+from contextlib import asynccontextmanager
 
 import numpy as np
 from fastapi import FastAPI, Response, HTTPException
@@ -17,34 +41,34 @@ from transformers import CLIPTokenizer
 
 from rknnlcm import RKNN2Model, RKNN2LatentConsistencyPipeline
 
-# --- Your imports (as in your script) ---
-# from your_pkg import RKNN2LatentConsistencyPipeline, RKNN2Model
-# NOTE: keep these as-is in your project.
-
 
 # -----------------------------
 # Request schema (HTTP)
 # -----------------------------
 class GenerateRequest(BaseModel):
     prompt: str
-    size: str = Field(default="512x512", pattern=r"^\d+x\d+$")
-    num_inference_steps: int = 4
-    guidance_scale: float = 1.0
-    seed: int = 1234
+    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
+@dataclass(frozen=True)
 class ModelPaths:
-    root: str  # args.i
+    root: str
+
     @property
     def scheduler_config(self) -> str:
-        return os.path.join(self.root, "scheduler/scheduler_config.json")
+        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")
@@ -57,69 +81,118 @@ class Job:
     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. Run this in a dedicated thread.
+    Owns ONE pipeline instance. Execute jobs sequentially on this worker.
     """
+
     def __init__(
         self,
         worker_id: int,
         paths: ModelPaths,
-        scheduler: LCMScheduler,
+        scheduler_config: Dict,
         tokenizer: CLIPTokenizer,
-        rknn_context_cfg: dict,
+        rknn_context_cfg: Optional[dict] = None,
+        use_rknn_context_cfgs: bool = True,
     ):
         self.worker_id = worker_id
         self.paths = paths
-        self.scheduler = scheduler
+        self.scheduler_config = scheduler_config
         self.tokenizer = tokenizer
-        self.rknn_context_cfg = rknn_context_cfg
+        self.rknn_context_cfg = rknn_context_cfg or {}
+        self.use_rknn_context_cfgs = use_rknn_context_cfgs
 
-        self.pipe = None  # built in init()
+        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: Each worker gets its *own* RKNN runtime context.
-        # You must map rknn_context_cfg to however your RKNN2Model supports it.
-        #
-        # Examples you might support in RKNN2Model:
-        #   RKNN2Model(path, core_mask=..., multi_context=True, device_id=..., ...)
-        #   RKNN2Model(path, runtime_options={...})
-        #
-        # Here: we pass **rknn_context_cfg as a flexible hook.
+        # 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=RKNN2Model(self.paths.text_encoder, **self.rknn_context_cfg),
-            unet=RKNN2Model(self.paths.unet, **self.rknn_context_cfg),
-            vae_decoder=RKNN2Model(self.paths.vae_decoder, **self.rknn_context_cfg),
-            scheduler=self.scheduler,
+            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) -> bytes:
-        h, w = (int(x) for x in job.req.size.split("x"))
+    def run_job(self, job: Job) -> Tuple[bytes, int]:
+        # Parse WIDTHxHEIGHT
+        width, height = parse_size(job.req.size)
 
-        # Deterministic per-request random generator
-        rng = np.random.RandomState(job.req.seed)
-
-        print("seed ", job.req.seed)
-        print("rng", rng)
+        # 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=h,
-            width=w,
+            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()
+        return buf.getvalue(), seed
 
 
 # -----------------------------
@@ -128,54 +201,57 @@ class PipelineWorker:
 class PipelineService:
     """
     Singleton-ish service that:
-      - loads scheduler/tokenizer once
+      - loads scheduler_config + tokenizer once
       - starts N worker threads
-      - provides a queued submit() API
+      - queues requests and runs them on worker-owned pipelines
     """
+
     _instance = None
     _instance_lock = threading.Lock()
 
     def __init__(
         self,
         paths: ModelPaths,
-        num_workers: int = 3,
-        queue_max: int = 64,
+        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 = num_workers
-        self.q: queue.Queue[Job] = queue.Queue(maxsize=queue_max)
+        self.num_workers = max(1, int(num_workers))
+        self.q: "queue.Queue[Job]" = queue.Queue(maxsize=int(queue_max))
 
-        # Load once (shared immutable objects)
+        # Load scheduler config once (immutable dict)
         with open(self.paths.scheduler_config, "r") as f:
-            scheduler_config = json.load(f)
-        self.scheduler = LCMScheduler.from_config(scheduler_config)
+            self.scheduler_config = json.load(f)
+
+        # Tokenizer is safe to share (read-only)
         self.tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch16")
 
-        # Build per-worker RKNN context configs
-        # If not provided, create N identical configs with multi_context enabled.
+        # Worker RKNN configs
         if rknn_context_cfgs is None:
-            rknn_context_cfgs = [{"multi_context": True, "worker_id": i} for i in range(num_workers)]
-        if len(rknn_context_cfgs) != num_workers:
+            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-owned pipelines
-        for i in range(num_workers):
-            worker = PipelineWorker(
+        # Create worker pipelines
+        for i in range(self.num_workers):
+            w = PipelineWorker(
                 worker_id=i,
                 paths=self.paths,
-                scheduler=self.scheduler,
+                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(worker)
+            self.workers.append(w)
 
-        # Start threads
-        for i in range(num_workers):
+        # 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)
@@ -184,9 +260,10 @@ class PipelineService:
     def get_instance(
         cls,
         paths: ModelPaths,
-        num_workers: int = 3,
-        queue_max: int = 64,
+        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:
@@ -195,12 +272,13 @@ class PipelineService:
                     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()
-        # Optionally drain queue with errors
+        # Drain queue with errors
         while True:
             try:
                 job = self.q.get_nowait()
@@ -210,10 +288,9 @@ class PipelineService:
                 job.fut.set_exception(RuntimeError("Service shutting down"))
             self.q.task_done()
 
-    def submit(self, req: GenerateRequest, timeout_s: float = 0.5) -> Future:
+    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:
@@ -233,9 +310,9 @@ class PipelineService:
                 continue
 
             try:
-                png = worker.run_job(job)
+                png, seed = worker.run_job(job)
                 if not job.fut.done():
-                    job.fut.set_result(png)
+                    job.fut.set_result((png, seed))
             except Exception as e:
                 if not job.fut.done():
                     job.fut.set_exception(e)
@@ -243,58 +320,46 @@ class PipelineService:
                 self.q.task_done()
 
 
-# -----------------------------
-# RKNN multi-context configuration
-# -----------------------------
-def build_rknn_context_cfgs_for_rk3588(num_workers: int) -> List[dict]:
-    """
-    Plug this into your RKNN2Model wrapper.
-    Typical approach on RK3588:
-      - bind each worker to a different NPU core (0/1/2)
-      - enable multi_context so each model instance has its own runtime context
-
-    You must map these fields inside RKNN2Model.
-    """
-    core_masks = ["NPU_CORE_0", "NPU_CORE_1", "NPU_CORE_2"]
-    cfgs = []
-    for i in range(num_workers):
-        cfgs.append({
-            "multi_context": True,
-            '''"core_mask":   core_masks[i % len(core_masks)],'''
-            "core_mask": "NPU_CORE_AUTO",
-            "context_name": f"w{i}",
-            "worker_id": i,
-        })
-    return cfgs
-
-
 # -----------------------------
 # FastAPI server
 # -----------------------------
-app = FastAPI()
-
-# Configure these for your deployment
 MODEL_ROOT = os.environ.get("MODEL_ROOT", "/models/lcm_rknn")
-NUM_WORKERS = int(os.environ.get("NUM_WORKERS", "3"))
+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)
 
-# Create singleton service at import time (fastest first request).
-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),
-)
+
+@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):
-    fut = service.submit(req, timeout_s=0.25)
+    service: PipelineService = app.state.service
 
+    fut = service.submit(req, timeout_s=0.25)
     try:
-        png_bytes = fut.result(timeout=120)  # you can tune this
+        png_bytes, seed = fut.result(timeout=REQUEST_TIMEOUT)
     except Exception as e:
         msg = str(e)
         if "Queue full" in msg:
@@ -306,14 +371,17 @@ def generate(req: GenerateRequest):
         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=int(os.environ.get("PORT", "4200")),
-        log_config=None,   # <-- key
+        port=PORT,
+        log_config=None,  # avoids logger dictConfig surprises
     )

rknnlcm.py

@@ -69,6 +69,15 @@ class RKNN2Model:
         self.verbose_shapes = verbose_shapes
         self.multi_context = multi_context
         self.runtime_kwargs = runtime_kwargs or {}
+        self.modelname = os.path.basename(model_dir.rstrip("/"))
+
+
+        # Known-good key orders (fallback)
+        self.key_orders = {
+            "text_encoder": ("input_ids",),
+            "unet": ("sample", "timestep", "encoder_hidden_states", "timestep_cond"),
+            "vae_decoder": ("latent_sample",),  # change to match your call
+        }        
 
         logger.info(f"Loading {model_dir}")
         start = time.time()
@@ -125,32 +134,46 @@ class RKNN2Model:
 
         raise TypeError(f"core_mask must be None, int, or str; got {type(core_mask)}")
 
-    def __call__(self, **kwargs) -> List[np.ndarray]:
-        # TODO We need deterministic ordering
-        input_list = [self._prep(v) for v in kwargs.values()]
-        results = self.rknnlite.inference(inputs=input_list, data_format=self.data_format)
-
-        logger.info("%s out[0] shape=%s dtype=%s",
-            self.modelname, results[0].shape, results[0].dtype)
+    def __call__(self, **kwargs):
+        import numpy as np
 
-        return results
+        def prep(x):
+            if isinstance(x, np.ndarray):
+                # dtype safety
+                if x.dtype == np.float64:
+                    x = x.astype(np.float32, copy=False)
+                elif x.dtype == np.float16:
+                    x = x.astype(np.float32, copy=False)
+
+                # layout safety: only transpose 4D tensors at RKNN boundary
+                if x.ndim == 4:
+                    if self.data_format == "nhwc" and x.shape[1] in (1, 3, 4):      # NCHW -> NHWC
+                        x = x.transpose(0, 2, 3, 1)
+                    elif self.data_format == "nchw" and x.shape[-1] in (1, 3, 4):   # NHWC -> NCHW
+                        x = x.transpose(0, 3, 1, 2)
+
+                x = np.ascontiguousarray(x)
+            return x
+
+        # deterministic per-model input ordering
+        if self.modelname == "text_encoder":
+            order = ("input_ids",)
+        elif self.modelname == "unet":
+            order = ("sample", "timestep", "encoder_hidden_states", "timestep_cond")
+        elif self.modelname == "vae_decoder":
+            order = ("latent_sample",)
+        else:
+            order = tuple(sorted(kwargs.keys()))
 
-    def _prep(self, x):
-        import numpy as np
-        if isinstance(x, np.ndarray):
-            # dtype safety
-            if self.force_fp32 and x.dtype in (np.float64, np.float16):
-                x = x.astype(np.float32, copy=False)
+        input_list = [prep(kwargs[k]) for k in order]
 
-            # layout safety for 4D tensors
-            if x.ndim == 4:
-                if self.data_format == "nhwc" and x.shape[1] in (1, 3, 4):  # likely NCHW
-                    x = x.transpose(0, 2, 3, 1)
-                elif self.data_format == "nchw" and x.shape[-1] in (1, 3, 4):  # likely NHWC
-                    x = x.transpose(0, 3, 1, 2)
+        if self.modelname == "vae_decoder":
+            x = input_list[0]
+            logger.info("vae in[0] shape=%s dtype=%s contiguous=%s", x.shape, x.dtype, x.flags['C_CONTIGUOUS'])
+        results = self.rknnlite.inference(inputs=input_list, data_format=self.data_format)
 
-            x = np.ascontiguousarray(x)
-        return x    
+        logger.info("%s out[0] shape=%s dtype=%s", self.modelname, results[0].shape, results[0].dtype)
+        return results
 
 class RKNN2LatentConsistencyPipeline(DiffusionPipeline):
 
@@ -554,7 +577,8 @@ class RKNN2LatentConsistencyPipeline(DiffusionPipeline):
         )
 
         # Adapted from diffusers to extend it for other runtimes than ORT
-        timestep_dtype = np.int64
+        #timestep_dtype = np.int64
+        timestep_dtype = np.int32
         
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
 
@@ -586,12 +610,18 @@ class RKNN2LatentConsistencyPipeline(DiffusionPipeline):
             image = denoised
             has_nsfw_concept = None
         else:
+            t0 = time.time()
             denoised /= self.vae_decoder.config["scaling_factor"]
-            # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
-            image = np.concatenate(
-                [self.vae_decoder(latent_sample=denoised[i : i + 1])[0] for i in range(denoised.shape[0])]
-            )
-            # image, has_nsfw_concept = self.run_safety_checker(image)
+            t1 = time.time()
+
+            t_inf0 = time.time()
+            outs = [self.vae_decoder(latent_sample=denoised[i:i+1])[0] for i in range(denoised.shape[0])]
+            t_inf1 = time.time()
+
+            t_cat0 = time.time()
+            image = np.concatenate(outs)
+            t_cat1 = time.time()
+
             has_nsfw_concept = None  # skip safety checker
 
         if has_nsfw_concept is None:
@@ -599,7 +629,12 @@ class RKNN2LatentConsistencyPipeline(DiffusionPipeline):
         else:
             do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
 
+        t_post0 = time.time()
         image = self.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+        t_post1 = time.time()
+
+        print("scale:", t1-t0, "vae_inf:", t_inf1-t_inf0, "concat:", t_cat1-t_cat0, "post:", t_post1-t_post0)
+
         decode_time = time.time() - decode_start
         print(f"Decode time: {decode_time:.2f}s")
 
@@ -672,9 +707,9 @@ def generate_png_bytes(args):
     user_specified_scheduler = LCMScheduler.from_config(scheduler_config)
 
     pipe = RKNN2LatentConsistencyPipeline(
-        text_encoder = RKNN2Model(os.path.join(args.i, "text_encoder"), data_format="nchw"),  # probably irrelevant
-        unet        = RKNN2Model(os.path.join(args.i, "unet"),        data_format="nhwc"),  # important
-        vae_decoder = RKNN2Model(os.path.join(args.i, "vae_decoder"), data_format="nhwc"),   # important        
+        text_encoder=RKNN2Model(self.paths.text_encoder, data_format="nchw", **self.rknn_context_cfg),
+        unet=RKNN2Model(self.paths.unet, data_format="nhwc", **self.rknn_context_cfg),
+        vae_decoder=RKNN2Model(self.paths.vae_decoder, data_format="nchw", **self.rknn_context_cfg),
         scheduler=user_specified_scheduler,
         tokenizer=CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch16"),
     )