Update README.md

README.md

@@ -36,7 +36,7 @@ from peft import LoraConfig, get_peft_model, PeftModel
 login(token="替换为你自己的",add_to_git_credential=True)
 
 weight_dtype = torch.bfloat16
-train_batch_size = 2
+train_batch_size = 4
 snr_gamma = 5  # SNR 参数，用于信噪比加权损失的调节系数
 # 设置随机数种子以确保可重复性
 seed = 1126  # 随机数种子
@@ -45,25 +45,23 @@ if torch.cuda.is_available():
     torch.cuda.manual_seed_all(seed)
 
 # 优化器参数
-unet_learning_rate = 1e-4  # UNet 的学习率，控制 UNet 参数更新的步长
+unet_learning_rate = 1e-6  # UNet 的学习率，控制 UNet 参数更新的步长
 text_encoder_learning_rate = 1e-4  # 文本编码器的学习率，控制文本嵌入层的参数更新步长
 
 # 学习率调度器参数
 lr_scheduler_name = "cosine_with_restarts"  # 设置学习率调度器为 Cosine annealing with restarts，逐渐减少学习率并定期重启
 lr_warmup_steps = 100  # 学习率预热步数，在最初的 100 步中逐渐增加学习率到最大值
-max_train_steps = 1000  # 总训练步数，决定了整个训练过程的迭代次数
-num_cycles = 3  # Cosine 调度器的周期数量，在训练期间会重复 3 次学习率周期性递减并重启
+max_train_steps = 500  # 总训练步数，决定了整个训练过程的迭代次数
+num_cycles = 1  # Cosine 调度器的周期数量，在训练期间会重复 3 次学习率周期性递减并重启
 
 pretrained_model_name_or_path = "stabilityai/stable-diffusion-2-1"
 
 # LoRA 配置
