Fixed the issue on flux dreambooth lora training

[jeongiin]

What does this PR do?

Fixes #9548

I resolved the issue by changing the autocast context from nullcontext() to torch.autocast(accelerator.device.type, dtype=torch_dtype). This adjustment ensures that the model properly uses the correct mixed-precision mode during validation, preventing the errors I encountered.

I performed the test with dog images as suggested in the README_flux.md, and I also confirmed that the results were successfully uploaded to wandb:

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?

@sayakpaul

[linoytsaban]

Thanks @jeongiin! looks like this will also fix #9476

[linoytsaban]

looks good to me, just unsure about the commented line with the condition over last validation (also mentioned in #9476 (comment))
@sayakpaul if it looks ok to you we can merge

[jeongiin]

Thank you for reviewing, @linoytsaban !!

I'm not sure if this will be helpful, but when I used just autocast_ctx = torch.autocast(accelerator.device.type) to solve the issue, a problem similar to issue #9558 occurred. I was using wandb, black images were uploaded. Like this:

[icsl-Jeon]

@jeongiin did you check non nan of prompt_embeds from T5?
This might due to nan value of your pipeline output.

[sayakpaul]

To use autocast successfully with Flux during validation inference, we need to pre-compute the text embeddings because, otherwise, T5-xxl doesn't really work under autocast.

See how it's handled in:

diffusers/examples/controlnet/train_controlnet_flux.py

Line 139 in 61d3764

# pre calculate prompt embeds, pooled prompt embeds, text ids because t5 does not support autocast

[jeongiin]

Thank you for good advice, @icsl-Jeon and @sayakpaul !
I will follow @sayakpaul 's advice and apply this to train_dreambooth_lora_flux.py.

[sayakpaul]

@jeongiin there's another adjacent PR here: #9565

cc: @icsl-Jeon

[sayakpaul]

@jeongiin thank you for the changes but as mentioned in #9549 (comment), we need to handle the autocasting a bit differently. Let me know if anything is unclear.

[jeongiin]

Hello! I apologize for the delay! @sayakpaul

If I understand correctly, you're suggesting that further revisions may be needed, referring to #9565 and #9549's comment.

Would the issue not be resolved with just the addition of torch.autocast(accelerator.device.type, dtype=torch_dtype)?
In my testing, there didn't seem to be any problems with my modification.

Would you mind clarifying if there's something I might have overlooked?

[sayakpaul]

@jeongiin thanks a lot for your contributions! Could you maybe check if #9565 solves the problems this PR is trying to address?

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[yiyixuxu]

Should we close this now if it is fixed by #9565? @sayakpaul @linoytsaban

[sayakpaul]

Yes this can be closed. Sorry @jeongiin for the delay on our side but we appreciate your willingness to help us.

[jeongiin]

I haven't had enough GPU available lately so I can't verify this. :(
It's disappointing, but I'll check it out next time I get a chance! Thank you for your help!