Skip to main content
Vane Data / Tutorials

Generating Images from Text

Text-to-image generation becomes a data pipeline problem when prompts, model settings, binary outputs, and metadata must stay aligned across many rows. This tutorial follows examples/image_generation.py to build that pipeline with a Vane batch UDF.

The script offers two independent choices:

  • Prompt source: built-in sample prompts or text and metadata from a LAION Parquet file.
  • Generation backend: a deterministic placeholder that requires no model download, or Stable Diffusion through Diffusers.

The default path combines built-in prompts with the placeholder backend. It lets you validate Relation construction, batch execution, PNG serialization, and output writing before using GPU resources.

1. Build a prompt Relation

The sample source produces id, prompt, source_url, and aesthetic_score columns directly. The LAION path loads DuckDB's HTTP/S3 extension, selects the same schema from Parquet, removes short or null prompts, and applies a limit.

example.py
def load_laion_relation(conn: Any, parquet_path: str, limit: int) -> Any:
    try:
        conn.execute("INSTALL httpfs")
        conn.execute("LOAD httpfs")
    except Exception:
        # httpfs may already be installed/loaded, or the path may be local.
        pass
    try:
        conn.execute("SET s3_region='us-west-2'")
        conn.execute("SET s3_url_style='path'")
    except Exception:
        pass


    return conn.sql(
        f"""
        select
            row_number() over () - 1 as id,
            TEXT as prompt,
            URL as source_url,
            cast(AESTHETIC_SCORE as double) as aesthetic_score
        from read_parquet({sql_literal(parquet_path)})
        where TEXT is not null
          and length(TEXT) > 50
        limit {int(limit)}
        """
    )

Keeping both sources schema-compatible means the generation UDF never needs to know whether a row came from local sample data or a remote corpus.

2. Verify the pipeline with deterministic PNG files

The placeholder hashes each prompt and derives a color palette from the digest. It writes a valid RGB PNG using only the Python standard library, so identical prompt text, width, and height produce identical bytes.

