Ai Toolkit diffuser issue in latest version

[jewelsonn]

What happened?

I tried reinstalling AI-toolkit etc etc. But the problem persisted. It worked Only after I downgraded to previous version it was after trying so much to get it to work,

I tried to start a job in this new version and got the following error about distutil..

Steps to reproduce

1. Update to latest version 2.16.3
2. Try to run a job in AI-Toolkit with python 2.11 and cuda toolkit 13 (in my case it was for Krea2 RAW)

Relevant logs

Traceback (most recent call last):
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\run.py", line 147, in <module>
    main()
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\run.py", line 132, in main
    raise e
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\run.py", line 119, in main
    job = get_job(config_file, args.name)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\toolkit\job.py", line 28, in get_job
    from jobs import ExtensionJob
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\jobs\__init__.py", line 1, in <module>
    from .BaseJob import BaseJob
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\jobs\BaseJob.py", line 5, in <module>
    from jobs.process import BaseProcess
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\jobs\process\__init__.py", line 8, in <module>
    from .TrainSliderProcess import TrainSliderProcess
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\jobs\process\TrainSliderProcess.py", line 26, in <module>
    from .BaseSDTrainProcess import BaseSDTrainProcess
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\jobs\process\BaseSDTrainProcess.py", line 30, in <module>
    from toolkit.custom_adapter import CustomAdapter
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\toolkit\custom_adapter.py", line 13, in <module>
    from toolkit.models.control_lora_adapter import ControlLoraAdapter
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\toolkit\models\control_lora_adapter.py", line 5, in <module>
    from toolkit.lora_special import LoRASpecialNetwork
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\toolkit\lora_special.py", line 12, in <module>
    from toolkit.models.lokr import LokrModule
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\toolkit\models\lokr.py", line 9, in <module>
    from toolkit.network_mixins import ToolkitModuleMixin
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\toolkit\network_mixins.py", line 7, in <module>
    from optimum.quanto import QTensor
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\optimum\quanto\__init__.py", line 18, in <module>
    from .library import *
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\optimum\quanto\library\__init__.py", line 15, in <module>
    from .extensions import *
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\optimum\quanto\library\extensions\__init__.py", line 17, in <module>
    from .cpp import *
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\optimum\quanto\library\extensions\cpp\__init__.py", line 19, in <module>
    from ..extension import Extension
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\optimum\quanto\library\extensions\extension.py", line 7, in <module>
    from torch.utils.cpp_extension import load
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\torch\utils\cpp_extension.py", line 11, in <module>
    import setuptools
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\setuptools\__init__.py", line 22, in <module>
    import _distutils_hack.override  # noqa: F401
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\_distutils_hack\override.py", line 1, in <module>
    __import__('_distutils_hack').do_override()
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\_distutils_hack\__init__.py", line 89, in do_override
    ensure_local_distutils()
  File "C:\STABLE-DIFFUSION\Data\Packages\AI-Toolkit\venv\Lib\site-packages\_distutils_hack\__init__.py", line 76, in ensure_local_distutils
    assert '_distutils' in core.__file__, core.__file__
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError: C:\STABLE-DIFFUSION\Data\Assets\Python\cpython-3.11.14-windows-x86_64-none\Lib\distutils\core.py

Version

v2.16.3

What Operating System are you using?

Windows

[NeuralFault]

Thanks for reporting this, This is a regression for pre-exisiting installs of AI-Toolkit introduced in 2.16.3.

The update moved new installs of AI-Toolkit over to Python 3.12, to match upstream, and part of that was changing a small packaging flag that controls where Python looks for a legacy module (distutils). On 3.12 that change is actually required, because the old module is gone.

The thing is now that flag gets applied to every AI-Toolkit install, and updating the app doesn't rebuild packages that are already installed. So an install that was set up under 2.16.2 or older is still running Python 3.11, and on 3.11 that flag trips a known quirk in Python's packaging tooling, which is what produces that distutils assertion error.

On a clean install with 2.16.3 the install configuration works perfectly, because fresh installs always get 3.12 by default, the one case where it's needed and behaves. It only breaks existing installs still on 3.11, which is exactly the situation that never shows up when testing against a fresh setup, so possibly a test case that slipped by the devs.

Your current solution of holding back a version is a temporary fix, but a real resolution would be to just install a new AI-Toolkit package using the new defaults.

I have submitted a PR that adds a conditional against the current running Python version being used in that instance, which uses stdlib for the distutils on <3.12, and local (the new change) for 3.12 or greater, as a transitory fix.