Problems with saved lora models when using rslora.

[ydniuyongjie]

Describe the bug

The SDXL model was fine-tuned using the rslora method and the training process was fine-tuned. After the training, the Lora model was saved and then the lora model was reloaded for the image generation test and it was found that the images generated in the test phase were wrong. It seems that the images generated during the testing phase do not use Lora, and the reason for analyzing this should be that the state of rslora was not saved when the model was saved.

Reproduction

Save the lora layers

accelerator.wait_for_everyone()
if accelerator.is_main_process:
    unet = unwrap_model(unet)
    unet = unet.to(torch.float32)
    unet_lora_layers = convert_state_dict_to_diffusers(get_peft_model_state_dict(unet))

    if args.train_text_encoder:
        text_encoder_one = unwrap_model(text_encoder_one)
        text_encoder_lora_layers = convert_state_dict_to_diffusers(
            get_peft_model_state_dict(text_encoder_one.to(torch.float32))
        )
        text_encoder_two = unwrap_model(text_encoder_two)
        text_encoder_2_lora_layers = convert_state_dict_to_diffusers(
            get_peft_model_state_dict(text_encoder_two.to(torch.float32))
        )
    else:
        text_encoder_lora_layers = None
        text_encoder_2_lora_layers = None

    StableDiffusionXLPipeline.save_lora_weights(
        save_directory=args.output_dir,
        unet_lora_layers=unet_lora_layers,
        text_encoder_lora_layers=text_encoder_lora_layers,
        text_encoder_2_lora_layers=text_encoder_2_lora_layers,
    )
    if args.output_kohya_format:
        lora_state_dict = load_file(f"{args.output_dir}/pytorch_lora_weights.safetensors")
        peft_state_dict = convert_all_state_dict_to_peft(lora_state_dict)
        kohya_state_dict = convert_state_dict_to_kohya(peft_state_dict)
        save_file(kohya_state_dict, f"{args.output_dir}/pytorch_lora_weights_kohya.safetensors")

    # Final inference with the newly saved lora weights
    # Load previous pipeline
    vae = AutoencoderKL.from_pretrained(
        vae_path,
        subfolder="vae" if args.pretrained_vae_model_name_or_path is None else None,
        revision=args.revision,
        variant=args.variant,
        torch_dtype=weight_dtype,
    )
    pipeline = StableDiffusionXLPipeline.from_pretrained(
        args.pretrained_model_name_or_path,
        vae=vae,
        revision=args.revision,
        variant=args.variant,
        torch_dtype=weight_dtype,
        configuration=unet_lora_config,
    )

    # load attention processors
    pipeline.load_lora_weights(args.output_dir,weight_name="pytorch_lora_weights.safetensors")

    # run inference运行推理
    images = []
    if args.validation_prompt and args.num_validation_images > 0:
        pipeline_args = {"prompt": args.validation_prompt, "num_inference_steps": 25}
        images = log_validation(
            pipeline,
            args,
            accelerator,
            pipeline_args,
            epoch,
            is_final_validation=True,
        )

Logs

No response

System Info

• diffusers version: 0.28.0.dev0
• Platform: Linux-5.4.0-42-generic-x86_64-with-glibc2.10
• Python version: 3.8.0
• PyTorch version (GPU?): 2.2.0+cu118 (True)
• Huggingface_hub version: 0.23.0
• Transformers version: 4.40.2
• Accelerate version: 0.29.3
• xFormers version: not installed
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

Who can help?

No response

[ydniuyongjie]

Has anyone else had this problem?

[asomoza]

Hi, does this happen to you if you train a normal dreambooth lora?

[dariant]

This seems to also be a problem when fine-tuning with the rsLoRA method and the DreamBooth script on other models (e.g., SD-2.1). Images generated during inference with fine-tuned rsLoRA weights do not resemble any of the validation images generated during training, as if the rsLoRA was not saved or loaded correctly.