Running a model with weights

Load a real model on a GPU: where the imports go, how to declare dependencies, and why weights belong on a volume.

Introduction

A text-to-image model on a GPU, with the weights on a volume so a restart does not fetch them again.

Z-Image Turbo is Apache 2.0 and ungated, so you deploy it with no token.

The app

model.py

"""Where the heavy parts of an app go, and why the weights belong on a volume."""

from __future__ import annotations

import base64
import io

from pydantic import BaseModel, Field

from runware_serverless import serve

MODEL = "Tongyi-MAI/Z-Image-Turbo"
"""Apache 2.0 and ungated, so this deploys with no token of any kind."""

CACHE = "/weights"
"""Mount a volume here. The worker's filesystem is part of what a checkpoint
captures, so weights left outside a volume are copied into every checkpoint and
fetched again on every cold start."""

DEFAULT_STEPS = 8
"""A turbo model, so the schedule is short. Raising this buys little."""


class Request(BaseModel):
    prompt: str = Field(description="What to generate.")
    steps: int = Field(default=DEFAULT_STEPS, ge=1, le=50)


class Generated(BaseModel):
    image: str = Field(description="The result as a base64 PNG.")
    steps: int


@serve
class LoadingWeights:
    requirements = [
        "torch==2.14.1",
        "diffusers==0.42.1",
        "transformers==5.21.0",
        "accelerate==1.15.0",
    ]

    def load(self) -> None:
        if hasattr(self, "pipeline"):
            return

        # Here and not at the top of the file. The build imports this file with the
        # platform's own dependencies, not yours, so nothing in requirements exists
        # until a worker runs.
        import torch
        from diffusers import DiffusionPipeline

        self.pipeline = DiffusionPipeline.from_pretrained(
            MODEL, torch_dtype=torch.bfloat16, cache_dir=CACHE
        ).to("cuda")

        # The first call compiles kernels and allocates, so no request should be
        # the one paying for it.
        self.pipeline(prompt="warmup", num_inference_steps=1)

    def predict(self, request: Request) -> Generated:
        image = self.pipeline(
            prompt=request.prompt, num_inference_steps=request.steps
        ).images[0]

        return Generated(image=_png(image), steps=request.steps)


def _png(image: object) -> str:
    buffer = io.BytesIO()
    image.save(buffer, format="PNG")  # pyright: ignore[reportAttributeAccessIssue]

    return base64.b64encode(buffer.getvalue()).decode("ascii")

Heavy imports

The build imports your file to find your endpoints, and it does that without your requirements installed, so a module-level import torch fails the build.

Inside load it runs in the worker, where torch is there. Every example that touches a real model imports this way.

Write load so it is safe to repeat. The platform retries it when the GPU is out of memory at startup, waiting for a larger allocation each time, and that retry re-instantiates your class.

Declaring dependencies

The build installs what you list in requirements, and keeping it next to the code means the next person reading your file finds it.

Pin the versions, and pin them to a release your GPU can run. Leave them open and a version you never chose lands on your next build.

Declaring the list on the class and passing --requirement on the deploy fails the build instead of picking one, because picking one would silently drop the other.

Weights on a volume

A checkpoint captures the worker's filesystem, so anything you download outside a volume is copied into every checkpoint and fetched again on every cold start.

Mount one at the cache directory and your second start reads from it.

Deploying it

runware serverless deploy model.py \
  --id loading-weights \
  --gpu-type rtx-pro-6000 \
  --volume /weights

Nothing in model.py names hardware. You choose the GPU when you deploy, so the same file runs on whatever you give it.

Calling it

The class declares no @endpoint, so @serve exposes predict at the default path.

import { createClient } from '@runware/sdk'

const client = await createClient({ apiKey: process.env.RUNWARE_API_KEY })

const task = await client.invoke({
  appId: 'loading-weights',
  endpointPath: 'predict',
  payload: {
    prompt: 'a red bicycle leaning on a wall'
  }
})

import asyncio
import os

from runware import Runware


async def main():
    client = Runware(api_key=os.environ["RUNWARE_API_KEY"])
    try:
        task = await client.invoke(
            "loading-weights",
            "predict",
            {
                "prompt": "a red bicycle leaning on a wall"
            },
        )
    finally:
        await client.close()


asyncio.run(main())

curl https://api.serverless.runware.ai/v1/apps/loading-weights/invoke-async/predict \
  -H "Authorization: Bearer $RUNWARE_API_KEY" \
  -H "Content-Type: application/json" \
  -d '{
    "taskId": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
    "payload": {
      "prompt": "a red bicycle leaning on a wall"
    }
  }'

cat <<'EOF' | runware serverless apps invoke loading-weights predict -f -
{
  "prompt": "a red bicycle leaning on a wall"
}
EOF

{
  "taskId": "7c9e6679-7425-40de-944b-e07fc1f90ae7",
  "payload": {
    "prompt": "a red bicycle leaning on a wall"
  }
}

The image comes back as a base64 PNG at output.predictions[0].image.

Measured on an RTX PRO 6000: 6m33s for the first call with an empty volume, 35 seconds for a restart that finds the weights already there, and 9 seconds on a worker already serving. The volume does not keep a worker warm, it makes a restart cheap.