There is a script broken for training in kohya_ss Please fix it

I've been able to make many training FLUX models (full tuning and loras) using dreambooth in kohya.. Until now. I need to check my model's going by sample images, but since the latest upgrades it gives me the error below and stop the trainining process. Please help me, and fix the code

Here is the error:

Using memory efficient save file: D:/AI-Tools/Training/Output/FullTuning/FLUX-Alemancu/FLUX-Alemancu.safetensors\model\FLUX-Alemancu-000010.safetensors
2025-06-05 16:18:51 INFO flux_train_utils.py:59
INFO generating sample images at step / サンプル画像生成 ステップ: 150 flux_train_utils.py:60
Traceback (most recent call last):
File "D:\AI-Tools\Training\kohya_ss\sd-scripts\flux_train.py", line 850, in
train(args)
File "D:\AI-Tools\Training\kohya_ss\sd-scripts\flux_train.py", line 763, in train
flux_train_utils.sample_images(
File "D:\AI-Tools\Training\kohya_ss\sd-scripts\library\flux_train_utils.py", line 70, in sample_images
text_encoders = [accelerator.unwrap_model(te) for te in text_encoders]
File "D:\AI-Tools\Training\kohya_ss\sd-scripts\library\flux_train_utils.py", line 70, in
text_encoders = [accelerator.unwrap_model(te) for te in text_encoders]
File "D:\AI-Tools\Training\kohya_ss\venv\lib\site-packages\accelerate\accelerator.py", line 2866, in unwrap_model
return extract_model_from_parallel(model, keep_fp32_wrapper, keep_torch_compile)
File "D:\AI-Tools\Training\kohya_ss\venv\lib\site-packages\accelerate\utils\other.py", line 176, in extract_model_from_parallel
has_compiled = has_compiled_regions(model)
File "D:\AI-Tools\Training\kohya_ss\venv\lib\site-packages\accelerate\utils\other.py", line 70, in has_compiled_regions
if module._modules:
AttributeError: 'NoneType' object has no attribute '_modules'
steps: 5%|██████████▏ | 150/3000 [32:50<10:24:00, 13.14s/it, avr_loss=0.475]
Traceback (most recent call last):
File "C:\Users\FabioROG\AppData\Local\Programs\Python\Python310\lib\runpy.py", line 196, in _run_module_as_main
return _run_code(code, main_globals, None,
File "C:\Users\FabioROG\AppData\Local\Programs\Python\Python310\lib\runpy.py", line 86, in run_code
exec(code, run_globals)
File "D:\AI-Tools\Training\kohya_ss\venv\Scripts\accelerate.EXE_main.py", line 7, in
sys.exit(main())
File "D:\AI-Tools\Training\kohya_ss\venv\lib\site-packages\accelerate\commands\accelerate_cli.py", line 50, in main
args.func(args)
File "D:\AI-Tools\Training\kohya_ss\venv\lib\site-packages\accelerate\commands\launch.py", line 1198, in launch_command
simple_launcher(args)
File "D:\AI-Tools\Training\kohya_ss\venv\lib\site-packages\accelerate\commands\launch.py", line 785, in simple_launcher
raise subprocess.CalledProcessError(returncode=process.returncode, cmd=cmd)
subprocess.CalledProcessError: Command '['D:\AI-Tools\Training\kohya_ss\venv\Scripts\python.exe', 'D:/AI-Tools/Training/kohya_ss/sd-scripts/flux_train.py', '--config_file', 'D:/AI-Tools/Training/Output/FullTuning/FLUX-Alemancu/FLUX-Alemancu.safetensors\model/config_dreambooth-20250605-154521.toml']' returned non-zero exit status 1.
16:18:56-883004 INFO Training has ended.

Start with sd-scripts/flux_train.py and library/flux_train_utils.py, following the sample-image generation path around the call to accelerator.unwrap_model. Reproduce training with sample generation enabled and investigate the NoneType error in accelerate's unwrap_model path. Done means sample images generate at the configured step and training continues without this traceback.

Written by the indexing model from the issue text.