Daft 实战教程:用 Stable Diffusion 在 GPU 上批量生成图像
Daft 实战教程用 Stable Diffusion 在 GPU 上批量生成图像【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft本文以 Daft 官方示例文档 Generate Images with Stable Diffusion 为主线讲解如何用 Stable Diffusion 模型将文本批量转换为图像并覆盖从 S3 读取 Parquet 数据集、下载并解码图像到通过daft.cls(gpusN)把扩散模型注册为 Daft UDF 的完整链路。读完本篇你可以掌握如何用 Daft 类 UDFclass UDF声明式地申请 GPU 资源、如何用 Daft 内建的download/decode_image表达式做多模态数据预处理以及这些 API 在 Daft 源码中的实现位置与行为细节。对应教程的可交互版本为仓库中的 text_to_image_generation.ipynb本文所有代码均与文档和 Notebook 保持一致可直接复制到本地运行。环境准备依赖安装与 GPU 运行时运行本教程需要一台带 GPU 的机器例如 Google Colab 的 GPU 运行时。在 Colab 中切换方式为菜单Runtime - Change runtime type - Hardware accelerator - GPU - Save。安装依赖pip install daft pip install transformers diffusers accelerate torch Pillowdaft数据引擎本体提供 DataFrame、UDF 与 IO 能力transformers/diffusers/accelerate/torchStable Diffusion 推理栈其中diffusers提供StableDiffusionPipelinePillow用于处理生成的PIL.Image对象。教程中通过一个USE_GPU开关来对比 CPU 与 GPU 的执行差异后续 UDF 声明会引用它。读取数据从 S3 加载 Parquet 并物化缓存教程使用的数据集是一个存放在 S3 上的 LAION 风格 Parquet 文件包含URL图片地址、TEXT图像描述文本、AESTHETIC_SCORE美学评分等列import daft # Flip this flag if you want to see the performance of running on CPU vs GPU USE_GPU True IO_CONFIG daft.io.IOConfig( s3daft.io.S3Config(anonymousTrue, region_nameus-west-2) ) # Use anonymous-mode for accessing AWS S3 PARQUET_PATH s3://daft-oss-public-data/tutorials/laion-parquet/train-00000-of-00001-6f24a7497df494ae.parquet parquet_df daft.read_parquet(PARQUET_PATH, io_configIO_CONFIG)要点说明daft.io.S3Config(anonymousTrue, region_nameus-west-2)以匿名模式访问公共桶因此无需配置 AWS 凭证私有桶则需填入对应凭证或 IAM 角色。读取后先执行parquet_df.collect()把 Parquet 文件下载到内存中物化使后续所有操作都命中缓存避免重复下载parquet_df.collect() parquet_df parquet_df.select(parquet_df[URL], parquet_df[TEXT], parquet_df[AESTHETIC_SCORE])select之后 DataFrame 只保留三列图片 URL、描述文本和美学分。其中TEXT列是后面送入 Stable Diffusion 的提示词prompt来源。下载并解码图像两行表达式完成多模态预处理与许多数据集相同该数据集并不直接存储图像二进制而是存储图像的 URL。Daft 提供内建的表达式级下载与解码能力几行代码即可把 URL 列转成可展示的 PIL 图像列# Filter for images with longer descriptions parquet_df_with_long_strings parquet_df.where(parquet_df[TEXT].length() 50) # Download images images_df ( parquet_df_with_long_strings.with_column( image, parquet_df_with_long_strings[URL].download(on_errornull).decode_image(on_errornull), ) .limit(5) .where(daft.col(image).not_null()) ) images_df.show(5)这里有两个值得展开的表达式download(on_errornull)把每一行字符串当作 URL下载其字节内容并生成 bytes 列。失败时on_errornull会把出错行置空而不是抛出异常。从源码看表达式方法定义在 download其参数为max_connections默认 32 个并发连接、on_errorraise/null默认raise和可选的io_config最终委托给daft.functions.download执行。对URL 可能失效的现实数据集而言on_errornull是保证批处理不中断的关键。decode_image(on_errornull)把 bytes 列解码为图像类型默认解码模式为 RGB。实现位于 decode_image签名参数on_error与modestr | ImageMode | None默认ImageMode.RGB与文档用法一一对应底层调用daft.functions.decode_image。随后的limit(5)加where(daft.col(image).not_null())组合是文档中演示数据集长相的实用技巧先限制样本量再过滤掉下载/解码失败的行确保show(5)总能展示出一批有效图像。用 GPU 跑 Stable Diffusiondaft.cls(gpusN)类 UDF接下来把 Stable Diffusion 模型作用在TEXT列上为每条文本生成图像。Daft 使用 GPU 的方式非常直接在daft.cls装饰器上指定gpusN其中N是每个 UDF 实例应预留的 GPU 数量import torch from diffusers import StableDiffusionPipeline daft.cls(gpus1 if USE_GPU else 0) class GenerateImageFromText: def __init__(self): model_id runwayml/stable-diffusion-v1-5 self.pipe StableDiffusionPipeline.from_pretrained( model_id, torch_dtypetorch.float32, ) self.pipe.enable_attention_slicing(1) def generate_image(self, prompt): return self.pipe(prompt, num_inference_steps20, height512, width512).images[0] daft.method.batch(return_dtypedaft.DataType.python()) def __call__(self, text_col): return [self.generate_image(t) for t in text_col] image_generator GenerateImageFromText() images_df.with_column( generated_image, image_generator(images_df[TEXT]), ).show(1)逐层拆解这段代码的设计意图__init__中加载模型只做一次StableDiffusionPipeline.from_pretrained(runwayml/stable-diffusion-v1-5, ...)的加载开销被摊薄到多个批次数据上。这正是 Daft 类 UDF 相对行级函数的核心优势——从 cls 装饰器文档字符串 可以看到官方定位即initialize a class instance once, and then reuse it for multiple rows of data并明确说明 Daft 类是惰性初始化的实例化GenerateImageFromText()时只保存参数__init__直到查询执行时才在每个 worker 实例上被调用。enable_attention_slicing(1)启用注意力切片以降低峰值显存使 v1-5 模型能在消费级 GPU 上以float32精度运行。generate_image中的推理参数num_inference_steps20采样步数步数越多越精细但越慢、height512, width512SD v1-5 的原生分辨率、取images[0]返回单张PIL.Image。daft.method.batch(return_dtypedaft.DataType.python())把__call__声明为批量方法接收daft.Series列逐行调用generate_image后返回 Python 对象列表return_dtypedaft.DataType.python()声明输出为任意 Python 对象列这里是 PIL 图像。批量装饰器的参数语义return_dtype、batch_size等见 batch 方法装饰器。最后一行即整条流水线images_df[TEXT]作为表达式传入image_generator得到新的表达式列generated_imageshow(1)触发执行并展示结果。gpus参数的取值规则与资源调度从 cls 装饰器实现 可以确认gpus的校验规则这也是使用 GPU UDF 前需要了解的行为边界默认值为0即不使用 GPU可如教程中gpus1 if USE_GPU else 0那样切换 CPU/GPU 对照实验支持(0, 1.0]区间的小数值如0.5官方文档字符串说明其用途是在同一张 GPU 上运行多个小模型大于1的取值必须是整数如2、4否则在装饰器阶段就抛出ValueError而不是等到调度时才报错。另外注意gpus是 Daft 层面的资源请求声明用于放置与分配在 Ray 执行器下如果同时在ray_options里设置num_gpus会直接报错——源码在 udf_v2.py 中提示 Cannot set num_gpus inray_options. Please use the gpus argument即资源声明应统一走gpus参数。在 GPU 可见性方面Daft 通过 daft/internal/gpu.py 的cuda_visible_devices()读取CUDA_VISIBLE_DEVICES或调用 NVML 枚举当前进程可见的 CUDA 设备调度器据此把请求了gpus1的 UDF 实例放置到有可用 GPU 的 worker 上。相关资源入口本教程在官方文档站点中的入口见 examples 目录其中对本文的描述为Open Source image generation model on your own GPUs using Daft UDFs。更完整的 UDF 模式行级、生成器、异步、有状态类可参考仓库中的 Dafts Four UDF Pattern Tutorial 文档daft.cls与daft.func的迁移背景见 custom-code 文档。小结这篇教程展示的是 Daft 处理多模态生成式任务的典型范式表达式算子负责数据获取与清洗downloaddecode_imageon_errornull保证容错类 UDF 负责把昂贵的模型加载与 GPU 推理封装成可并行执行的算子daft.cls(gpusN)daft.method.batch。整套流程从读 S3 Parquet 到展示生成图像不超过几十个 Python 语句且 GPU 资源以声明式参数显式表达为后续扩展到 Ray 集群多 GPU 并行留出了直接的路径。【免费下载链接】DaftHigh-performance data engine for AI and multimodal workloads. Process images, audio, video, and structured data at any scale项目地址: https://gitcode.com/GitHub_Trending/da/Daft创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考