Upload 2 files

run.sh

@@ -7,12 +7,11 @@
 # --resume_from_checkpoint="feifei-zimage-lora/checkpoint-100"
 
 accelerate launch --mixed_precision="bf16" train_zimage_lora.py \
-  --pretrained_model_name_or_path="./Z-Image-Turbo" \
-  --train_data_dir="./fiefei_pic" \
+  --pretrained_model_name_or_path="../../../smodels/Z-Image-Turbo" \
+  --train_data_dir="../../../datasets/fiefei_pic" \
   --resolution=1024 \
   --train_batch_size=1 \
   --gradient_accumulation_steps=4 \
-  --max_train_steps=1000 \
   --learning_rate=1e-4 \
   --mixed_precision="bf16" \
   --output_dir="feifei-zimage-lora" \
@@ -21,4 +20,5 @@ accelerate launch --mixed_precision="bf16" train_zimage_lora.py \
   --gradient_checkpointing \
   --use_8bit_adam \
   --checkpointing_steps=100 \
-  --checkpoints_total_limit=3
+  --checkpoints_total_limit=3 \
+  --max_train_steps=200

train_zimage_lora.py

@@ -4,7 +4,7 @@ import argparse
 import logging
 import math
 import os
-from safetensors.torch import save_file  # Add this import
+from safetensors.torch import save_file 
 import random
 import shutil
 import glob
@@ -58,7 +58,7 @@ check_min_version("0.24.0")
 logger = get_logger(__name__, log_level="INFO")
 
 def parse_args():
-    parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser = argparse.ArgumentParser(description="Fixed Training script V3.")
     parser.add_argument("--pretrained_model_name_or_path", type=str, required=True)
     parser.add_argument("--revision", type=str, default=None)
     parser.add_argument("--variant", type=str, default=None)
@@ -242,7 +242,6 @@ def main():
     noise_scheduler = pipe.scheduler
     del pipe
     
-    # Analyze Params
     forward_signature = inspect.signature(transformer.forward)
     params = forward_signature.parameters
     param_names = list(params.keys())
@@ -252,11 +251,14 @@ def main():
     if "x" in param_names: input_arg = "x"
     elif "sample" in param_names: input_arg = "sample"
 
+    time_arg = "t"
+    if "timestep" in param_names: time_arg = "timestep"
+
     cond_arg = "encoder_hidden_states"
     if "cap_feats" in param_names: cond_arg = "cap_feats"
     elif "context" in param_names: cond_arg = "context"
 
-    logger.info(f"Mapping: Input='{input_arg}', Cond='{cond_arg}'")
+    logger.info(f"Mapping: Input='{input_arg}', Time='{time_arg}', Cond='{cond_arg}'")
 
     transformer.requires_grad_(False)
     vae.requires_grad_(False)
@@ -360,25 +362,17 @@ def main():
         transformer, optimizer, train_dataloader, lr_scheduler
     )
 
-    # 计算每个 epoch 需要多少步
     num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
 
-    # === 【核心修复：智能恢复训练状态】 ===
-    # 默认从 0 开始
     global_step = 0
     first_epoch = 0
     resume_step = 0
 
     if args.resume_from_checkpoint:
         logger.info(f"Resuming from checkpoint {args.resume_from_checkpoint}")
-        
-        # 【核心修复 1】总是调用 load_state，这样无论是否使用 accelerate launch 都能工作
         accelerator.load_state(args.resume_from_checkpoint)
         
-        # 【核心修复 2】更可靠的 global_step 恢复逻辑
         checkpoint_path = args.resume_from_checkpoint.rstrip('/')
-        
-        # 方法 1：尝试从文件夹名提取
         folder_name = os.path.basename(checkpoint_path)
         if folder_name.startswith("checkpoint-") and "-" in folder_name:
             try:
@@ -386,7 +380,6 @@ def main():
             except:
                 pass
         
-        # 方法 2：如果无法提取，尝试从 saved_state.json 读取（accelerate 保存的）
         if global_step == 0:
             saved_state_path = os.path.join(checkpoint_path, "saved_state.json")
             if os.path.exists(saved_state_path):
@@ -397,33 +390,16 @@ def main():
                         global_step = saved_state.get("global_step", 0)
                 except:
                     pass
-        
-        # 方法 3：从进度条保存的文件读取（如果有）
-        if global_step == 0:
-            progress_path = os.path.join(checkpoint_path, "progress.json")
-            if os.path.exists(progress_path):
-                try:
-                    import json
-                    with open(progress_path, "r") as f:
-                        progress_data = json.load(f)
-                        global_step = progress_data.get("step", 0)
-                except:
-                    pass
-        
-        logger.info(f"恢复的 Global step: {global_step}")
-        
-        # 计算我们应该从哪个 epoch 和哪一步开始
+
         first_epoch = global_step // num_update_steps_per_epoch
         resume_step = global_step % num_update_steps_per_epoch
         logger.info(f"Resuming from epoch {first_epoch}, step {resume_step}")
     
+    # 训练循环
     for epoch in range(first_epoch, args.num_train_epochs):
         transformer.train()
-        
-        # 定义一个累加器，用来计算平均 Loss
         loss_accumulator = 0.0
 
