Skip to content

Stream progress and a finished image separately

Image generation produces two results with different lifecycles: progress while the model runs and one image when it finishes. Giving each result its own output port lets a caller drain them concurrently, apply different size limits, and display either one without decoding a mixed event stream.

The same port structure supports video previews and completed assets, speech frames and alignment records, or image masks and rendered output. Each reader subscribes only to its required result.

Before you start

The demo below draws against a hosted Stable Diffusion 1.5 backend. The hosted demo needs no local packages.

To run the action yourself, the backend machine needs diffusers, transformers, torch and a checkpoint. The first run downloads the checkpoint:

pip install 'a11-kit[diffusion]'
python -m a11.demos.web_demos_server

A page loaded over HTTPS may refuse a plaintext ws:// socket even to localhost (Chrome allows it, Firefox does not), so give a local backend the --certificate / --private-key flags and a trusted certificate — mkcert makes one — if the browser blocks it.

A backend without the diffusion stack fails the action with FAILED_PRECONDITION; the demo shows its message in the error region.

Try it

The bar moves once per denoising step, from the progress port; the image appears when image closes.

The action is a11/demos/text_to_image.py and the page is js/demo/generative_media.ts. HTTP as separate streams applies the same design to HTTP protocol fields.

1. Define the contract

TEXT_TO_IMAGE_SCHEMA = a11.ActionSchema(
    name="text_to_image",
    description="Draw an image from a prompt, reporting progress as it goes.",
    inputs={
        "request": a11.ActionPortSchema(
            name="request", type="application/json", typeinfo=DiffusionRequest,
            unary=True, required=True,
            description="`{prompt, num_inference_steps, height, width, seed}`.",
        )
    },
    outputs={
        "image": a11.ActionPortSchema(
            name="image", type="image/png", unary=True, required=True),
        "progress": a11.ActionPortSchema(
            name="progress", type="application/json", typeinfo=dict,
            description="`{step, steps}` once per denoising step."),
    },
)

unary=True declares a port that carries one complete value. The progress port omits it and carries a stream. The declared DiffusionRequest type adds its JSON Schema to the action contract and decodes the browser's plain JSON into the model:

request = await action["request"].consume(DiffusionRequest)

2. Report from a worker thread

The pipeline is blocking, so it runs on a thread. Some schedulers invoke the callback once beyond the requested step count, so the progress value is capped at the declared total:

loop = asyncio.get_running_loop()

def on_step(_pipeline, step, _timestep, kwargs):
    done = min(step + 1, request.num_inference_steps)
    asyncio.run_coroutine_threadsafe(
        progress.put({"step": done, "steps": request.num_inference_steps},
                     mimetype="application/json"),
        loop,
    )
    return kwargs

result = await asyncio.to_thread(
    pipeline,
    request.prompt,
    num_inference_steps=request.num_inference_steps,
    height=request.height,
    width=request.width,
    callback_on_step_end=on_step,
)

The handler omits each progress tick's confirmation future, allowing denoising to continue while storage accepts the record. Payloads that require confirmed storage await both stages with await (await node.put(value)), as a11.gateway.conversations.ConversationStore.record does.

Both ports are closed however the handler ends:

try:
    ...
finally:
    await progress.close()
    await image_out.close()

Close every output port so readers can observe the end of the stream.

3. Write the encoded image

The handler encodes the image in its chosen format and labels the chunk with the corresponding media type. The chunk is built and written directly:

def _png_chunk(png: bytes) -> a11.Chunk:
    return a11.Chunk(
        data=png, metadata=a11.ChunkMetadata(mimetype="image/png")
    )

png = await asyncio.to_thread(_png_bytes, result.images[0])
await image_out.put_chunk(_png_chunk(png), final=True)

put encodes a value through a registered (type, media type) codec. Raw PNG bytes have no application value codec, so put returns NOT_FOUND. Write an already encoded payload as a chunk, as a11.sdk.http.client does for request bodies.

The browser reads the PNG as a chunk because it has no registered application type:

const node = need(await call.getOutput('image', false));
const chunk = need(await node.nextChunk(MAX_IMAGE_BYTES));
image.src = URL.createObjectURL(new Blob([chunk.data], {type: chunk.mimetype}));

4. Read both ports at once

const progress = readPort(call, 'progress', (value) => {
    bar.value = (value as Progress).step;
});
const chunk = need(await node.nextChunk(MAX_IMAGE_BYTES));
await progress;

nextChunk takes the largest payload the reader will accept. The page sets a 4 MB ceiling for this action's PNG output.

An undrained output port stalls its producer. Read progress while waiting for the image.