-lora_config = LoraConfig(
+unet_lora_config = LoraConfig(
     r=32,  # LoRA 的秩，即低秩矩阵的维度，决定了参数调整的自由度
     lora_alpha=16,  # 缩放系数，控制 LoRA 权重对模型的影响
-    target_modules=[
-        # "q_proj", "v_proj", "k_proj", "out_proj",  # 指定 Text encoder 的 LoRA 应用对象（用于调整注意力机制中的投影矩阵）
-        "to_k", "to_q", "to_v", "to_out.0"  # 指定 UNet 的 LoRA 应用对象（用于调整 UNet 中的注意力机制）
-    ],
+    init_lora_weights="gaussian",
+    target_modules=["to_k", "to_q", "to_v", "to_out.0"],
     lora_dropout=0  # LoRA dropout 概率，0 表示不使用 dropout
 )
 
@@ -110,13 +108,13 @@ train_dataloader = DataLoader(dataset, shuffle=True, collate_fn=collate_fn, batc
 
 from diffusers import SD3Transformer2DModel
 
-def prepare_lora_model(lora_config, pretrained_model_name_or_path, model_path=None, resume=False, merge_lora=False):
+def prepare_lora_model(unet_lora_config, pretrained_model_name_or_path, model_path=None, resume=False, merge_lora=False):
     """
     (1) 目标:
         - 加载完整的 Stable Diffusion 模型，包括 LoRA 层，并根据需要合并 LoRA 权重。这包括 Tokenizer、噪声调度器、UNet、VAE 和文本编码器。
 
     (2) 参数:
-        - lora_config: LoraConfig, LoRA 的配置对象
+        - unet_lora_config: LoraConfig, LoRA 的配置对象
         - pretrained_model_name_or_path: str, Hugging Face 上的模型名称或路径
         - model_path: str, 预训练模型的路径
         - resume: bool, 是否从上一次训练中恢复
@@ -130,7 +128,7 @@ def prepare_lora_model(lora_config, pretrained_model_name_or_path, model_path=No
         - text_encoder: CLIPTextModel
     """
     # 加载噪声调度器，用于控制扩散模型的噪声添加和移除过程
-    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+    noise_scheduler = DDIMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
 
     # 加载 Tokenizer，用于将文本标注转换为 tokens
     tokenizer = CLIPTokenizer.from_pretrained(
@@ -157,51 +155,56 @@ def prepare_lora_model(lora_config, pretrained_model_name_or_path, model_path=No
         torch_dtype=weight_dtype,
         subfolder="unet"
     )
+
+    # 冻结 VAE 参数
+    vae.requires_grad_(False)
+    text_encoder.requires_grad_(False)
+    unet.requires_grad_(False)
     
     # 如果设置为继续训练，则加载上一次的模型权重
     if resume:
         if model_path is None or not os.path.exists(model_path):
             raise ValueError("当 resume 设置为 True 时，必须提供有效的 model_path")
         # 使用 PEFT 的 from_pretrained 方法加载 LoRA 模型
-        text_encoder = PeftModel.from_pretrained(text_encoder, os.path.join(model_path, "text_encoder"))
+        # text_encoder = PeftModel.from_pretrained(text_encoder, os.path.join(model_path, "text_encoder"))
         unet = PeftModel.from_pretrained(unet, os.path.join(model_path, "unet"))
 
-        # 确保 UNet 的可训练参数的 requires_grad 为 True
-        for param in unet.parameters():
-            if param.requires_grad is False:
-                param.requires_grad = True
+        # 确保 LoRA 参数是可训练的，仅将指定的模块参数设为可训练
+        target_modules = ["to_k", "to_q", "to_v", "to_out.0"]
+    
+        for name, param in unet.named_parameters():
+            # 只对指定的目标模块设置 requires_grad 为 True
+            if any(target_module in name for target_module in target_modules):
+                param.requires_grad = True  # 仅将 LoRA 参数设为可训练
         
-        # 确保文本编码器的可训练参数的 requires_grad 为 True
-        for param in text_encoder.parameters():
-            if param.requires_grad is False:
-                param.requires_grad = True
                 
         print(f"✅ 已从 {model_path} 恢复模型权重")
 
     else:
-        # 将 LoRA 配置应用到 text_encoder 和 unet
-        # text_encoder = get_peft_model(text_encoder, lora_config)
-        unet = get_peft_model(unet, lora_config)
+
+        # 将 LoRA 配置应用到unet
+        unet.add_adapter(unet_lora_config)
 
         # 打印可训练参数数量
-        # print("📊 Text Encoder 可训练参数:")
-        # text_encoder.print_trainable_parameters()
         print("📊 UNet 可训练参数:")
-        unet.print_trainable_parameters()
+        trainable_params = 0
+        for name, param in unet.named_parameters():
+            if param.requires_grad:
+                param_count = param.numel()  # 计算该参数张量的元素数量
+                trainable_params += param_count
+                # print(f"可训练参数: {name}, 形状: {param.shape}, 参数数量: {param_count}")
+
+        print(f"总的 LoRA 可训练参数数量: {trainable_params}")
     
     if merge_lora:
         # 合并 LoRA 权重到基础模型，仅在推理时调用
-        text_encoder = text_encoder.merge_and_unload()
+        # text_encoder = text_encoder.merge_and_unload()
         unet = unet.merge_and_unload()
 
         # 切换为评估模式
         text_encoder.eval()
         unet.eval()
 
-    # 冻结 VAE 参数
-    vae.requires_grad_(False)
-    text_encoder.requires_grad_(False)
-
     # 将模型移动到 GPU 上并设置权重的数据类型
     unet.to(device, dtype=weight_dtype)
     vae.to(device, dtype=weight_dtype)
@@ -213,13 +216,9 @@ def prepare_optimizer(unet, text_encoder, unet_learning_rate=5e-4, text_encoder_
     # 筛选出 UNet 中需要训练的 Lora 层参数
     unet_lora_layers = [p for p in unet.parameters() if p.requires_grad]
     
-    # 筛选出文本编码器中需要训练的 Lora 层参数
-    text_encoder_lora_layers = [p for p in text_encoder.parameters() if p.requires_grad]
-    
     # 将需要训练的参数分组并设置不同的学习率
     trainable_params = [
         {"params": unet_lora_layers, "lr": unet_learning_rate},
-        {"params": text_encoder_lora_layers, "lr": text_encoder_learning_rate}
     ]
     
     # 使用 AdamW 优化器
@@ -237,7 +236,7 @@ project_name = "fupo"
 dataset_name = "fupo"
 # 根目录和主要目录
 root_dir = "./"  # 当前目录
-main_dir = os.path.join(root_dir, "SD-2-2")  # 主目录
+main_dir = os.path.join(root_dir, "SD-2-1")  # 主目录
 # 项目目录
 project_dir = os.path.join(main_dir, project_name)
 model_path = os.path.join(project_dir, "logs", "checkpoint-last")
@@ -248,7 +247,7 @@ model_path = os.path.join(project_dir, "logs", "checkpoint-last")
 
 # 准备模型
 tokenizer, noise_scheduler, unet, vae, text_encoder = prepare_lora_model(
-    lora_config,
+    unet_lora_config,
     pretrained_model_name_or_path,
     model_path,
     resume=False,
@@ -278,14 +277,18 @@ import math
 from huggingface_hub import HfApi, Repository
 from tqdm.auto import tqdm
 import torch.nn.functional as F
+from peft.utils import get_peft_model_state_dict
+from diffusers.utils import convert_state_dict_to_diffusers
 
+accumulation_steps = 4  # 梯度累积步数
+max_norm = 0.5
 output_folder = os.path.join(project_dir, "logs")
 # 禁用并行化，避免警告
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # 初始化
 global_step = 0
-best_face_score = float("inf")  # 初始化为正无穷大，存储最佳面部相似度分数
+best_loss = float("inf")  # 初始化为正无穷大，存储最佳损失值
 
 # 进度条显示训练进度
 progress_bar = tqdm(
@@ -297,7 +300,6 @@ progress_bar = tqdm(
 for epoch in range(math.ceil(max_train_steps / len(train_dataloader))):
     # 如果你想在训练中增加评估，那在循环中增加 train() 是有必要的
     unet.train()
-    text_encoder.train()
     
     for step, batch in enumerate(train_dataloader):
         if global_step >= max_train_steps:
@@ -313,7 +315,7 @@ for epoch in range(math.ceil(max_train_steps / len(train_dataloader))):
         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # 获取文本的嵌入表示
-        encoder_hidden_states = text_encoder(batch["input_ids"].to(device))[0]
+        encoder_hidden_states = text_encoder(batch["input_ids"].to(device),return_dict=False)[0]
         assert encoder_hidden_states is not None, "Encoder hidden states should not be None"
 
         # 计算目标值
@@ -323,7 +325,8 @@ for epoch in range(math.ceil(max_train_steps / len(train_dataloader))):
             target = noise_scheduler.get_velocity(latents, noise, timesteps)  # 预测速度向量
 
         # UNet 模型预测
-        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states)[0]
+        with torch.autograd.detect_anomaly():
+            model_pred = unet(noisy_latents, timesteps, encoder_hidden_states, return_dict=False)[0]
         assert model_pred is not None, "Model prediction should not be None"
 
         # 计算损失
@@ -345,34 +348,41 @@ for epoch in range(math.ceil(max_train_steps / len(train_dataloader))):
 
         # 反向传播
         loss.backward()
-        optimizer.step()
-        lr_scheduler.step()
-        optimizer.zero_grad()
+        torch.nn.utils.clip_grad_norm_(unet.parameters(), max_norm)
+        # 梯度累积
+        if (global_step + 1) % accumulation_steps == 0:
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
         progress_bar.update(1)
         global_step += 1
 
-        # 打印训练损失
-        if global_step % 100 == 0 or global_step == max_train_steps:
-            print(f"🔥 步骤 {global_step}, 损失: {loss.item()}")
-
-        # 保存中间检查点，当前简单设置为每 500 步保存一次
-        if global_step % 500 == 0:
-            save_path = os.path.join(output_folder, f"checkpoint-{global_step}")
-            os.makedirs(save_path, exist_ok=True)
-
-            # 使用 save_pretrained 保存 PeftModel
-            unet.save_pretrained(os.path.join(save_path, "unet"))
-            print(f"💾 已保存中间模型到 {save_path}")
+        if global_step %100 == 0:
+            # 保存当前损失最低的模型
+            if loss.item() < best_loss:
+                best_loss = loss.item()
+                save_path = os.path.join(output_folder, "best_checkpoint")
+                os.makedirs(save_path, exist_ok=True)
+
+                # 使用 save_pretrained 保存 PeftModel
+                unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
+                StableDiffusionPipeline.save_lora_weights(
+                    save_directory=save_path,
+                    unet_lora_layers=unet_lora_state_dict,
+                    safe_serialization=True,
+                    )
+                # text_encoder.save_pretrained(os.path.join(save_path, "text_encoder"))
+                print(f"💾 损失最小模型已保存到 {save_path}, 当前损失: {best_loss}")
 
 # 保存最终模型到 checkpoint-last
 save_path = os.path.join(output_folder, "checkpoint-last")
 os.makedirs(save_path, exist_ok=True)
-unet.save_pretrained(os.path.join(save_path, "unet"))
-# text_encoder.save_lora_weights(os.path.join(save_path, "text_encoder"))
+unet_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))
+StableDiffusionPipeline.save_lora_weights(
+    save_directory=save_path,
+    unet_lora_layers=unet_lora_state_dict,
+    safe_serialization=True,
+    )
 print(f"💾 已保存最终模型到 {save_path}")
 
-print("🎉 微调完成！")
-
-# 上传到 Hugging Face Hub
-
 ```