-        # === 【数据加载器跳过逻辑】 ===
         if args.resume_from_checkpoint and epoch == first_epoch:
             train_dataloader_skip = accelerator.skip_first_batches(train_dataloader, resume_step * args.train_batch_size)
         else:
@@ -431,40 +407,56 @@ def main():
         
         for step, batch in enumerate(train_dataloader_skip):
             with accelerator.accumulate(transformer):
-                # --- 数据准备和模型前向传播 (你的代码保持不变) ---
                 latents = vae.encode(batch["pixel_values"].to(dtype=weight_dtype)).latent_dist.sample()
                 latents = latents * vae.config.scaling_factor
+                
                 noise = torch.randn_like(latents)
                 bsz = latents.shape[0]
                 timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device).long()
+                
                 sigmas = timesteps.flatten() / noise_scheduler.config.num_train_timesteps
                 while sigmas.ndim < latents.ndim:
                     sigmas = sigmas.unsqueeze(-1)
+                
                 noisy_latents = (1.0 - sigmas) * latents + sigmas * noise
                 target = noise - latents 
+                
+                # Z-Image 输入需要 5D
                 noisy_latents = noisy_latents.unsqueeze(2)
+                
                 encoder_hidden_states = batch["encoder_hidden_states"].to(dtype=weight_dtype)
+                
                 forward_kwargs = {
                     input_arg: noisy_latents,
-                    "t": timesteps,
+                    time_arg: timesteps, 
                     cond_arg: encoder_hidden_states,
                     "return_dict": False
                 }
+                
+                if "pooled_projections" in param_names:
+                    forward_kwargs["pooled_projections"] = torch.zeros(
+                        (bsz, 2048), device=latents.device, dtype=weight_dtype
+                    )
+
                 model_pred = transformer(**forward_kwargs)[0]
                 
-                # --- 各种形状修复 (你的代码保持不变) ---
                 if isinstance(model_pred, list): model_pred = model_pred[0]
-                if model_pred.ndim == 5: model_pred = model_pred.squeeze(2)
+                
+                # === 【最后一步核心修复】处理维度颠倒 ===
+                
+                # 1. 如果输出是 5 维的，先把那个 1 维压扁
+                if model_pred.ndim == 5:
+                    model_pred = model_pred.squeeze(2)
+
+                # 2. 如果输出形状和 Target 形状颠倒了 (Channel <-> Batch)，则转置回来
+                # Target: [1, 16, 128, 128]  (Batch, Channel, H, W)
+                # Pred:   [16, 1, 128, 128]  (Channel, Batch, H, W) -> 需要修正
                 if model_pred.shape != target.shape:
                     if model_pred.shape[0] == target.shape[1] and model_pred.shape[1] == target.shape[0]:
+                        # 执行转置，把 Channel 和 Batch 换回来
                         model_pred = model_pred.transpose(0, 1)
-                    elif model_pred.numel() == target.numel():
-                        model_pred = model_pred.reshape(target.shape)
 
-                # --- Loss 计算和累加 ---
                 loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
-                
-                # 【核心修改 1】：把每一小步的 loss 加到累加器里
                 loss_accumulator += loss.detach().item()
 
                 accelerator.backward(loss)
@@ -474,7 +466,6 @@ def main():
                 lr_scheduler.step()
                 optimizer.zero_grad()
 
-            # --- 模型更新后的操作 ---
             if accelerator.sync_gradients:
                 global_step += 1
                 
@@ -483,45 +474,32 @@ def main():
                         save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                         accelerator.save_state(save_path)
 
-                # 【核心修改 2】：只在模型更新时，打印平均 Loss
                 if accelerator.is_main_process:
-                    # 计算过去 N 步的平均 Loss
                     avg_loss = loss_accumulator / args.gradient_accumulation_steps
-                    
-                    # 打印格式化的结果
                     print(f"Steps: {global_step}/{args.max_train_steps} | Loss: {avg_loss:.4f}")
-                    
-                    # 清零累加器，为下一轮做准备
                     loss_accumulator = 0.0
                 
             if global_step >= args.max_train_steps:
                 break
         
-        # === 【核心修改：这里加两行】 ===
-        # 如果步数到了，强制跳出外层的 Epoch 循环，否则它会空转
         if global_step >= args.max_train_steps:
             break
-        # ============================
 
-    # === 这里的缩进是 4 个空格 ===
+    # 保存权重
     if accelerator.is_main_process:
         transformer = accelerator.unwrap_model(transformer)
         transformer = transformer.to(torch.float32)
         transformer_lora_state_dict = convert_state_dict_to_diffusers(get_peft_model_state_dict(transformer))
         
-        # === 【核心修复】给参数名加上 "transformer." 前缀 ===
-        # 这一步非常关键，没有它，load_lora_weights 就认不出来
         new_state_dict = {}
         for k, v in transformer_lora_state_dict.items():
             new_state_dict[f"transformer.{k}"] = v
-        # ================================================
 
-        # 使用 safetensors 保存加上前缀后的权重
         save_path = os.path.join(args.output_dir, "pytorch_lora_weights.safetensors")
-        save_file(new_state_dict, save_path)  # 注意这里传的是 new_state_dict
+        save_file(new_state_dict, save_path)
         logger.info(f"Saved LoRA weights to {save_path}")
 
     accelerator.end_training()
 
 if __name__ == "__main__":
-    main()
+    main()