图像扩散生成
更新时间：2026-09-14
简介
Stable Diffusion img2img 图像增广算子：以原图为条件、caption 为 prompt 走 img2img 重绘，为每张图生成 aug_num 张增广图并落盘，strength 控制改动强度、guidance_scale 控制 CFG 强度。默认权重 CompVis/stable-diffusion-v1-4。
功能描述
- 推理前把图 BILINEAR resize 到 512×512，出图后再 BILINEAR 缩回原图尺寸，输出图与输入图同分辨率
- caption 是可选的第二个输入列：该行 caption 为空、非字符串或整列不传时用
default_prompt；算子内不会再套一个大模型现算 caption - 命中 pipeline 自带的 NSFW safety checker 时重跑采样，最多 3 次；3 次仍命中该行按失败处理（safety checker 被权重裁掉时跳过该判断）
- 列式算子行数恒定，输出列是该行生成的
aug_num个图片路径列表，不含原图 aug_num > 1时逐张独立采样，耗时随aug_num线性增长- 输出文件名
<输入图 stem>_diffusion_<序号>.jpg；输入不是路径（base64 / binary）时 stem 取row<行内下标> - 权重只从本地
{model_path}/{model_name}加载，目录不存在时在构造期抛FileNotFoundError，不联网下载、不做运行时 pip 安装 - 构造期校验
image_src_type合法、strength ∈ [0, 1]、aug_num > 0 - 输入为 null 直接返回空列表；单行异常（解码失败、NSFW 重试用尽等）记日志后返回空列表，不影响其它行
算子参数
输入
| 输入 | 含义 |
|---|---|
| image | 图像输入，内容类型由 image_src_type 决定（本地/BOS/HTTP 路径、Base64 字符串、二进制数据） |
| caption | 可选，img2img 使用的 prompt 文本列；该行为空或整列不传时使用 default_prompt |
输出
| 输出 | 含义 |
|---|---|
| aug_images | list<large_string>：本行生成的 aug_num 个增广图路径（output_bosdir 非空时为 BOS 路径）；输入为空或该行失败时为空列表 |
参数
| 参数名称 | 类型 | 默认值 | 描述 |
|---|---|---|---|
| image_src_type | str | "image_url" | 图像输入类型：image_url / image_base64 / image_binary |
| model_path | str | "/opt/aihc/model" | 权重根目录 |
| model_name | str | "CompVis/stable-diffusion-v1-4" | 相对 model_path 的 SD 权重子目录 |
| dtype | str | "float16" | 权重精度：float32 / float16 / bfloat16。GPU 上建议 float16 |
| strength | float | 0.8 | img2img 改动强度，取值 [0, 1] |
| guidance_scale | float | 7.5 | CFG 强度 |
| aug_num | int | 1 | 每张图生成几张增广图 |
| default_prompt | str | "A photo of a scene" | caption 缺省时使用的 prompt |
| num_inference_steps | int | 50 | 采样步数 |
| output_dir | str | "/tmp/aihc_image_diffusion" | 本地输出目录 |
| output_bosdir | str | "" | 非空时把生成图上传到该 BOS 目录，并返回 BOS 路径 |
| rank | int | 0 | 多卡场景 worker 序号，设备取 cuda:(rank % 可见卡数) |
注意事项
- 权重要求 safetensors 格式（
use_safetensors=True）。 - 只有
aihc_udf(num_gpus=...)大于 0 时才用 GPU；否则整条 SD 采样跑在 CPU 上，速度不可用。 - 输出文件名只由输入图 stem 与序号决定：不同目录下同名的输入图、或使用 base64/binary 输入时不同微批的相同行内下标，都会写到同一个输出路径而互相覆盖。批量场景建议按批设置不同的
output_dir。 num_inference_steps默认 50，配合aug_num会直接决定单行耗时；回归测试里取 20 步。
调用示例
Python
1from __future__ import annotations
2
3import os
4
5import daft
6from daft import col
7
8from daft.aihc.common.udf import aihc_udf
9from daft.aihc.functions.image.image_diffusion import ImageDiffusion
10
11if __name__ == "__main__":
12 if os.getenv("DAFT_RUNNER", "native") == "ray":
13 import ray
14 ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
15 daft.set_runner_ray()
16 daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
17
18samples = {
19 "image": ["bos://your-bucket/sample.jpg"],
20 "caption": ["a photo of a street scene"],
21 }
22 ds = daft.from_pydict(samples)
23 ds = ds.with_column(
24 "aug_images",
25 aihc_udf(
26 ImageDiffusion,
27 construct_args={
28 "image_src_type": "image_url",
29 "model_path": "/path/to/models",
30 "model_name": "CompVis/stable-diffusion-v1-4",
31 "dtype": "float16",
32 "aug_num": 2,
33 "num_inference_steps": 20,
34 "output_dir": "/tmp/aihc_image_diffusion",
35 },
36 num_cpus=1,
37 num_gpus=1,
38 concurrency=1,
39 batch_size=1,
40 )(col("image"), col("caption")),
41 )
42 ds.show()评价此篇文章