From 89e5e91f40a3d75cd0ec989fcf52f9f98e1d2c18 Mon Sep 17 00:00:00 2001 From: Steve Tietze Date: Wed, 17 Dec 2025 22:09:11 +0700 Subject: [PATCH] fix: define rgb_bytes before torch.frombuffer in image decode --- kie_api/nanobanana.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/kie_api/nanobanana.py b/kie_api/nanobanana.py index eb9e1e3..5902d98 100644 --- a/kie_api/nanobanana.py +++ b/kie_api/nanobanana.py @@ -268,8 +268,14 @@ def _image_bytes_to_tensor(image_bytes: bytes) -> torch.Tensor: try: with Image.open(BytesIO(image_bytes)) as img: rgb_image = img.convert("RGB") + rgb_bytes = rgb_image.tobytes() + tensor = torch.frombuffer(rgb_bytes, dtype=torch.uint8).clone() - tensor = tensor.view(rgb_image.height, rgb_image.width, 3).float() / 255.0 + tensor = tensor.view( + rgb_image.height, rgb_image.width, 3 + ).float() / 255.0 + return tensor.unsqueeze(0) except Exception as exc: raise RuntimeError("Failed to decode result image.") from exc +