从文本生成图像

当提示词、模型配置、二进制输出和元数据需要在大量数据行中保持对齐时，文生图就成了数据处理流程问题。本教程沿着 examples/image_generation.py，使用 Vane 批量 UDF 构建这条处理流程。

脚本提供两个相互独立的选择：

• 提示词来源：内置样例提示词，或 LAION Parquet 文件中的文本与元数据。
• 生成后端：无需下载模型的确定性占位生成器，或通过 Diffusers 使用 Stable Diffusion。

默认输入把内置提示词与占位后端组合起来。你可以先验证 Relation 构建、批量执行、PNG 序列化和输出写入，再使用 GPU 资源。

1. 构建提示词 Relation

样例数据源直接产生 id、prompt、source_url 和 aesthetic_score 列。LAION 输入会加载 DuckDB 的 HTTP/S3 扩展，从 Parquet 中选择相同列结构，删除过短或为空的提示词，并应用行数上限。

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)}
        """
    )

让两个数据源保持结构兼容，意味着生成 UDF 无需知道数据行来自本地样例还是远程语料。

2. 用确定性 PNG 验证处理流程

占位后端会对每条提示词做哈希，并从摘要中派生配色。它只使用 Python 标准库写出有效 RGB PNG，因此相同的提示词、宽度和高度会产生完全相同的字节。

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"")
    )

这个后端不是文生图模型，而是用于验证模型周边各个处理阶段的确定性契约测试后端。

3. 每个 UDF 实例只加载一次 Stable Diffusion

真实生成路径中的有状态 UDF 会推迟到首次使用时才加载模型，并在实例上缓存模型管线，从而避免该实例处理后续批次时再次加载模型。

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

模型 ID、revision、设备、数据类型和仅本地文件行为都可以通过参数配置。如果模型加载失败，错误信息会给出具体的模型下载和脚本调用命令。

生成方法会把图像尺寸、推理步数、引导系数，以及为每行设置随机种子的生成器传给 Diffusers，再把返回的第一张图编码为 PNG 字节。

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}")

把行 ID 加到基础随机种子上，可以在保持输出可复现的同时，避免每条提示词使用完全相同的随机流。

4. 声明批量输出契约

UDF 返回全部输入元数据以及 generated_image 字节。调用处声明每个 DuckDB 类型，并且只在显式提供 GPU 数量时添加资源请求。

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",
        },
    )

默认批大小为一，是 GPU 显存方面的安全起点。只有在综合检查所选模型、分辨率、精度和执行节点资源后，才应增大它。

5. 检查生成文件与元数据

每个输出行会变成一个带序号的 PNG 文件。metadata.csv 记录行 ID、原始提示词和生成路径；如果源数据存在美学分数，终端 Relation 预览也会显示它。

命令行默认使用 512×512 输出和 20 个扩散推理步骤；执行方式由 Vane 配置的 runner 决定。占位后端会忽略仅与扩散模型有关的设置。进行真实推理时，请选择 Diffusers 后端，确保每个执行节点都能访问模型，选择合适精度，并显式请求 GPU 资源。

完整 UDF 构造器、样例提示词生成、文件命名与全部运行参数请查看完整源码。