fix: NHWC data_format for UNet/VAE decoder under librknnrt 2.3.2

run_rknn-lcm.py

@@ -30,9 +30,21 @@ from rknnlite.api import RKNNLite
 class RKNN2Model:
     """ Wrapper for running RKNPU2 models """
 
-    def __init__(self, model_dir):
+    def __init__(self, model_dir, data_format: str = "nchw"):
+        """
+        Args:
+            model_dir: directory containing config.json + model.rknn
+            data_format: "nchw" or "nhwc" — how the RKNN model expects
+                its 4-D inputs. Under librknnrt 2.3.2 (2025-04-09) the
+                runtime's automatic NCHW→NHWC conversion path is no
+                longer reliable for this UNet, so models that expect
+                NHWC (the UNet and the VAE decoder in this pipeline)
+                must be told explicitly and their input tensors
+                transposed in Python before the inference call.
+        """
         logger.info(f"Loading {model_dir}")
         start = time.time()
+        self.data_format = data_format.lower()
         self.config = json.load(open(os.path.join(model_dir, "config.json")))
         assert os.path.exists(model_dir) and os.path.exists(os.path.join(model_dir, "model.rknn"))
         self.rknnlite = RKNNLite()
@@ -44,15 +56,27 @@ class RKNN2Model:
         self.inference_time = 0
 
     def __call__(self, **kwargs) -> List[np.ndarray]:
-        # np.savez(f"rknn_out/{self.modelname}_input_{self.inference_time}.npz", **kwargs)
-        # self.inference_time += 1
-        #print(kwargs)
-        input_list = [value for key, value in kwargs.items()]
+        def prep(x):
+            if isinstance(x, np.ndarray):
+                # dtype safety — the runtime wants float32
+                if x.dtype in (np.float16, np.float64):
+                    x = x.astype(np.float32, copy=False)
+                # layout safety: transpose 4-D tensors to match the
+                # declared data_format at the RKNN boundary
+                if x.ndim == 4:
+                    if self.data_format == "nhwc" and x.shape[1] in (1, 3, 4):
+                        x = x.transpose(0, 2, 3, 1)  # NCHW -> NHWC
+                    elif self.data_format == "nchw" and x.shape[-1] in (1, 3, 4):
+                        x = x.transpose(0, 3, 1, 2)  # NHWC -> NCHW
+                x = np.ascontiguousarray(x)
+            return x
+
+        input_list = [prep(v) for v in kwargs.values()]
         for i, input in enumerate(input_list):
             if isinstance(input, np.ndarray):
                 print(f"input {i} shape: {input.shape}")
 
-        results = self.rknnlite.inference(inputs=input_list, data_format='nchw')
+        results = self.rknnlite.inference(inputs=input_list, data_format=self.data_format)
         for res in results:
             print(f"output shape: {res.shape}")
         return results
@@ -573,9 +597,9 @@ def main(args):
     print("user_specified_scheduler", user_specified_scheduler)
 
     pipe = RKNN2LatentConsistencyPipeline(
-        text_encoder=RKNN2Model(os.path.join(args.i, "text_encoder")),
-        unet=RKNN2Model(os.path.join(args.i, "unet")),
-        vae_decoder=RKNN2Model(os.path.join(args.i, "vae_decoder")),
+        text_encoder=RKNN2Model(os.path.join(args.i, "text_encoder"), data_format="nchw"),
+        unet=RKNN2Model(os.path.join(args.i, "unet"), data_format="nhwc"),
+        vae_decoder=RKNN2Model(os.path.join(args.i, "vae_decoder"), data_format="nhwc"),
         scheduler=user_specified_scheduler,
         tokenizer=CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch16"),
     )