Receive progress and a PNG from one image request

Run this guide to save a real PNG while structured progress arrives from the same action. The Python client keeps the prompt, timeout and file-writing policy; the hosted GPU keeps the model and its weights. Replacing the backend does not require the image UI to reinterpret a provider event stream.

Generate your first image#

Select Generate in your browser, enter A small blue lighthouse on a quiet island, watercolor, set Steps to 10 and Seed to 7, then select Draw. The existing hosted demo runs Stable Diffusion and displays the result. It needs no local GPU or API key.

To save the image directly from Python, use Run this guide above in a fresh directory with an active Python 3.11+ environment.

Your prompt goes to the shared backend; model weights stay there. The program writes a11-lighthouse.png in your current directory, replacing that filename if it already exists. Use --output another-name.png to choose another destination.

Expected terminal shape:

Step 1/10
Step 2/10
…
Step 10/10
Saved a11-lighthouse.png (… bytes)

The image appears only after the action succeeds. The model's scheduler can report the last step twice. The image output still contains one PNG.

This is a saved result from the documented prompt, seed 7 and ten steps. It illustrates the output; the launch commands perform a new generation. Different hardware or model versions can change pixels even with the same seed.

Generate a cabin and save it to a different file#

You have a lighthouse image. To request a cabin, use a new prompt and output filename; the client still reads the same progress and image ports:

python image_client.py "A red cabin beside a snowy forest, watercolor" --seed 12 --steps 20 --output cabin.png

After connecting and discovering the schema, the client dispatches the image action:

registry = a11.ActionRegistry()
await install_peer_actions(connection, registry)
action = connection.action(
    "text_to_image", registry.get_schema("text_to_image")
)
await action.call()

The script then sends a single structured request:

await action["request"].finalize({
    "prompt": args.prompt, "num_inference_steps": args.steps,
    "height": 512, "width": 512, "seed": args.seed,
})

The backend validates dimensions and step count. This client fixes the dimensions at 512×512 and accepts 10–50 steps. Its default of ten keeps the first run small. Shared GPU capacity can make a run slower; the client stops waiting after 180 seconds instead of waiting indefinitely.

Update a progress bar before the PNG exists#

A user has clicked Draw, but the GPU has not finished an image yet. Your UI needs small progress updates now and the PNG later. Separate outputs let the progress bar update without making the image viewer inspect every status message.

The action exposes one request and two outputs:

The two readers handle different representations:

async def read_progress():
    async for update in action["progress"]:
        print(f"Step {update['step']}/{update['steps']}", flush=True)

async def read_image():
    # PNG is already encoded. Read bytes, not the text/JSON decoder.
    return await action["image"].next_chunk()

The image reader uses next_chunk(). The PNG is already encoded; passing it through a text decoder would destroy the contract. bytes(chunk.data) is exactly what the file writer needs.

Start both readers together, then check the action’s final status:

async with asyncio.timeout(args.timeout):
    async with asyncio.TaskGroup() as readers:
        readers.create_task(read_progress())
        image = readers.create_task(read_image())
    await action.wait()

image is the task whose result contains the PNG chunk. Both readers run inside the same TaskGroup. They are concurrent because they have different arrival times. Reading only the image and ignoring progress would be a poor general pattern for bounded streams: an undrained output can hold its producer up.

Load the model once and forward its denoising updates#

If every click reloaded the model weights, users would pay that startup cost repeatedly. The backend keeps the pipeline loaded and forwards its step callbacks while inference runs.

The model handler contains the schema, validation, model loading, inference and PNG encoding. The important sequence is:

1. Consume and validate the request.
2. Load a diffusion pipeline once and reuse it.
3. Run blocking model inference on a worker thread.
4. Forward denoising callbacks to the async loop and the progress output.
5. Encode the finished image and write its PNG chunk to image.
6. Close both outputs, including when an operation fails.

The server’s diffusion callback runs on a worker thread. This excerpt from the linked model handler forwards each step back to the async event loop:

loop = asyncio.get_running_loop()

def on_step(_pipeline, step: int, _timestep: int, kwargs: dict) -> dict:
    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

progress is the server action’s output node; request is its validated generation request. The callback is passed as callback_on_step_end when the handler calls the pipeline via asyncio.to_thread. This server code explains where the client’s progress values originate; it is not part of the client download.

A lock admits one generation at a time on that pipeline. The handler owns this resource policy. The UI depends on the output contract and does not import the diffusion library.

