[bugfix]Fix bug in Lora checkpoint saving step

[leisuzz]

What does this PR do?

Fix bug in checkpoint saving steps:

1. Free variable referenced before assignment in the original code if it is only elif isinstance(model, type(unwrap_model(text_encoder)))
2. IndexError: pop from empty list

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[leisuzz]

@sayakpaul Please take a look at this PR, I'm having trouble in dreambooth lora checkpoint saving. After these implementation, the issue has been solved

[sayakpaul]

Just so I understand, this is applicable only when we're using DeepSpeed?

[leisuzz]

@sayakpaul I'm using deepspeed to train the model, I'm not sure whether it works or not without deepspeed

[sayakpaul]

Same comments from your other PR that you closed apply here. We're missing:

• Incorporating LoRA/model loading support when using DeepSpeed.
• Other necessary checks when using DeepSpeed.

As mentioned in the other PR of yours, if you go here and search for DistributedType.DEEPSPEED, the outline of the changes should become evident and they will help you understand what we're missing in this PR. Without those changes, it'd be hard to merge this PR.

[leisuzz]

@sayakpaul I don' have any issue related with loading with deepspeed, the issue only occurs in saving steps

[sayakpaul]

We have extensively tested the changes I mentioned in my comment. Hence, I kindly ask you to propagate them.

[leisuzz]

Thanks, let me think how to modify them

[leisuzz]

@sayakpaul I've updated the save mode for the correct version, please take a look. This should work properly for loading function as well

[sayakpaul]

Let's keep the changes limited to the Flux, script perhaps?

[leisuzz]

@sayakpaul The issue happened both in sd3 and flux actually, I've also fixed the issue with distributed in condition with checkpoints_total_limit is not none

[leisuzz]

@sayakpaul I've fixed both loading and saving step, please take a look

[sayakpaul]

We'd prefer to tackle this one script at a time. So, please keep the changes to a single script, only.

[leisuzz]

@sayakpaul I've edited and now the changes are only based on train_dreambooth_lora_sd3.py (sd3 and flux are similar and I tested both of them), please take a look. Let me know if that's ok for me to edit on flux as well.

[sayakpaul]

Just a few more changes.

[leisuzz]

@sayakpaul Please take a look at this newest version, thanks

[sayakpaul]

@leisuzz where you not using a non LoRA version of the script? I think I got confused because of the repeated reviews.

Maybe we could close this PR and I could instead open a PR with you as a co-author?

[leisuzz]

@sayakpaul Sorry about this, I will close this PR then

[sayakpaul]

@leisuzz thanks for your understanding. Could you please provide your GitHub commit email address so that I can add you as a co-author?

[leisuzz]

@sayakpaul My email address is: [email protected]. Thanks for your help!