Diffusers

• FLUX
  • 模型介绍
  • 微调
    • 环境搭建
    • 微调
    • 性能
  • 推理
    • 环境搭建及运行
    • 性能
  • 环境变量声明
• 引用
  • 公网地址说明

FLUX

模型介绍

FLUX.1 dev 是一种基于Rectified Flow Transformers (矫正流) 的生成模型。

• 参考实现：

  url=https://github.com/huggingface/diffusers
  commit_id=a98a839de75f1ad82d8d200c3bc2e4ff89929081
  

微调

环境搭建

【模型开发时推荐使用配套的环境版本】

请参考安装指南

1. 软件与驱动安装

   # 安装 torch 和 torch_npu，参考上述安装指南进行安装
   
   # 将shell脚本中的环境变量路径修改为真实路径，下面为参考路径
   source /usr/local/Ascend/cann/set_env.sh
   

2. 克隆仓库到本地服务器

   git clone --branch 26.0.0 https://gitcode.com/Ascend/MindSpeed-MM.git
   

3. 模型搭建

   3.1 【下载 FLUX GitHub参考实现 在模型根目录下执行以下命令，安装模型对应PyTorch版本需要的依赖】

   git clone https://github.com/huggingface/diffusers.git
   cd diffusers
   git checkout a98a839de75f1ad82d8d200c3bc2e4ff89929081
   cp -r ../MindSpeed-MM/examples/diffusers/flux/* ./examples/dreambooth
   

   【主要代码路径】

   code_path=examples/dreambooth/
   

   3.2【安装其余依赖库】

   pip install -e .
   vim examples/dreambooth/requirements_flux.txt #修改版本： torch==2.7.1, accelerate==1.7.0, transformers==4.47.1, 添加deepspeed==0.17.2, peft==0.7.1
   pip install -r examples/dreambooth/requirements_flux.txt # 安装对应依赖
   

微调

1. 【准备微调数据集】

   • 用户需自行获取并解压pokemon-blip-captions数据集，并在以下启动shell脚本中将dataset_name参数设置为本地数据集的绝对路径

   dataset_name="pokemon-blip-captions" # 数据集 路径
   

   • pokemon-blip-captions数据集格式如下:

   pokemon-blip-captions
   ├── dataset_infos.json
   ├── README.MD
   └── data
         └── train-001.parquet
   

   • 只包含图片的训练数据集，如deepspeed脚本使用训练数据集dog:下载地址，并将dog文件夹转移到examples/dreambooth/目录下

   input_dir="dog" # 数据集路径
   

   dog
   ├── alvan-nee-*****.jpeg
   ├── alvan-nee-*****.jpeg
   

   说明： 该数据集的训练过程脚本只作为一种参考示例。

   • 如用自己的微调数据集，需在shell脚本中将input_dir修改为dataset_name：

   dataset_name="/path/customized_datasets" # 数据集路径
   

   在shell脚本accelerate launch目录下（70行左右）将instance_data_dir=$instance_dir \ 修改为 dataset_name=$dataset_name，并将instance_prompt与validation_prompt改为与自己数据集所匹配的prompt:

   # Example
   accelerate launch --config_file ${config_file} \
     ./train_dreambooth_flux.py \
     --pretrained_model_name_or_path=$model_name  \
     --dataset_name=$dataset_name \
     --instance_prompt="a prompt that is suitable for your own dataset" \
     --validation_prompt="a validation prompt based on your own dataset" \
   

2. 【配置 FLUX 微调脚本】

   联网情况下，微调模型可通过以下步骤下载。无网络时，用户可访问huggingface官网自行下载FLUX.1-dev模型 model_name模型

   export model_name="black-forest-labs/FLUX.1-dev" # 预训练模型路径
   

   获取对应的微调模型后，在以下shell启动微调脚本中将model_name参数设置为本地预训练模型绝对路径

   model_name="black-forest-labs/FLUX.1-dev" # 预训练模型路径
   batch_size=16
   max_train_steps=5000
   mixed_precision="bf16" # 混精
   resolution=256
   config_file="bf16_accelerate_config.yaml"
   
   # accelerate launch --config_file ${config_file} \ 目录下
   --dataloader_num_workers=0 \ # 请基于系统配置与数据大小进行调整num workers
   

3. 【修改代码文件】

   1. 在 src/diffusers/models/embeddings.py 文件里，在 class FluxPosEmbed(nn.Module): 下的 第813行左右 找到代码： freqs_dtype = torch.float32 if is_mps else torch.float64 进行修改, 请修改为：freqs_dtype = torch.float32

      # 修改为freqs_dtype = torch.float32
      vim src/diffusers/models/embeddings.py
      

      freqs_dtype = torch.float32 # 813行附近
      # freqs_dtype = torch.float32 if is_mps else torch.float64 # 原代码
      

   2. 打开train_dreambooth_flux.py或train_dreambooth_lora_flux_advanced.py文件

      cd examples/dreambooth/ # 从diffusers目录进入dreambooth目录
      vim train_dreambooth_flux.py # 进入Python文件
      # 如是flux lora，需先进入advanced_diffusion_training目录
      vim ../advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py # 进入Python文件
      

      • 在import栏if is_wandb_available():上方（62行附近添加代码）

      # 添加代码到train_dreambooth_flux.py 62行附近
      from patch_flux import TorchPatcher, config_gc, create_save_model_hook
      TorchPatcher.apply_patch()
      config_gc()
      
      if is_wandb_available(): # 原代码
        import wandb
      

      • 在log_validation里修改pipeline = pipeline.to(accelerator.device)，train_dreambooth_flux.py在171行附近

      # 修改pipeline为：
      pipeline = pipeline.to(accelerator.device, dtype=torch_dtype)
      # pipeline = pipeline.to(accelerator.device) # 原代码
      

   3. 【Optional】Ubuntu系统需在1701行附近 添加 accelerator.print("")

      if global_step >= args.max_train_steps: # 原代码
        break
      accelerator.print("") # 添加
      

   4. 【Optional】模型checkpoint saving保存

      【因模型较大 如不需要checkpointing_steps，请设置其大于max_train_steps, 避免开启】

      --checkpointing_steps=50000 \ # 修改50000步为所需要步数
      

      【如需保存checkpointing请修改代码】

      vim train_dreambooth_flux.py #（1669行附近）
      vim ../advanced_diffusion_training/train_dreambooth_lora_flux_advanced.py #（2322行附近）
      

      • 在文件上方的import栏增加DistributedType在from accelerate import Accelerator后 （30行附近）
      • 在if accelerator.is_main_process后增加 or accelerator.distributed_type == DistributedType.DEEPSPEED (1669/2322行附近)，并在if args.checkpoints_total_limit is not None后增加and accelerator.is_main_process

      from accelerate import Accelerator, DistributedType
      # from accelerate import Accelerator # 原代码
      
      if accelerator.is_main_process or accelerator.distributed_type == DistributedType.DEEPSPEED:
      # if accelerator.is_main_process: # 原代码
        if global_step % args.checkpointing_steps == 0:  # 原代码 不进行修改
          if args.checkpoints_total_limit is not None and accelerator.is_main_process: # 添加
      

      Lora任务需调用patch任务进行权重保存： 在train_dreambooth_lora_flux_advanced.py文件中找到代码accelerator.register_save_state_pre_hook(save_model_hook)进行修改(1712行附近)，复制粘贴以下代码：

      # 添加
      save_Model_Hook = create_save_model_hook(
            accelerator=accelerator,
            unwrap_model=unwrap_model,
            transformer=transformer,
            text_encoder_one=text_encoder_one,
            args=args,
            weight_dtype=weight_dtype
      )
      accelerator.register_save_state_pre_hook(save_Model_Hook) # 修改
      # accelerator.register_save_state_pre_hook(save_model_hook) # 原代码
      accelerator.register_load_state_pre_hook(load_model_hook) # 原代码 不修改
      

      更改shell脚本：

      export HCCL_CONNECT_TIMEOUT=1200 # 大幅调高HCCL_CONNECT_TIMEOUT (如5000)
      export HCCL_EXEC_TIMEOUT=17000
      --checkpointing_steps=50000 \ # 修改50000步为所需要步数
      

   5. 【Optional】多机运行

      修改config文件

      vim bf16_accelerate_config.yaml
      

      将文件中的deepspeed_multinode_launcher, main_process_ip, 以及main_process_port消除注释而进行使用。

          zero_stage: 2
        #  deepspeed_multinode_launcher: standard
        # main_process_ip: localhost  # 主节点IP
        # main_process_port: 6000     # 主节点port
        machine_rank: 0             # 当前机器的rank
        num_machines: 1             # 总共的机器数
        num_processes: 8            # 总共的卡数
      

      如运行双机：

      • 将两台机器的yaml文件的main_process_ip与main_process_port设置成一样的主节点与port
      • 一台节点machine_rank: 0，另一台machine_rank: 1
      • 两台机器均设置num_machines: 2，num_processes: 16

4. 【启动 FLUX 微调脚本】

   本任务主要提供flux_dreambooth与flux_dreambooth_lora微调脚本，支持多卡训练。

   启动FLUX dreambooth微调脚本

   bash finetune_flux_dreambooth_deepspeed_bf16.sh 
   

   启动FLUX dreambooth_lora微调脚本

   bash finetune_flux_dreambooth_lora_deepspeed_bf16.sh
   

性能

吞吐

FLUX 在 昇腾芯片 和 参考芯片 上的性能对比：

芯片	卡数	任务	FPS	batch_size	AMP_Type	Torch_Version	deepspeed
Atlas 900 A2 PODc	8p	Flux-全参微调	55.23	16	bf16	2.1	✔
竞品A	8p	Flux-全参微调	53.65	16	bf16	2.1	✔

推理

环境搭建及运行

同微调对应章节

cd examples/dreambooth/ # 从diffusers目录进入dreambooth目录

【FLUX模型推理】

vim infer_flux_text2img_bf16.py # 进入运行推理的Python文件

1. 修改路径

   MODEL_PATH = "/black-forest-labs/FLUX.1-dev"  # FLUX模型路径
   

2. 运行代码

   python infer_flux_text2img_bf16.py
   

【DREAMBOOTH微调FLUX模型推理】

vim infer_flux_text2img_dreambooth_bf16.py

1. 修改路径

   MODEL_PATH = "./output_FLUX_dreambooth"  # Dreambooth微调保存模型路径
   

2. 运行代码

   python infer_flux_text2img_dreambooth_bf16.py
   

【lora微调FLUX模型推理】

vim infer_flux_text2img_lora_bf16.py

1. 修改路径

   MODEL_PATH = "./FLUX"  # Flux 模型路径
   LORA_WEIGHTS = "./output/pytorch_lora_weights.safetensors"  # LoRA权重路径
   

2. 运行代码

   python infer_flux_text2img_lora_bf16.py
   

【分布式推理】

vim infer_flux_text2img_distrib.py

• 修改模型权重路径 model_path为模型权重路径或微调后的权重路径

• 如lora微调 可将lora_weights修改为Lora权重路径

  model_path = "/black-forest-labs/FLUX.1-dev"  # 模型权重/微调权重路径
  lora_weights = "/pytorch_lora_weights.safetensors"  # Lora权重路径
  

• 启动分布式推理脚本

  • 因使用accelerate进行分布式推理，config可设置：--num_processes=卡数，num_machines=机器数等

  accelerate launch --num_processes=4 infer_flux_text2img_distrib.py # 单机四卡进行分布式推理
  

性能

芯片	卡数	任务	E2E（it/s）	AMP_Type	Torch_Version
Atlas 900 A2 PODc	8p	文生图	1.16	bf16	2.1
竞品A	8p	文生图	1.82	bf16	2.1
Atlas 900 A2 PODc	8p	文生图微调	1.12	bf16	2.1
竞品A	8p	文生图微调	1.82	bf16	2.1

环境变量声明

环境变量	描述	取值说明
ASCEND_SLOG_PRINT_TO_STDOUT	是否开启日志打印	0: 关闭日志打屏 1: 开启日志打屏
ASCEND_GLOBAL_LOG_LEVEL	设置应用类日志的日志级别及各模块日志级别，仅支持调试日志	0: 对应DEBUG级别 1: 对应INFO级别 2: 对应WARNING级别 3: 对应ERROR级别 4: 对应NULL级别，不输出日志
TASK_QUEUE_ENABLE	用于控制开启task_queue算子下发队列优化的等级	0: 关闭 1: 开启Level 1优化 2: 开启Level 2优化
COMBINED_ENABLE	设置combined标志。设置为0表示关闭此功能；设置为1表示开启，用于优化非连续两个算子组合类场景	0: 关闭 1: 开启
CPU_AFFINITY_CONF	控制CPU端算子任务的处理器亲和性，即设定任务绑核	设置0或未设置: 表示不启用绑核功能 1: 表示开启粗粒度绑核 2: 表示开启细粒度绑核
HCCL_CONNECT_TIMEOUT	用于限制不同设备之间socket建链过程的超时等待时间	需要配置为整数，取值范围[120,7200]，默认值为120，单位s
PYTORCH_NPU_ALLOC_CONF	控制缓存分配器行为	expandable_segments:<value>: 使能内存池扩展段功能，即虚拟内存特征
HCCL_EXEC_TIMEOUT	控制设备间执行时同步等待的时间，在该配置时间内各设备进程等待其他设备执行通信同步	需要配置为整数，取值范围[68,17340]，默认值为1800，单位s
ACLNN_CACHE_LIMIT	配置单算子执行API在Host侧缓存的算子信息条目个数	需要配置为整数，取值范围[1, 10,000,000]，默认值为10000
TOKENIZERS_PARALLELISM	用于控制Hugging Face的transformers库中的分词器（tokenizer）在多线程环境下的行为	False: 禁用并行分词 True: 开启并行分词
OMP_NUM_THREADS	设置执行期间使用的线程数	需要配置为整数

引用

公网地址说明

代码涉及公网地址参考 公网地址