If you already have a suitable GPU and want to run generation without relying on the shared demo server’s capacity, install the provider dependencies on that machine:

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

Then, in another terminal:

python image_client.py --endpoint ws://127.0.0.1:9010/a11-demos

This path downloads substantial model weights and needs suitable compute. The default walkthrough uses the hosted GPU.

Keep a failed render from looking like a finished image#

Imagine the lighthouse run succeeded, then you request a red cabin using the same output filename. The second run reaches step eight and fails. The old lighthouse file still exists. Displaying that file as the cabin result would tell the user the wrong request succeeded; leaving the progress indicator spinning would make them wait for an image that will never arrive.

The await action.wait() above must succeed before this file-writing code runs:

chunk = image.result()
if chunk is None:
    raise Status(
        code=StatusCode.DATA_LOSS,
        message="The action completed without an image.",
    ).to_exception()
destination = Path(args.output)
destination.write_bytes(bytes(chunk.data))
print(f"Saved {destination} ({destination.stat().st_size} bytes)")

A missing chunk raises DATA_LOSS; it does not create an empty image file. Cleanup still runs on success, failure or timeout:

finally:
    if action is not None:
        action.cancel()
    await connection.aclose()

The Python client therefore waits for the action’s successful final status before writing the new PNG. A UI needs the same distinction:

Press Ctrl+C in the terminal to cancel a run. A timeout or server failure before the image is saved leaves the previous file untouched. The terminal’s Saved … message identifies a successful new write.

If the server reports a rate limit, retry later. If it reports missing model dependencies, restarting the same request cannot install them; configure the backend or choose a working endpoint. Increase --timeout when a known slow machine needs more time. A reported failure still requires its own fix.

This example demonstrates structured text and image bytes. Audio can use the same pattern with its own media type and timing information; it does not need to be squeezed into the image or progress stream. The voice walkthrough applies the same separation to audio buffers and transcript text.

Complete source#

Download image_client.py. The browser's own reader is also available in generative_media.ts.

"""Generate a PNG remotely while reading denoising progress separately."""

import argparse
import asyncio
from pathlib import Path

import a11
from a11 import timing
from a11.client.connection import GatewayConnection
from a11.client.discovery import install_peer_actions
from a11.status import Status, StatusCode


async def main(args):
    connection = await GatewayConnection.connect(
        args.endpoint, timeout=timing.Duration.seconds(15)
    )
    action = None
    try:
        registry = a11.ActionRegistry()
        await install_peer_actions(connection, registry)
        action = connection.action(
            "text_to_image", registry.get_schema("text_to_image")
        )
        await action.call()
        await action["request"].finalize(
            {
                "prompt": args.prompt,
                "num_inference_steps": args.steps,
                "height": 512,
                "width": 512,
                "seed": args.seed,
            }
        )

        async def read_progress():
            async for update in action["progress"]:
                print(f"Step {update['step']}/{update['steps']}", flush=True)

        async def read_image():
            # PNG is already encoded. Read bytes, not the text/JSON decoder.
            return await action["image"].next_chunk()

        async with asyncio.timeout(args.timeout):
            async with asyncio.TaskGroup() as readers:
                readers.create_task(read_progress())
                image = readers.create_task(read_image())
            await action.wait()
        chunk = image.result()
        if chunk is None:
            raise Status(
                code=StatusCode.DATA_LOSS,
                message="The action completed without an image.",
            ).to_exception()
        destination = Path(args.output)
        destination.write_bytes(bytes(chunk.data))
        print(f"Saved {destination} ({destination.stat().st_size} bytes)")
    finally:
        if action is not None:
            action.cancel()
        await connection.aclose()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "prompt",
        nargs="?",
        default="A small blue lighthouse on a quiet island, watercolor",
    )
    parser.add_argument("--endpoint", default="wss://a11.to/ws/demoserver")
    parser.add_argument(
        "--steps", type=int, choices=range(10, 51), default=10, metavar="10..50"
    )
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--output", default="a11-lighthouse.png")
    parser.add_argument("--timeout", type=float, default=180)
    args = parser.parse_args()
    try:
        asyncio.run(main(args))
    except KeyboardInterrupt:
        print("\nStopped. No new image was saved.")
    except TimeoutError:
        raise SystemExit(
            "Timed out; the action was cancelled. Retry later or increase "
            "--timeout."
        )