example.py
def placeholder_png(width: int, height: int, prompt: str) -> bytes:
    """Create a deterministic PNG thumbnail from prompt text."""
    digest = hashlib.sha256(prompt.encode("utf-8")).digest()
    top = digest[0], digest[1], digest[2]
    bottom = digest[3], digest[4], digest[5]
    accent = digest[6], digest[7], digest[8]
    stripe_width = max(8, width // 12)


    rows: list[bytes] = []
    for y in range(height):
        t = y / max(1, height - 1)
        base = tuple(int(top[i] * (1 - t) + bottom[i] * t) for i in range(3))
        row = bytearray(b"\x00")
        for x in range(width):
            wave = 0.5 + 0.5 * math.sin((x + digest[9]) / max(1, width) * math.tau * 3)
            stripe = 1 if (x // stripe_width + y // stripe_width) % 2 == 0 else 0
            rgb = []
            for i in range(3):
                value = base[i] * 0.78 + accent[i] * 0.22 * wave
                if stripe:
                    value += 22
                rgb.append(max(0, min(255, int(value))))
            row.extend(rgb)
        rows.append(bytes(row))


    raw = b"".join(rows)
    return (
        b"\x89PNG\r\n\x1a\n"
        + png_chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
        + png_chunk(b"IDAT", zlib.compress(raw, level=9))
        + png_chunk(b"IEND", b"")
    )

This backend is not a text-to-image model. It is a deterministic contract-test backend for every pipeline stage around the model.

3. Load Stable Diffusion once per UDF instance

For real generation, the stateful UDF delays model loading until its first use and caches the pipeline on the instance, avoiding another model load for later batches handled by that instance.

example.py
    def _load_pipe(self) -> Any:
        if self._pipe is not None:
            return self._pipe


        try:
            import torch
            from diffusers import StableDiffusionPipeline
        except ImportError as exc:
            raise RuntimeError(
                "Install image generation dependencies first: "
                "pip install diffusers transformers accelerate torch Pillow"
            ) from exc


        torch_dtype = getattr(torch, self.dtype)
        try:
            pipe = StableDiffusionPipeline.from_pretrained(
                self.model_id,
                torch_dtype=torch_dtype,
                revision=self.revision,
                local_files_only=self.local_files_only,
            )
        except Exception as exc:
            raise RuntimeError(
                "Could not load the diffusion model. Download it first, then pass "
                "the local directory with --model-id. For example:\n\n"
                "  hf download "
                f"{DEFAULT_MODEL_ID} --local-dir ~/.cache/vane/models/stable-diffusion-v1-5\n"
                "  python "
                "examples/image_generation.py --backend diffusers "
                "--model-id ~/.cache/vane/models/stable-diffusion-v1-5 "
                "--source sample --limit 2 --device cuda --dtype float16\n\n"
                f"Original error: {type(exc).__name__}: {exc}"
            ) from exc
        pipe.enable_attention_slicing(1)
        if self.device:
            pipe = pipe.to(self.device)
        self._pipe = pipe
        return pipe

The model ID, revision, device, data type, and local-files-only behavior are all arguments. If model loading fails, the error includes concrete model-download and script-invocation commands.

The generation method forwards image dimensions, inference steps, guidance scale, and a random generator seeded for each row to Diffusers, then encodes the first returned image as PNG bytes.

example.py
    def _generate_with_diffusers(self, prompt: str, index: int) -> bytes:
        import torch


        pipe = self._load_pipe()
        generator = None
        if self.seed is not None:
            generator = torch.Generator(device=self.device or "cpu").manual_seed(self.seed + index)
        image = pipe(
            prompt,
            num_inference_steps=self.num_inference_steps,
            height=self.height,
            width=self.width,
            guidance_scale=self.guidance_scale,
            generator=generator,
        ).images[0]
        return pil_to_png_bytes(image)


    def _generate_one(self, prompt: str, index: int) -> bytes:
        if self.backend == "placeholder":
            return placeholder_png(self.width, self.height, prompt)
        if self.backend == "diffusers":
            return self._generate_with_diffusers(prompt, index)
        raise ValueError(f"Unsupported backend: {self.backend}")

Adding the row ID to the base seed keeps outputs reproducible while avoiding the same random stream for every prompt.

4. Declare the batch output contract

The UDF returns all input metadata plus generated_image bytes. The call site declares every DuckDB type and adds a GPU resource request only when one was supplied.

example.py
    map_kwargs: dict[str, Any] = {
        "schema": {
            "id": vane.sqltypes.BIGINT,
            "prompt": vane.sqltypes.VARCHAR,
            "source_url": vane.sqltypes.VARCHAR,
            "aesthetic_score": vane.sqltypes.DOUBLE,
            "generated_image": vane.sqltypes.BLOB,
        },
        "batch_size": args.batch_size,
    }
    if args.gpus is not None:
        map_kwargs["gpus"] = args.gpus


    generated = rel.map_batches(generator.__call__, **map_kwargs)
    generated_table = collect_relation(generated)
    written_table = save_generated_images(generated_table, Path(args.output_dir))
    written = relation_from_rows(
        conn,
        [
            {
                "id": row["id"],
                "prompt": row["prompt"],
                "aesthetic_score": row["aesthetic_score"],
                "generated_path": row["generated_path"],
            }
            for row in written_table.to_pylist()
        ],
        {
            "id": "BIGINT",
            "prompt": "VARCHAR",
            "aesthetic_score": "DOUBLE",
            "generated_path": "VARCHAR",
        },
    )

The default batch size is one, which is a safe starting point for GPU memory. Increase it only after checking the chosen model, resolution, precision, and worker resources together.

5. Inspect generated files and metadata

Each output row becomes a numbered PNG file. metadata.csv records the row ID, original prompt, and generated path, while the terminal Relation preview also shows the source aesthetic score when one exists.

The CLI defaults to 512×512 output and 20 diffusion inference steps; execution follows Vane's configured runner. The placeholder backend ignores the diffusion-only settings. For real inference, select the Diffusers backend, make the model available on every worker, choose an appropriate precision, and request GPU resources explicitly.

See the complete source for the full UDF constructor, sample prompt generation, file naming, and all runtime arguments.