Datawhale X 魔搭 AI夏令营 “AIGC”方向 task1

一、任务要求

task1 的任务和上一期的类似，都是跑通给出的代码即可，没有太大难度。

具体要求是训练 Lora 模型，实现文生图，额外的要求是8张图片必须组成一个连贯的故事，需要一定的“写小作文”能力。

二、代码解析

1. 下载数据集 这一步不用分析，就是保存图片和元数据。
2. 数据处理 题目已经提示了使用 data-juicer 处理数据。之所以选择 data-juicer，可能因为有自动化、批量处理、统一配置管理以及允许多核并行的优势。 代码其实很简单，只需要编写配置文件，就可以自动处理数据。具体解释如下： # data_juicer 配置data_juicer_config = """# global parameters（全局参数）project_name: 'data-process'dataset_path: './data/data-juicer/input/metadata.jsonl' # path to your dataset directory or file# 指定用于处理数据集的子进程数目np: 4 # number of subprocess to process your dataset# 定义了数据集中用于存储文本的键text_keys: 'text'# 定义了数据集中用于存储图像的键image_key: 'image'# 定义了图像的特殊标记，用于标记或识别图像数据image_special_token: '<__dj__image>'# 指定处理后的数据将导出的路径export_path: './data/data-juicer/output/result.jsonl'# process schedule# a list of several process operators with their arguments（定义了处理数据的操作列表，这里包含两个操作）process: # 过滤图像尺寸 - image_shape_filter: # 要求图像的最小宽度和高度分别为 1024 像素 min_width: 1024 min_height: 1024 # 选择 any 表示只要宽度或高度任意一项满足条件即通过 any_or_all: any # 过滤图像的宽高比 - image_aspect_ratio_filter: # 要求宽高比在 0.5 到 2.0 之间 min_ratio: 0.5 max_ratio: 2.0 # 选择 any 表示只要宽高比任意一项满足条件即通过 any_or_all: any"""# 将配置写入文件with open("data/data-juicer/data_juicer_config.yaml", "w") as file: file.write(data_juicer_config.strip())# 这行代码使用了 dj-process 命令，并指定了配置文件路径，用来启动数据处理任务# 其中 dj-process 是一个命令行工具，读取指定的配置文件，根据其中定义的参数和操作来处理数据集!dj-process --config data/data-juicer/data_juicer_config.yaml
3. 训练模型 这一部分的核心内容是通过一句命令行来执行的，具体代码如下： import oscmd = """python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py \ --pretrained_unet_path models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors \ --pretrained_text_encoder_path models/kolors/Kolors/text_encoder \ --pretrained_fp16_vae_path models/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors \ --lora_rank 16 \ --lora_alpha 4.0 \ --dataset_path data/lora_dataset_processed \ --output_path ./models \ --max_epochs 1 \ --center_crop \ --use_gradient_checkpointing \ --precision "16-mixed" """.strip()os.system(cmd) 参数的含义： python DiffSynth-Studio/examples/train/kolors/train_kolors_lora.py：运行指定的 Python 脚本。 --pretrained_unet_path：指定了预训练的 UNet 模型的路径。 --pretrained_text_encoder_path：指定了预训练的文本编码器路径。 --pretrained_fp16_vae_path：指定了预训练的 16 位混合精度 VAE 模型的路径。 --lora_rank：决定了新加入的低秩矩阵的维度。它控制了新参数矩阵的复杂度，较高的 lora_rank 值增加了模型的表达能力，但也会增加计算复杂度和内存使用；较低的 lora_rank 值则可能减少表达能力，但会更加高效。 --lora_alpha：低秩矩阵更新的缩放系数，用来平衡原始模型参数和 LoRA 新加入的参数之间的影响。较高的 lora_alpha 值会放大 LoRA 参数的影响力，可能导致模型更快收敛或更容易过拟合；较低的值则相反。 --dataset_path：指定了训练数据集的路径。 --output_path：指定了训练输出模型的保存路径。 --max_epochs：设置了最大训练轮数。 --center_crop：启用居中裁剪图像的选项。 --use_gradient_checkpointing：启用了梯度检查点，用于减少内存使用。 --precision "16-mixed"：使用混合精度训练，以提高计算效率。 具体的模型代码显然就在 train_kolors_lora.py 中了，简单分析一下。 from diffsynth import ModelManager, SDXLImagePipelinefrom diffsynth.trainers.text_to_image import LightningModelForT2ILoRA, add_general_parsers, launch_training_taskimport torch, os, argparse# 设置环境变量 TOKENIZERS_PARALLELISM 为 True，以允许多线程并行化分词处理os.environ["TOKENIZERS_PARALLELISM"] = "True"# LightningModel 继承自 LightningModelForT2ILoRA，这是一个用于文本到图像任务的 LoRA 训练模型类class LightningModel(LightningModelForT2ILoRA): # 初始化: # torch_dtype 指定了数据类型（如 float16 或 float32），主要用于控制模型的精度。 # pretrained_weights 是一个包含预训练模型路径的列表，用于加载预训练的 UNet、文本编码器、VAE 等模型。 # learning_rate 是学习率，用于训练过程中的优化。 # use_gradient_checkpointing 控制是否使用梯度检查点来节省内存。 # lora_rank 和 lora_alpha 是 LoRA 相关的参数，用于调整低秩矩阵的秩和缩放系数。 # lora_target_modules 指定了哪些模块（如 to_q、to_k 等）将应用 LoRA 技术。 def __init__( self, torch_dtype=torch.float16, pretrained_weights=[], learning_rate=1e-4, use_gradient_checkpointing=True, lora_rank=4, lora_alpha=4, lora_target_modules="to_q,to_k,to_v,to_out" ): super().__init__(learning_rate=learning_rate, use_gradient_checkpointing=use_gradient_checkpointing) # Load models # 模型管理和加载: # ModelManager 用于管理模型加载和设备分配 # SDXLImagePipeline 用于定义图像生成的流水线 # self.pipe.scheduler.set_timesteps(1100) 设置了调度器的时间步数 model_manager = ModelManager(torch_dtype=torch_dtype, device=self.device) model_manager.load_models(pretrained_weights) self.pipe = SDXLImagePipeline.from_model_manager(model_manager) self.pipe.scheduler.set_timesteps(1100) # Convert the vae encoder to torch.float16 # 参数调整和冻结: # 将 VAE 编码器转换为指定的数据类型 # 冻结模型的部分参数，以减少训练过程中的计算需求 # 将 LoRA 层添加到指定的模型模块中 self.pipe.vae_encoder.to(torch_dtype) self.freeze_parameters() self.add_lora_to_model(self.pipe.denoising_model(), lora_rank=lora_rank, lora_alpha=lora_alpha, lora_target_modules=lora_target_modules)# 解析命令行参数def parse_args(): # 定义了一个 argparse.ArgumentParser 对象，用于解析命令行参数 parser = argparse.ArgumentParser(description="Simple example of a training script.") # 下面是若干必需的参数 parser.add_argument( "--pretrained_unet_path", type=str, default=None, required=True, help="Path to pretrained model (UNet). For example, `models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors`.", ) parser.add_argument( "--pretrained_text_encoder_path", type=str, default=None, required=True, help="Path to pretrained model (Text Encoder). For example, `models/kolors/Kolors/text_encoder`.", ) parser.add_argument( "--pretrained_fp16_vae_path", type=str, default=None, required=True, help="Path to pretrained model (VAE). For example, `models/kolors/Kolors/sdxl-vae-fp16-fix/diffusion_pytorch_model.safetensors`.", ) parser.add_argument( "--lora_target_modules", type=str, default="to_q,to_k,to_v,to_out", help="Layers with LoRA modules.", ) parser = add_general_parsers(parser) args = parser.parse_args() return args# 启动训练if __name__ == '__main__': args = parse_args() model = LightningModel( torch_dtype=torch.float32 if args.precision == "32" else torch.float16, pretrained_weights=[ args.pretrained_unet_path, args.pretrained_text_encoder_path, args.pretrained_fp16_vae_path, ], learning_rate=args.learning_rate, use_gradient_checkpointing=args.use_gradient_checkpointing, lora_rank=args.lora_rank, lora_alpha=args.lora_alpha, lora_target_modules=args.lora_target_modules ) launch_training_task(model, args)
4. 加载模型 顾名思义，将前面训练好的模型加载，用于下面生成图像的任务。def load_lora(model, lora_rank, lora_alpha, lora_path): # 配置 LoRA 模块，包括秩、缩放系数、初始化方式等 lora_config = LoraConfig( r=lora_rank, lora_alpha=lora_alpha, init_lora_weights="gaussian", target_modules=["to_q", "to_k", "to_v", "to_out"], ) # 将 LoRA 模块注入到指定的模型中 model = inject_adapter_in_model(lora_config, model) # 加载训练好的 LoRA 权重 state_dict = torch.load(lora_path, map_location="cpu") # 将加载的权重应用到模型中，并返回更新后的模型 model.load_state_dict(state_dict, strict=False) return model# Load models# 管理和加载多个模型组件，比如文本编码器（Text Encoder）、UNet、VAE 等model_manager = ModelManager(torch_dtype=torch.float16, device="cuda", file_path_list=[ "models/kolors/Kolors/text_encoder", "models/kolors/Kolors/unet/diffusion_pytorch_model.safetensors", "models/kolors/Kolors/vae/diffusion_pytorch_model.safetensors" ])pipe = SDXLImagePipeline.from_model_manager(model_manager)# Load LoRA# 加载 LoRA 模块到 UNet 模型# 两者的关系：LoRA 提供了一种高效的方法来微调大型预训练模型，比如 UNet，使其能够更好地适应特定任务或数据集pipe.unet = load_lora( pipe.unet, lora_rank=16, # This parameter should be consistent with that in your training script.（参数需要与训练时的一致，以确保模型的行为一致） lora_alpha=2.0, # lora_alpha can control the weight of LoRA. lora_path="models/lightning_logs/version_0/checkpoints/epoch=0-step=500.ckpt")
5. 生成图像 这里指的是文生图，8张图片的生成方式是一样的，以第一张为例，代码如下：# 设置随机数生成器的种子，保证结果的可重复性torch.manual_seed(0)image = pipe( # 模型生成图像的描述性文本（即正向提示） prompt="水墨画风格，一个黑头发的青年男子，站在窗户前，双手背在身后，很惆怅，全身，白色长袍", # 反向提示，指定了生成图像时需要避免的特征 negative_prompt="丑陋、变形、嘈杂、模糊、低对比度、现代", # 指导生成模型在多大程度上遵循提示的配置参数；值越大，模型越倾向于严格遵循 prompt 提示，可能会导致图像风格更加确定但多样性降低；值较小时，生成的图像可能更具多样性但可能偏离提示 cfg_scale=4, # 生成图像时进行的推理步骤数，数值越高，图像质量通常越高，但生成时间也会增加；指定生成图像的分辨率为 1024x1024 像素 num_inference_steps=50, height=1024, width=1024,)# 保存图片image.save("1.jpg")
6. 展示图像 将多个图像拼接成一张大图，并调整其尺寸。images = [np.array(Image.open(f"{i}.jpg")) for i in range(1, 9)]# 拼接图像，每两张拼接为一行，形成4*2的大图image = np.concatenate([ np.concatenate(images[0:2], axis=1), np.concatenate(images[2:4], axis=1), np.concatenate(images[4:6], axis=1), np.concatenate(images[6:8], axis=1),], axis=0)# 调整大图尺寸为宽1024，高2048image = Image.fromarray(image).resize((1024, 2048))image

三、效果展示

我本来想以某个历史故事为原型，风格是古风，但是我发现给定的训练集几乎都是二次元的图片，我以为效果不会太好，不过还是挺有感觉的。当然我自己爬虫搞了几张古风图片，之后试试看，希望能提高图片质量。

还是展示一张吧！