Upload folder using huggingface_hub (#3)

- 0cf1359d1746adb23074cc2027b559d84c71823a194f6d10f9827ab4de94eea1 (d33c2a18f2b9f6d5a26f7ac83031998695e09a86)

README.md

@@ -0,0 +1,372 @@
+---
+license: mit
+tags:
+- pytorch
+- stable-diffusion
+- text2Image
+- stabilityai/stable-diffusion-2-1
+---
+
+# This LoRA is trained based on stabilityai/stable-diffusion-2-1.
+
+## Training code
+
+```python
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+from datasets import load_dataset
+
+dataset = load_dataset("xchuan/text2image-fupo",split="train")
+
+from transformers import CLIPTokenizer
+from huggingface_hub import login
+# ========== LoRA 模型库 ==========
+from peft import LoraConfig, get_peft_model, PeftModel
+
+
+login(token="替换为你自己的",add_to_git_credential=True)
+
+weight_dtype = torch.bfloat16
+train_batch_size = 2
+snr_gamma = 5  # SNR 参数，用于信噪比加权损失的调节系数
+# 设置随机数种子以确保可重复性
+seed = 1126  # 随机数种子
+torch.manual_seed(seed)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(seed)
+
+# 优化器参数
+unet_learning_rate = 1e-4  # 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 次学习率周期性递减并重启
+
+pretrained_model_name_or_path = "stabilityai/stable-diffusion-2-1"
+
+# LoRA 配置
+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 中的注意力机制）
+    ],
+    lora_dropout=0  # LoRA dropout 概率，0 表示不使用 dropout
+)
+
+from torchvision import transforms
+from torch.utils.data import DataLoader
+
+resolution = 512
+
+
+train_transform = transforms.Compose([
+    transforms.Resize(resolution, interpolation=transforms.InterpolationMode.BILINEAR),  # 调整图像大小
+        transforms.CenterCrop(resolution),  # 中心裁剪图像
+        transforms.RandomHorizontalFlip(),  # 随机水平翻转
+        transforms.ToTensor(),  # 将图像转换为张量
+])
+
+def collate_fn(examples):
+    pixel_values = []
+    input_ids = []
+
+    for example in examples:
+        image_tensor = train_transform(example["image"])
+        if not isinstance(image_tensor, torch.Tensor):
+            print(f"Expected Tensor, got {type(image_tensor)} instead.")
+            continue
+        pixel_values.append(image_tensor)
+
+        input_text = "fupo:" + example["text"]
+        tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer")
+        encode_text = tokenizer(input_text, return_tensors="pt",padding="max_length",truncation=True)
+        inputs_id = encode_text["input_ids"].squeeze(0)
+        input_ids.append(inputs_id)
+    
+        # 如果没有有效的图像，则返回空的字典
+    if len(pixel_values) == 0:
+        return {"pixel_values": torch.empty(0), "input_ids": torch.empty(0)}
+    
+    pixel_values = torch.stack(pixel_values, dim=0).float()
+    input_ids = torch.stack(input_ids, dim=0)
+    return {"pixel_values": pixel_values, "input_ids": input_ids}
+
+
+train_dataloader = DataLoader(dataset, shuffle=True, collate_fn=collate_fn, batch_size=train_batch_size)
+
+from diffusers import SD3Transformer2DModel
+
+def prepare_lora_model(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 的配置对象
+        - pretrained_model_name_or_path: str, Hugging Face 上的模型名称或路径
+        - model_path: str, 预训练模型的路径
+        - resume: bool, 是否从上一次训练中恢复
+        - merge_lora: bool, 是否在推理时合并 LoRA 权重
+
+    (3) 返回:
+        - tokenizer: CLIPTokenizer
+        - noise_scheduler: DDPMScheduler
+        - unet: UNet2DConditionModel
+        - vae: AutoencoderKL
+        - text_encoder: CLIPTextModel
+    """
+    # 加载噪声调度器，用于控制扩散模型的噪声添加和移除过程
+    noise_scheduler = DDPMScheduler.from_pretrained(pretrained_model_name_or_path, subfolder="scheduler")
+
+    # 加载 Tokenizer，用于将文本��注转换为 tokens
+    tokenizer = CLIPTokenizer.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="tokenizer"
+    )
+
+    # 加载 CLIP 文本编码器，用于将文本标注转换为特征向量
+    text_encoder = CLIPTextModel.from_pretrained(
+        pretrained_model_name_or_path,
+        torch_dtype=weight_dtype,
+        subfolder="text_encoder"
+    )
+
+    # 加载 VAE 模型，用于在扩散模型中处理图像的潜在表示
+    vae = AutoencoderKL.from_pretrained(
+        pretrained_model_name_or_path,
+        subfolder="vae"
+    )
+
+    # 加载 UNet 模型，负责处理扩散模型中的图像生成和推理过程
+    unet = UNet2DConditionModel.from_pretrained(
+        pretrained_model_name_or_path,
+        torch_dtype=weight_dtype,
+        subfolder="unet"
+    )
+    
+    # 如果设置为继续训练，则加载上一次的模型权重
+    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"))
+        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
+        
+        # 确保文本编码器的可训练参数的 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)
+
+        # 打印可训练参数数量
+        # print("📊 Text Encoder 可训练参数:")
+        # text_encoder.print_trainable_parameters()
+        print("📊 UNet 可训练参数:")
+        unet.print_trainable_parameters()
+    
+    if merge_lora:
+        # 合并 LoRA 权重到基础模型，仅在推理时调用
+        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)
+    text_encoder.to(device, dtype=weight_dtype)
+    
+    return tokenizer, noise_scheduler, unet, vae, text_encoder
+
+def prepare_optimizer(unet, text_encoder, unet_learning_rate=5e-4, text_encoder_learning_rate=1e-4):
+    # 筛选出 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 优化器
+    optimizer = torch.optim.AdamW(trainable_params)
+    
+    return optimizer
+
+import os
+from diffusers.optimization import get_scheduler
+from diffusers.training_utils import compute_snr
+from diffusers import DDPMScheduler,AutoencoderKL,UNet2DConditionModel
+from transformers import CLIPTextModel
+
+project_name = "fupo"
+dataset_name = "fupo"
+# 根目录和主要目录
+root_dir = "./"  # 当前目录
+main_dir = os.path.join(root_dir, "SD-2-2")  # 主目录
+# 项目目录
+project_dir = os.path.join(main_dir, project_name)
+model_path = os.path.join(project_dir, "logs", "checkpoint-last")
+
+# 项目目录
+project_dir = os.path.join(main_dir, project_name)
+model_path = os.path.join(project_dir, "logs", "checkpoint-last")
+
+# 准备模型
+tokenizer, noise_scheduler, unet, vae, text_encoder = prepare_lora_model(
+    lora_config,
+    pretrained_model_name_or_path,
+    model_path,
+    resume=False,
+    merge_lora=False
+)
+
+# 准备优化器
+optimizer = prepare_optimizer(
+    unet, 
+    text_encoder, 
+    unet_learning_rate=unet_learning_rate, 
+    text_encoder_learning_rate=text_encoder_learning_rate
+)
+
+# 设置学习率调度器
+lr_scheduler = get_scheduler(
+    lr_scheduler_name,
+    optimizer=optimizer,
+    num_warmup_steps=lr_warmup_steps,
+    num_training_steps=max_train_steps,
+    num_cycles=num_cycles
+)
+
+print("✅ 模型和优化器准备完成！可以开始训练。")
+
+import math
+from huggingface_hub import HfApi, Repository
+from tqdm.auto import tqdm
+import torch.nn.functional as F
+
+output_folder = os.path.join(project_dir, "logs")
+# 禁用并行化，避免警告
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# 初始化
+global_step = 0
+best_face_score = float("inf")  # 初始化为正无穷大，存储最佳面部相似度分数
+
+# 进度条显示训练进度
+progress_bar = tqdm(
+    range(max_train_steps),  # 根据 num_training_steps 设置
+    desc="训练步骤",
+)
+
+# 训练循环
+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:
+            break
+        
+        # 编码图像为潜在表示（latent）
+        latents = vae.encode(batch["pixel_values"].to(device, dtype=weight_dtype)).latent_dist.sample()
+        latents = latents * vae.config.scaling_factor  # 根据 VAE 的缩放因子调整潜在空间
+
+        # 为潜在表示添加噪声，生成带噪声的图像
+        noise = torch.randn_like(latents)  # 生成与潜在表示相同形状的随机噪声
+        timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (latents.shape[0],), device=device).long()
+        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+
+        # 获取文本的嵌入表示
+        encoder_hidden_states = text_encoder(batch["input_ids"].to(device))[0]
+        assert encoder_hidden_states is not None, "Encoder hidden states should not be None"
+
+        # 计算目标值
+        if noise_scheduler.config.prediction_type == "epsilon":
+            target = noise  # 预测噪声
+        elif noise_scheduler.config.prediction_type == "v_prediction":
+            target = noise_scheduler.get_velocity(latents, noise, timesteps)  # 预测速度向量
+
+        # UNet 模型预测
+        model_pred = unet(noisy_latents, timesteps, encoder_hidden_states)[0]
+        assert model_pred is not None, "Model prediction should not be None"
+
+        # 计算损失
+        if not snr_gamma:
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
+        else:
+            # 计算信噪比 (SNR) 并根据 SNR 加权 MSE 损失
+            snr = compute_snr(noise_scheduler, timesteps)
+            mse_loss_weights = torch.stack([snr, snr_gamma * torch.ones_like(timesteps)], dim=1).min(dim=1)[0]
+            if noise_scheduler.config.prediction_type == "epsilon":
+                mse_loss_weights = mse_loss_weights / snr
+            elif noise_scheduler.config.prediction_type == "v_prediction":
+                mse_loss_weights = mse_loss_weights / (snr + 1)
+            
+            # 计算加权的 MSE 损失
+            loss = F.mse_loss(model_pred.float(), target.float(), reduction="none")
+            loss = loss.mean(dim=list(range(1, len(loss.shape)))) * mse_loss_weights
+            loss = loss.mean()
+
+        # 反向传播
+        loss.backward()
+        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"))
+            text_encoder.save_pretrained(os.path.join(save_path, "text_encoder"))
+            print(f"💾 已保存中间模型到 {save_path}")
+
+# 保存最终模型到 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"))
+print(f"💾 已保存最终模型到 {save_path}")
+
+print("🎉 微调完成！")
+
+# 上传到 Hugging Face Hub
+
+```