qwen edit 2509 训练lora效果很差

[hello-program]

你好，我参考了 https://github.com/inclusionAI/TwinFlow/tree/main/src 的lora使用介绍去蒸馏qwen edit，但是效果很差,能否帮忙看看是哪里的问题呀，谢谢
这是我的推理代码：

import os
import sys
sys.path.append("/home/workspace/code/TwinFlow/src")
import torch
from functools import partial
from torchvision.utils import save_image
from PIL import Image
from torchvision import transforms

from src.networks.qwen_image.modeling_qwen_image import QwenImage
from peft import PeftModel
from unified_sampler import UnifiedSampler
from safetensors.torch import load_file
from torch.amp import autocast as torch_autocast

seed = 42
torch.manual_seed(seed)
device = torch.device("cuda")

base_model_path = "/home/workspace/hf_model/Qwen-Image-Edit-2509"
lora_checkpoint_path = "../outputs/qwenimage_task/qwenimage_lora_2order/checkpoints/global_step_30/model"

input_image_path = "/home/workspace/code/DiffSynth-Studio/example_image_dataset/edit/image1_F.jpg"
prompt = "将裙子改为粉色"

height = 512
width = 512

input_image = Image.open(input_image_path).convert("RGB").resize((width, height), Image.LANCZOS)
transform = transforms.Compose([
    transforms.Resize(min(height, width), interpolation=transforms.InterpolationMode.BILINEAR),
    transforms.ToTensor(),
    transforms.Normalize([0.5], [0.5]),
])
input_tensor = transform(input_image).unsqueeze(0).to(device, dtype=torch.bfloat16)

print("load base model...")
model = QwenImage(
    base_model_path,
    model_type="edit",
    aux_time_embed=False,
    text_dtype=torch.bfloat16,
    imgs_dtype=torch.bfloat16,
    device=device
).to(device)

print("load lora...")
base_transformer = model.model.transformer

if lora_checkpoint_path and os.path.exists(lora_checkpoint_path):
    print(f"Loading LoRA from {lora_checkpoint_path}...")
    
    lora_transformer = PeftModel.from_pretrained(
        base_transformer,
        lora_checkpoint_path,
        adapter_name="default",
        is_trainable=False
    )
    
    adapter_path = os.path.join(lora_checkpoint_path, "adapter_model.safetensors")
    raw_state_dict = load_file(adapter_path)
    final_state_dict = {}
    
    for key, value in raw_state_dict.items():
        new_key = key
        if "base_model.model.transformer." in new_key:
            new_key = new_key.replace("base_model.model.transformer.", "base_model.model.")
        if "lora_A.weight" in new_key and ".default.weight" not in new_key:
            new_key = new_key.replace("lora_A.weight", "lora_A.default.weight")
        elif "lora_B.weight" in new_key and ".default.weight" not in new_key:
            new_key = new_key.replace("lora_B.weight", "lora_B.default.weight")
        elif "lora_embedding_A" in new_key and ".default.weight" not in new_key:
            new_key = new_key.replace("lora_embedding_A.weight", "lora_embedding_A.default.weight")
        elif "lora_embedding_B" in new_key and ".default.weight" not in new_key:
            new_key = new_key.replace("lora_embedding_B.weight", "lora_embedding_B.default.weight")
        
        final_state_dict[new_key] = value
    
    missing, unexpected = lora_transformer.load_state_dict(final_state_dict, strict=False)
    

    real_missing = [k for k in missing if "lora" in k]
    real_unexpected = [k for k in unexpected if "lora" in k]
    
    if len(real_missing) > 0:
        print("missing")
    else:
        print("success！")

    lora_transformer.set_adapter("default")
    lora_transformer.to(device, dtype=torch.bfloat16)
    lora_transformer.eval()
    
    # 5. 替换模型中的 transformer
    model.model.transformer = lora_transformer
    model.transformer.transformer = lora_transformer
else:
    print("can not find lora")
model.eval()

sampler_config = {
    "sampling_steps": 4,
    "stochast_ratio": 1.0,
    "extrapol_ratio": 0.0,
    "sampling_order": 1,
    "time_dist_ctrl": [1.0, 1.0, 1.0],
    "rfba_gap_steps": [0.001, 0.5],
}
sampler = partial(UnifiedSampler().sampling_loop, **sampler_config)

with (
    torch.no_grad(),
    torch_autocast(enabled=True, dtype=torch.bfloat16, device_type="cuda"),
):
        edited_image = model.sample(
            prompts=[prompt],
            images=input_tensor,
            cfg_scale=0.0, 
            seed=seed,
            height=height,
            width=width,
            sampler=sampler,
            return_traj=False,
        )

