Diffusers

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

FLUX

模型介绍

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

• 参考实现：

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

微调

环境搭建

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

请参考安装指南

1. 软件与驱动安装

   # python3.10
   conda create -n test python=3.10
   conda activate test
   
   # 安装 torch 和 torch_npu，注意要选择对应python版本、x86或arm的torch、torch_npu及apex包
   pip install torch-2.7.1-cp310-cp310-manylinux_2_28_aarch64.whl
   pip install torch_npu-2.7.1*-cp310-cp310-manylinux_2_28_aarch64.whl
   
   # apex for Ascend 参考 https://gitee.com/ascend/apex
   # 建议从原仓编译安装
   
   # 将shell脚本中的环境变量路径修改为真实路径，下面为参考路径
   source /usr/local/Ascend/ascend-toolkit/set_env.sh
   

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

   git clone https://gitee.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 #修改版本：torchvision==0.16.0, torch==2.7.1, accelerate==0.33.0, transformers==4.47.1, 添加deepspeed==0.15.2
   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="pretrain_${mixed_precision}_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 Acceleratore后 （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 在 昇腾芯片 和 参考芯片 上的性能对比：

推理

环境搭建及运行

同微调对应章节

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 # 单机四卡进行分布式推理
  

性能

环境变量声明

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表示开启，用于优化非连续两个算子组合类场景
HCCL_WHITELIST_DISABLE： 配置在使用HCCL时是否开启通信白名单，0：开启白名单，1：关闭白名单
CPU_AFFINITY_CONF： 控制CPU端算子任务的处理器亲和性，即设定任务绑核，设置0或未设置：表示不启用绑核功能， 1：表示开启粗粒度绑核， 2：表示开启细粒度绑核
HCCL_CONNECT_TIMEOUT: 用于限制不同设备之间socket建链过程的超时等待时间，需要配置为整数，取值范围[120,7200]，默认值为120，单位s
ACLNN_CACHE_LIMIT： 配置单算子执行API在Host侧缓存的算子信息条目个数
TOKENIZERS_PARALLELISM： 用于控制Hugging Face的transformers库中的分词器（tokenizer）在多线程环境下的行为
PYTORCH_NPU_ALLOC_CONF： 控制缓存分配器行为
OMP_NUM_THREADS： 设置执行期间使用的线程数

引用

公网地址说明

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