Skip to content

Commit

Permalink
SDXL app: move await device to fix empty results (#975)
Browse files Browse the repository at this point in the history
The output array was coming out empty. We discovered that the `await
device` needs to be moved after the transfer to host but before the
device array's `items` is read.
  • Loading branch information
daveliddell authored Feb 17, 2025
1 parent ab42f0c commit 84a2a3a
Showing 1 changed file with 10 additions and 3 deletions.
13 changes: 10 additions & 3 deletions shortfin/python/shortfin_apps/sd/components/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -386,9 +386,11 @@ async def run(self):
await self._denoise(device=device)
if phases[InferencePhase.DECODE]["required"]:
await self._decode(device=device)
# Postprocessing needs the output data to be on the host. Even
# without postprocessing, we're done with the GPU, so we wait for
# it to finish here.
else:
# Decode and postprocess both need the output data to be on the host.
# With decode enabled, decode itself will wait for the data.
# With decode disabled, whether or not we're postprocessing,
# we're done with the GPU, so we wait for it to finish here.
await device
if phases[InferencePhase.POSTPROCESS]["required"]:
await self._postprocess(device=device)
Expand Down Expand Up @@ -546,6 +548,11 @@ async def _decode(self, device):
)
(cb.images,) = await fn(cb.latents, fiber=self.fiber)
cb.images_host.copy_from(cb.images)

# Wait for the device-to-host transfer, so that we can read the
# data with .items.
await device

image_array = cb.images_host.items
dtype = image_array.typecode
if cb.images_host.dtype == sfnp.float16:
Expand Down

0 comments on commit 84a2a3a

Please sign in to comment.