diff --git a/diffsynth/diffusion/base_pipeline.py b/diffsynth/diffusion/base_pipeline.py index 8961284..2bec693 100644 --- a/diffsynth/diffusion/base_pipeline.py +++ b/diffsynth/diffusion/base_pipeline.py @@ -178,15 +178,26 @@ class BasePipeline(torch.nn.Module): def get_vram(self): return torch.cuda.mem_get_info(self.device)[1] / (1024 ** 3) + def get_module(self, model, name): + if "." in name: + name, suffix = name[:name.index(".")], name[name.index(".") + 1:] + if name.isdigit(): + return self.get_module(model[int(name)], suffix) + else: + return self.get_module(getattr(model, name), suffix) + else: + return getattr(model, name) def freeze_except(self, model_names): - for name, model in self.named_children(): - if name in model_names: - model.train() - model.requires_grad_(True) - else: - model.eval() - model.requires_grad_(False) + self.eval() + self.requires_grad_(False) + for name in model_names: + module = self.get_module(self, name) + if module is None: + print(f"No {name} models in the pipeline. We cannot enable training on the model. If this occurs during the data processing stage, it is normal.") + continue + module.train() + module.requires_grad_(True) def blend_with_mask(self, base, addition, mask): diff --git a/docs/Model_Details/Qwen-Image.md b/docs/Model_Details/Qwen-Image.md index 1671e67..717b4e0 100644 --- a/docs/Model_Details/Qwen-Image.md +++ b/docs/Model_Details/Qwen-Image.md @@ -2,7 +2,55 @@ ![Image](https://github.com/user-attachments/assets/738078d8-8749-4a53-a046-571861541924) -Qwen-Image 是由阿里巴巴通义实验室开源的图像生成模型。 +Qwen-Image 是由阿里巴巴通义实验室通义千问团队训练并开源的图像生成模型。 + +## 安装 + +在使用本项目进行模型推理和训练前，请先安装 DiffSynth-Studio。 + +```shell +git clone https://github.com/modelscope/DiffSynth-Studio.git +cd DiffSynth-Studio +pip install -e . +``` + +更多关于安装的信息，请参考[安装依赖](/docs/Pipeline_Usage/Setup.md)。 + +## 快速开始 + +运行以下代码可以快速加载 [Qwen/Qwen-Image](https://www.modelscope.cn/models/Qwen/Qwen-Image) 模型并进行推理。显存管理已启动，框架会自动根据剩余显存控制模型参数的加载，最低 8G 显存即可运行。 + +```python +from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig +import torch + +vram_config = { + "offload_dtype": "disk", + "offload_device": "disk", + "onload_dtype": torch.float8_e4m3fn, + "onload_device": "cpu", + "preparing_dtype": torch.float8_e4m3fn, + "preparing_device": "cuda", + "computation_dtype": torch.bfloat16, + "computation_device": "cuda", +} +pipe = QwenImagePipeline.from_pretrained( + torch_dtype=torch.bfloat16, + device="cuda", + model_configs=[ + ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors", **vram_config), + ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors", **vram_config), + ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **vram_config), + ], + tokenizer_config=ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"), + vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, +) +prompt = "精致肖像，水下少女，蓝裙飘逸，发丝轻扬，光影透澈，气泡环绕，面容恬静，细节精致，梦幻唯美。" +image = pipe(prompt, seed=0, num_inference_steps=40) +image.save("image.jpg") +``` + +## 模型总览
@@ -30,35 +78,6 @@ graph LR;
-## 快速开始 - -通过运行以下代码可以快速加载 [Qwen/Qwen-Image](https://www.modelscope.cn/models/Qwen/Qwen-Image) 模型并进行推理 - -```python -from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig -from PIL import Image -import torch - -pipe = QwenImagePipeline.from_pretrained( - torch_dtype=torch.bfloat16, - device="cuda", - model_configs=[ - ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"), - ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors"), - ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), - ], - tokenizer_config=ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"), -) -prompt = "精致肖像，水下少女，蓝裙飘逸，发丝轻扬，光影透澈，气泡环绕，面容恬静，细节精致，梦幻唯美。" -image = pipe( - prompt, seed=0, num_inference_steps=40, - # edit_image=Image.open("xxx.jpg").resize((1328, 1328)) # For Qwen-Image-Edit -) -image.save("image.jpg") -``` - -## 模型总览 - |模型 ID|推理|低显存推理|全量训练|全量训练后验证|LoRA 训练|LoRA 训练后验证| |-|-|-|-|-|-|-| |[Qwen/Qwen-Image](https://www.modelscope.cn/models/Qwen/Qwen-Image)|[code](/examples/qwen_image/model_inference/Qwen-Image.py)|[code](/examples/qwen_image/model_inference_low_vram/Qwen-Image.py)|[code](/examples/qwen_image/model_training/full/Qwen-Image.sh)|[code](/examples/qwen_image/model_training/validate_full/Qwen-Image.py)|[code](/examples/qwen_image/model_training/lora/Qwen-Image.sh)|[code](/examples/qwen_image/model_training/validate_lora/Qwen-Image.py)| @@ -75,6 +94,13 @@ image.save("image.jpg") |[DiffSynth-Studio/Qwen-Image-In-Context-Control-Union](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-In-Context-Control-Union)|[code](/examples/qwen_image/model_inference/Qwen-Image-In-Context-Control-Union.py)|[code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-In-Context-Control-Union.py)|-|-|[code](/examples/qwen_image/model_training/lora/Qwen-Image-In-Context-Control-Union.sh)|[code](/examples/qwen_image/model_training/validate_lora/Qwen-Image-In-Context-Control-Union.py)| |[DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix](https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Edit-Lowres-Fix)|[code](/examples/qwen_image/model_inference/Qwen-Image-Edit-Lowres-Fix.py)|[code](/examples/qwen_image/model_inference_low_vram/Qwen-Image-Edit-Lowres-Fix.py)|-|-|-|-| +特殊训练脚本： + +* 差分 LoRA 训练：[doc](/docs/Training/Differential_LoRA.md)、[code](/examples/qwen_image/model_training/special/differential_training/) +* FP8 精度训练：[doc](/docs/Training/FP8_Precision.md)、[code](/examples/qwen_image/model_training/special/fp8_training/) +* 两阶段拆分训练：[doc](/docs/Training/Split_Training.md)、[code](/examples/qwen_image/model_training/special/split_training/) +* 端到端直接蒸馏：[doc](/docs/Training/Direct_Distill.md)、[code](/examples/qwen_image/model_training/lora/Qwen-Image-Distill-LoRA.sh) + ## 模型推理 模型通过 `QwenImagePipeline.from_pretrained` 加载，详见[加载模型](/docs/Pipeline_Usage/Model_Inference.md#加载模型)。 @@ -108,15 +134,58 @@ image.save("image.jpg") * `tile_stride`: VAE 编解码阶段的分块步长，默认为 64，仅在 `tiled=True` 时生效，需保证其数值小于或等于 `tile_size`。 * `progress_bar_cmd`: 进度条，默认为 `tqdm.tqdm`。可通过设置为 `lambda x:x` 来屏蔽进度条。 -如果显存不足，请开启[显存管理](/docs/Pipeline_Usage/VRAM_management.md)。 +如果显存不足，请开启[显存管理](/docs/Pipeline_Usage/VRAM_management.md)，我们在示例代码中提供了每个模型推荐的低显存配置，详见前文“模型总览”中的表格。 ## 模型训练 -模型训练脚本位于 `examples/qwen_image/model_training/train.py`，脚本的输入参数包括[基础脚本参数](/docs/Pipeline_Usage/Model_Training.md#脚本参数)以及以下额外参数： +Qwen-Image 系列模型统一通过 [`examples/qwen_image/model_training/train.py`](/examples/qwen_image/model_training/train.py) 进行训练，脚本的参数包括： -* `--tokenizer_path`: tokenizer 的路径，适用于文生图模型，留空则自动从远程下载。 -* `--processor_path`: processor 的路径，适用于图像编辑模型，留空则自动从远程下载。 +* 通用训练参数 + * 数据集基础配置 + * `--dataset_base_path`: 数据集的根目录。 + * `--dataset_metadata_path`: 数据集的元数据文件路径。 + * `--dataset_repeat`: 每个 epoch 中数据集重复的次数。 + * `--dataset_num_workers`: 每个 Dataloder 的进程数量。 + * `--data_file_keys`: 元数据中需要加载的字段名称，通常是图像或视频文件的路径，以 `,` 分隔。 + * 模型加载配置 + * `--model_paths`: 要加载的模型路径。JSON 格式。 + * `--model_id_with_origin_paths`: 带原始路径的模型 ID，例如 `"Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors"`。用逗号分隔。 + * `--extra_inputs`: 模型 Pipeline 所需的额外输入参数，例如训练图像编辑模型 Qwen-Image-Edit 时需要额外参数 `edit_image`，以 `,` 分隔。 + * `--fp8_models`：以 FP8 格式加载的模型，格式与 `--model_paths` 或 `--model_id_with_origin_paths` 一致，目前仅支持参数不被梯度更新的模型（不需要梯度回传，或梯度仅更新其 LoRA）。 + * 训练基础配置 + * `--learning_rate`: 学习率。 + * `--num_epochs`: 轮数（Epoch）。 + * `--trainable_models`: 可训练的模型，例如 `dit`、`vae`、`text_encoder`。 + * `--find_unused_parameters`: DDP 训练中是否存在未使用的参数，少数模型包含不参与梯度计算的冗余参数，需开启这一设置避免在多 GPU 训练中报错。 + * `--weight_decay`：权重衰减大小，详见 [torch.optim.AdamW](https://docs.pytorch.org/docs/stable/generated/torch.optim.AdamW.html)。 + * `--task`: 训练任务，默认为 `sft`，部分模型支持更多训练模式，请参考每个特定模型的文档。 + * 输出配置 + * `--output_path`: 模型保存路径。 + * `--remove_prefix_in_ckpt`: 在模型文件的 state dict 中移除前缀。 + * `--save_steps`: 保存模型的训练步数间隔，若此参数留空，则每个 epoch 保存一次。 + * LoRA 配置 + * `--lora_base_model`: LoRA 添加到哪个模型上。 + * `--lora_target_modules`: LoRA 添加到哪些层上。 + * `--lora_rank`: LoRA 的秩（Rank）。 + * `--lora_checkpoint`: LoRA 检查点的路径。如果提供此路径，LoRA 将从此检查点加载。 + * `--preset_lora_path`: 预置 LoRA 检查点路径，如果提供此路径，这一 LoRA 将会以融入基础模型的形式加载。此参数用于 LoRA 差分训练。 + * `--preset_lora_model`: 预置 LoRA 融入的模型，例如 `dit`。 + * 梯度配置 + * `--use_gradient_checkpointing`: 是否启用 gradient checkpointing。 + * `--use_gradient_checkpointing_offload`: 是否将 gradient checkpointing 卸载到内存中。 + * `--gradient_accumulation_steps`: 梯度累积步数。 + * 图像宽高配置（适用于图像生成模型和视频生成模型） + * `--height`: 图像或视频的高度。将 `height` 和 `width` 留空以启用动态分辨率。 + * `--width`: 图像或视频的宽度。将 `height` 和 `width` 留空以启用动态分辨率。 + * `--max_pixels`: 图像或视频帧的最大像素面积，当启用动态分辨率时，分辨率大于这个数值的图片都会被缩小，分辨率小于这个数值的图片保持不变。 +* Qwen-Image 专有参数 + * `--tokenizer_path`: tokenizer 的路径，适用于文生图模型，留空则自动从远程下载。 + * `--processor_path`: processor 的路径，适用于图像编辑模型，留空则自动从远程下载。 -`--task` 参数支持 `sft`（[标准监督训练](/docs/Training/Supervised_Fine_Tuning.md)）与 `direct_distill`（[直接蒸馏](/docs/Training/Direct_Distill.md)），两者都支持[两阶段拆分训练](/docs/Training/Split_Training.md)和[FP8 精度](/docs/Training/FP8_Precision.md)。 +我们构建了一个样例图像数据集，以方便您进行测试，通过以下命令可以下载这个数据集： -使用命令 `modelscope download --dataset DiffSynth-Studio/example_image_dataset --local_dir ./data/example_image_dataset` 可下载样例数据集。我们为每个模型编写了推荐的训练命令，详见[模型总览](#模型总览)中的表格。详细的训练流程，请参考[模型训练](/docs/Pipeline_Usage/Model_Training.md)。 +```shell +modelscope download --dataset DiffSynth-Studio/example_image_dataset --local_dir ./data/example_image_dataset +``` + +我们为每个模型编写了推荐的训练脚本，请参考前文“模型总览”中的表格。关于如何编写模型训练脚本，请参考[模型训练](/docs/Pipeline_Usage/Model_Training.md)；更多高阶训练算法，请参考[训练框架详解](/docs/Training/)。 diff --git a/docs/Pipeline_Usage/Environment_Variables.md b/docs/Pipeline_Usage/Environment_Variables.md index be222b8..3dfea66 100644 --- a/docs/Pipeline_Usage/Environment_Variables.md +++ b/docs/Pipeline_Usage/Environment_Variables.md @@ -33,3 +33,7 @@ DIFFSYNTH_MODEL_BASE_PATH="./path_to_my_models" python xxx.py ## `DIFFSYNTH_DISK_MAP_BUFFER_SIZE` 硬盘直连中的 Buffer 大小，默认是 1B（1000000000），数值越大，占用内存越大，速度越快。 + +## `DIFFSYNTH_DOWNLOAD_RESOURCE` + +远程模型下载源，可设置为 `modelscope` 或 `huggingface`，控制模型下载的来源，默认值为 `modelscope`。 diff --git a/docs/Pipeline_Usage/Model_Inference.md b/docs/Pipeline_Usage/Model_Inference.md index 24bee13..e006eed 100644 --- a/docs/Pipeline_Usage/Model_Inference.md +++ b/docs/Pipeline_Usage/Model_Inference.md @@ -38,6 +38,14 @@ pipe = QwenImagePipeline.from_pretrained( > > 默认情况下，即使模型已经下载完毕，程序仍会向远程查询是否有遗漏文件，如果要完全关闭远程请求，请将[环境变量 DIFFSYNTH_SKIP_DOWNLOAD](/docs/Pipeline_Usage/Environment_Variables.md#diffsynth_skip_download) 设置为 `True`。 +如需从 [HuggingFace](https://huggingface.co/) 下载模型，可通过设置[环境变量](Environment_Variables.md)实现： + +```shell +import os +os.environ["DIFFSYNTH_DOWNLOAD_RESOURCE"] = "huggingface" +import diffsynth +``` +
diff --git a/examples/z_image/model_training/lora/Z-Image-Turbo.sh b/examples/z_image/model_training/lora/Z-Image-Turbo.sh index a00d57e..4f539b4 100644 --- a/examples/z_image/model_training/lora/Z-Image-Turbo.sh +++ b/examples/z_image/model_training/lora/Z-Image-Turbo.sh @@ -13,3 +13,27 @@ accelerate launch examples/z_image/model_training/train.py \ --lora_rank 32 \ --use_gradient_checkpointing \ --dataset_num_workers 8 + + +# Z-Image-Turbo is a distilled model. +# After training, it loses its distillation-based acceleration capability, +# leading to degraded generation quality at fewer inference steps. +# This issue can be mitigated by using a pre-trained LoRA model to assist the training process. + +# accelerate launch examples/z_image/model_training/train.py \ +# --dataset_base_path data/example_image_dataset \ +# --dataset_metadata_path data/example_image_dataset/metadata.csv \ +# --max_pixels 1048576 \ +# --dataset_repeat 50 \ +# --model_id_with_origin_paths "Tongyi-MAI/Z-Image-Turbo:transformer/*.safetensors,Tongyi-MAI/Z-Image-Turbo:text_encoder/*.safetensors,Tongyi-MAI/Z-Image-Turbo:vae/diffusion_pytorch_model.safetensors" \ +# --learning_rate 1e-4 \ +# --num_epochs 5 \ +# --remove_prefix_in_ckpt "pipe.dit." \ +# --output_path "./models/train/Z-Image-Turbo_lora" \ +# --lora_base_model "dit" \ +# --lora_target_modules "to_q,to_k,to_v,to_out.0,w1,w2,w3" \ +# --lora_rank 32 \ +# --preset_lora_path "models/ostris/zimage_turbo_training_adapter/zimage_turbo_training_adapter_v1.safetensors" \ +# --preset_lora_model "dit" \ +# --use_gradient_checkpointing \ +# --dataset_num_workers 8