feature(shortfin_apps): move png construction to server side (#1025)

Author: bjacobgordon

shortfin/python/shortfin_apps/sd/components/TextToImageInferenceOutput.py

@@ -0,0 +1,14 @@
+# Copyright 2025 Advanced Micro Devices, Inc.
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+from PIL.Image import Image
+
+from dataclasses import dataclass
+
+
+@dataclass
+class TextToImageInferenceOutput:
+    image: Image

shortfin/python/shortfin_apps/sd/components/generate.py

@@ -9,9 +9,16 @@
 import json
 
 from typing import (
+    TypeVar,
     Union,
 )
 
+from shortfin_apps.types.Base64CharacterEncodedByteSequence import (
+    Base64CharacterEncodedByteSequence,
+)
+
+from shortfin_apps.utilities.image import png_from
+
 import shortfin as sf
 
 # TODO: Have a generic "Responder" interface vs just the concrete impl.
@@ -21,6 +28,7 @@
 from .messages import SDXLInferenceExecRequest
 from .service import SDXLGenerateService
 from .metrics import measure
+from .TextToImageInferenceOutput import TextToImageInferenceOutput
 
 logger = logging.getLogger("shortfin-sd.generate")
 
@@ -45,13 +53,38 @@ def __init__(
         self.client = client
         self.gen_req = gen_req
         self.index = index
-        self.result_image: Union[str, None] = None
+        self.output: Union[TextToImageInferenceOutput, None] = None
 
     async def run(self):
         exec = SDXLInferenceExecRequest.from_batch(self.gen_req, self.index)
         self.client.batcher.submit(exec)
         await exec.done
-        self.result_image = exec.result_image
+
+        self.output = (
+            TextToImageInferenceOutput(exec.response_image)
+            if exec.response_image
+            else None
+        )
+
+
+Item = TypeVar("Item")
+
+
+def from_batch(
+    given_subject: list[Item] | Item | None,
+    given_batch_index,
+) -> Item:
+    if given_subject is None:
+        raise Exception("Expected an item or batch of items but got `None`")
+
+    if not isinstance(given_subject, list):
+        return given_subject
+
+    # some args are broadcasted to each prompt, hence overriding index for single-item entries
+    if len(given_subject) == 1:
+        return given_subject[0]
+
+    return given_subject[given_batch_index]
 
 
 class ClientGenerateBatchProcess(sf.Process):
@@ -99,8 +132,20 @@ async def run(self):
 
             # TODO: stream image outputs
             logging.debug("Responding to one shot batch")
-            response_data = {"images": [p.result_image for p in gen_processes]}
-            json_str = json.dumps(response_data)
-            self.responder.send_response(json_str)
+
+            png_images: list[Base64CharacterEncodedByteSequence] = []
+
+            for index_of_each_process, each_process in enumerate(gen_processes):
+                if each_process.output is None:
+                    raise Exception(
+                        f"Expected output for process {index_of_each_process} but got `None`"
+                    )
+
+                each_png_image = png_from(each_process.output.image)
+                png_images.append(each_png_image)
+
+            response_body = {"images": png_images}
+            response_body_in_json = json.dumps(response_body)
+            self.responder.send_response(response_body_in_json)
         finally:
             self.responder.ensure_response()

shortfin/python/shortfin_apps/sd/components/messages.py

@@ -10,6 +10,8 @@
     Union,
 )
 
+from PIL.Image import Image
+
 import logging
 
 import shortfin as sf
@@ -90,7 +92,7 @@ def __init__(
         # Decode phase.
         self.image_array = image_array
 
-        self.result_image: Union[str, None] = None
+        self.response_image: Union[Image, None] = None
 
         self.done = sf.VoidFuture()
 

shortfin/python/shortfin_apps/sd/components/service.py

@@ -12,7 +12,6 @@
 from pathlib import Path
 from PIL import Image
 from collections import namedtuple
-import base64
 
 import shortfin as sf
 import shortfin.array as sfnp
@@ -435,10 +434,8 @@ async def _postprocess(self, device):
         # TODO: reimpl with sfnp
         permuted = np.transpose(self.exec_request.image_array, (0, 2, 3, 1))[0]
         cast_image = (permuted * 255).round().astype("uint8")
-        image_bytes = Image.fromarray(cast_image).tobytes()
-
-        image = base64.b64encode(image_bytes).decode("utf-8")
-        self.exec_request.result_image = image
+        processed_image = Image.fromarray(cast_image)
+        self.exec_request.response_image = processed_image
         return
 
 

shortfin/python/shortfin_apps/sd/simple_client.py

@@ -5,16 +5,21 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from datetime import datetime as dt
-import os
 import time
 import json
 import argparse
-import base64
 import asyncio
 import aiohttp
 import requests
 
-from PIL import Image
+from shortfin_apps.types.Base64CharacterEncodedByteSequence import (
+    Base64CharacterEncodedByteSequence,
+)
+
+from shortfin_apps.utilities.image import (
+    save_to_file,
+    image_from,
+)
 
 sample_request = {
     "prompt": [
@@ -31,30 +36,6 @@
 }
 
 
-def bytes_to_img(in_bytes, outputdir, idx=0, width=1024, height=1024):
-    timestamp = dt.now().strftime("%Y-%m-%d_%H-%M-%S")
-    image = Image.frombytes(
-        mode="RGB", size=(width, height), data=base64.b64decode(in_bytes)
-    )
-    if not os.path.isdir(outputdir):
-        os.mkdir(outputdir)
-    im_path = os.path.join(outputdir, f"shortfin_sd_output_{timestamp}_{idx}.png")
-    image.save(im_path)
-    print(f"Saved to {im_path}")
-
-
-def get_batched(request, arg, idx):
-    if isinstance(request[arg], list):
-        # some args are broadcasted to each prompt, hence overriding idx for single-item entries
-        if len(request[arg]) == 1:
-            indexed = request[arg][0]
-        else:
-            indexed = request[arg][idx]
-    else:
-        indexed = request[arg]
-    return indexed
-
-
 async def send_request(session: aiohttp.ClientSession, rep, args, data):
     print("Sending request batch #", rep)
     url = f"{args.host}:{args.port}/generate"
@@ -66,13 +47,22 @@ async def send_request(session: aiohttp.ClientSession, rep, args, data):
             response.raise_for_status()  # Raise an error for bad responses
             res_json = await response.json(content_type=None)
             if args.save:
-                for idx, item in enumerate(res_json["images"]):
-                    width = get_batched(data, "width", idx)
-                    height = get_batched(data, "height", idx)
-                    print("Saving response as image...")
-                    bytes_to_img(
-                        item.encode("utf-8"), args.outputdir, idx, width, height
+                for idx, each_png in enumerate(res_json["images"]):
+                    if not isinstance(each_png, str):
+                        raise ValueError(f"png was not string at index {idx}")
+
+                    each_image = image_from(
+                        Base64CharacterEncodedByteSequence(each_png)
+                    )
+
+                    timestamp = dt.now().strftime("%Y-%m-%d_%H-%M-%S")
+                    each_file_name = f"shortfin_sd_output_{timestamp}_{idx}.png"
+
+                    each_file_path = save_to_file(
+                        each_image, args.outputdir, each_file_name
                     )
+
+                    print(f"Saved to {each_file_path}")
             latency = end - start
             print("Responses processed.")
             return latency, len(data["prompt"])

shortfin/python/shortfin_apps/types/Base64CharacterEncodedByteSequence.py

@@ -0,0 +1,40 @@
+import re
+import binascii
+
+
+class Base64CharacterEncodedByteSequence(str):
+    """
+    A sequence of 8-bit integers encoded using the [Base64 alphabet](https://www.rfc-editor.org/rfc/rfc4648.html#section-4).
+    """
+
+    pattern = (
+        r"^(?:[A-Za-z0-9+\/]{4})*(?:[A-Za-z0-9+\/]{3}={1}|[A-Za-z0-9+\/]{2}={2})?$"
+    )
+
+    def __new__(Self, given_subject: str):
+        if re.match(Self.pattern, given_subject) == None:
+            raise ValueError("String cannot be interpreted as a byte sequence")
+
+        return super().__new__(Self, given_subject)
+
+    @property
+    def as_bytes(self):
+        base64_integer_encoded_byte_sequence = self.encode()
+
+        derived_raw_byte_sequence = binascii.a2b_base64(
+            base64_integer_encoded_byte_sequence
+        )
+
+        return derived_raw_byte_sequence
+
+    @classmethod
+    def decoded_from(Self, given_raw_byte_sequence: bytes):
+        base64_integer_encoded_byte_sequence = binascii.b2a_base64(
+            given_raw_byte_sequence, newline=False
+        )
+
+        base64_character_encoded_byte_sequence = (
+            base64_integer_encoded_byte_sequence.decode()
+        )
+
+        return Self(base64_character_encoded_byte_sequence)

shortfin/python/shortfin_apps/utilities/image.py

@@ -0,0 +1,35 @@
+import os
+
+from io import (
+    BytesIO,
+)
+
+from PIL import Image
+
+from shortfin_apps.types.Base64CharacterEncodedByteSequence import (
+    Base64CharacterEncodedByteSequence,
+)
+
+
+def save_to_file(
+    given_image: Image.Image,
+    given_directory: str,
+    given_file_name: str,
+) -> str:
+    if not os.path.isdir(given_directory):
+        os.mkdir(given_directory)
+    derived_file_path = os.path.join(given_directory, given_file_name)
+    given_image.save(derived_file_path)
+    return derived_file_path
+
+
+def png_from(given_image: Image.Image) -> Base64CharacterEncodedByteSequence:
+    memory_for_png = BytesIO()
+    given_image.save(memory_for_png, format="PNG")
+    png_from_memory = memory_for_png.getvalue()
+    return Base64CharacterEncodedByteSequence.decoded_from(png_from_memory)
+
+
+def image_from(given_png: Base64CharacterEncodedByteSequence) -> Image.Image:
+    memory_for_png = BytesIO(given_png.as_bytes)
+    return Image.open(memory_for_png, formats=["PNG"])

shortfin/tests/apps/sd/e2e_test.py

@@ -4,10 +4,8 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-import json
 import requests
 import time
-import base64
 import pytest
 import subprocess
 import os
@@ -16,7 +14,15 @@
 import copy
 from contextlib import closing
 
-from PIL import Image
+from PIL.Image import Image
+
+from shortfin_apps.types.Base64CharacterEncodedByteSequence import (
+    Base64CharacterEncodedByteSequence,
+)
+
+from shortfin_apps.utilities.image import (
+    image_from,
+)
 
 BATCH_SIZES = [1]
 
@@ -197,17 +203,10 @@ def __del__(self):
             process.wait()
 
 
-def bytes_to_img(bytes, idx=0, width=1024, height=1024):
-    image = Image.frombytes(
-        mode="RGB", size=(width, height), data=base64.b64decode(bytes)
-    )
-    return image
-
-
 def send_json_file(url="http://0.0.0.0:8000", num_copies=1):
     # Read the JSON file
     data = copy.deepcopy(sample_request)
-    imgs = []
+    imgs: list[Image] = []
     # Send the data to the /generate endpoint
     data["prompt"] = (
         [data["prompt"]]
@@ -217,30 +216,23 @@ def send_json_file(url="http://0.0.0.0:8000", num_copies=1):
     try:
         response = requests.post(url + "/generate", json=data)
         response.raise_for_status()  # Raise an error for bad responses
-        request = json.loads(response.request.body.decode("utf-8"))
+        response_body = response.json()
+
+        for idx, each_png in enumerate(response_body["images"]):
+            if not isinstance(each_png, str):
+                raise ValueError(
+                    f"Expected string-encoded png at index {idx}, found {each_png}"
+                )
 
-        for idx, item in enumerate(response.json()["images"]):
-            width = getbatched(request, idx, "width")
-            height = getbatched(request, idx, "height")
-            img = bytes_to_img(item.encode("utf-8"), idx, width, height)
-            imgs.append(img)
+            each_image = image_from(Base64CharacterEncodedByteSequence(each_png))
+            imgs.append(each_image)
 
     except requests.exceptions.RequestException as e:
         print(f"Error sending the request: {e}")
 
     return imgs, response.status_code
 
 
-def getbatched(req, idx, key):
-    if isinstance(req[key], list):
-        if len(req[key]) == 1:
-            return req[key][0]
-        elif len(req[key]) > idx:
-            return req[key][idx]
-    else:
-        return req[key]
-
-
 def find_free_port():
     """This tries to find a free port to run a server on for the test.