Use enable_gradient_checkpointing() for flux

[jhj0517]

Thank you for your amazing work 😚.

Related Issues / PRs

• Fix Error running job: 'NoneType' object is not callable #243
• Refactor gradient checkpointing huggingface/diffusers#10611 ( This PR was merged a few hours ago and caused the issue above )

Summarize Changes

• With the Refactor gradient checkpointing huggingface/diffusers#10611 changes, enable_gradient_checkpointing() also needs to be called for Flux, unless it gets a None type error for the _gradient_checkpointing_func attribute.

[firofame]

I have verified that this PR fixes the issue.

git clone -b fix/gradient-checkpointing --single-branch https://github.com/jhj0517/ai-toolkit.git /root/ai-toolkit

[kx221005]

I have verified that this PR fixes the issue.

git clone -b fix/gradient-checkpointing --single-branch https://github.com/jhj0517/ai-toolkit.git /root/ai-toolkit

Can I know what is the proper step? I run these code first only run setup code?

git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
git submodule update --init --recursive
python -m venv venv
source venv/bin/activate
pip install torch
pip install -r requirements.txt
pip install --upgrade accelerate transformers diffusers huggingface_hub #Optional, run it if you run into issues

which line come first?