save_image((edited_image.squeeze(0) + 1) / 2, "edited_output.jpg")
print("done")

我保存模型的实现

def save_ckpt(
    ckpt_root_dir,
    model_to_save,
    global_step,
):
    model_dir = os.path.join(ckpt_root_dir, f"global_step_{global_step}", "model")
    os.makedirs(model_dir, exist_ok=True)

    if hasattr(model_to_save.transformer, 'module'):
        model_to_save.transformer.module.save_pretrained(
            model_dir,
            safe_serialization=True  
        )
    else:
        model_to_save.transformer.save_pretrained(
            model_dir,
            safe_serialization=True
        )

左图是蒸馏前的图，右边是蒸馏后的，我只训练了一张图，跑了120个step，无法过拟合，配置文件没咋改，只是改了这2个地方

  model_name: QwenImageEdit #Flux
  aux_time_embed: false
  lora_rank: 64
  lora_alpha: 64

[QAQdev]

@hello-program 您好，可以看下配置文件吗，以及使用的是哪个method文件（twinflow or twinflow_lora)

我目前观察到的你这里的可能问题是如果用的是lora，且t-tt/2这种方式的话，是不能用from unified_sampler import UnifiedSampler这里的sampler的，因为这里是给全参训练的用的，没有t-tt/2的逻辑

[QAQdev]

TwinFlow/src/methodes/twinflow/twinflow_lora.py

Lines 399 to 400 in 0c53373

	# NOTE: when using lora, you could use a single t-tt/2 to approximate 2 time embedders
	F_t = (-1) ** (1 - self.integ_st) * model(x_t, t=q - tt / 2, **model_kwargs)

[hello-program]

@hello-program 您好，可以看下配置文件吗，以及使用的是哪个method文件（twinflow or twinflow_lora)

我目前观察到的你这里的可能问题是如果用的是lora，且t-tt/2这种方式的话，是不能用from unified_sampler import UnifiedSampler这里的sampler的，因为这里是给全参训练的用的，没有t-tt/2的逻辑

感谢回复，我使用的是lora，配置文件如下：

model:
  model_name: QwenImageEdit #Flux
  aux_time_embed: false
  # model_path: /dllab/share/data/zhenglin/model_zoo/FLUX.1-dev
  model_path: /home/workspace/hf_model/Qwen-Image-Edit-2509
  lora_rank: 64
  lora_alpha: 64
  lora_target_modules: to_k,to_q,to_v,to_out.0,img_mlp.net.0.proj,img_mlp.net.2,add_k_proj,add_q_proj,add_v_proj,txt_mlp.net.0.proj,txt_mlp.net.2,attn.to_add_out,timestep_embedder.linear_1,timestep_embedder.linear_2

data:
  height: 512
  width: 512
  train_dirs:
    # - path/to/data
    - /home/workspace/code/DiffSynth-Studio/example_image_dataset/edit

# NOTE: modify probs in twinflow_lora.py
method:
  method_type: TwinFlow_LoRA
  transport_type: Linear
  consistc_ratio: 1.0
  enhanced_ratio: 0.8
  ema_decay_rate: 0.99
  enhanced_range: [0.0, 1.0]
  time_dist_ctrl: [1.0, 1.0, 1.0]
  estimate_order: 2
  loss_func_type: {"type": "barron_reweighting", "alpha": 1.0}
  dist_match_cof: 0.5

# few-step sampling
sample:
  ckpt: "700"
  cfg_scale: 0
  cfg_interval: [0.00, 0.00]
  sampling_steps: 2 # 1
  stochast_ratio: 1.0 # 0.8
  extrapol_ratio: 0.0
  sampling_order: 1
  time_dist_ctrl: [1.0, 1.0, 1.0]
  rfba_gap_steps: [0.001, 0.7]
  sampling_style: few

train:
  load_checkpoint_path: ""
  output_dir: ../outputs
  seed: 42
  grad_accumulation_steps: 1
  micro_batch_size: 4
  num_train_epochs: 50
  optimizer: adamw
  lr: 1.0e-4
  betas:
    - 0.9
    - 0.99
  weight_decay: 0.0
  max_grad_norm: 1.0

另外，我还是用的F_t = (-1) ** (1 - self.integ_st) * model(x_t, t=q - tt / 2, **model_kwargs),没有使用t-tt/2，用的是from unified_sampler import UnifiedSampler这里的sampler

[QAQdev]

@hello-program 好的，q-tt/2就是t-tt/2。另外，你可以从TwinFlow_LoRA里的sampling_loop来采样，这样才会正确使用训练时使用的q-tt/2，不要用from unified_sampler import UnifiedSampler这里的sampler

[hello-program]

@QAQdev 谢谢，我去试试