kohya-ss / kohya-ss/sd-scripts

Can not resume with the Stage folder, Flux model

Open
#2,092 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
7.2k
Forks
1.2k
Avg merge
11m
Merged PRs (30d)
2

Description

This error I have met about 2-3 months, but until now has not been overcome, I often resume before

```
2025-05-22 05:36:42 INFO resume training from train_util.py:4684
local state:
/content/drive/Shareddr
ives/StableDiffusion/Lo
ra/lip_color_skin_1/lip
_color_skin_1-state
INFO Loading states from accelerator.py:3085
/content/drive/Sharedd
rives/StableDiffusion/
Lora/lip_color_skin_1/
lip_color_skin_1-state
INFO load train state from train_network.py:947
/content/drive/Shared
drives/StableDiffusio
n/Lora/lip_color_skin
_1/lip_color_skin_1-s
tate/train_state.json
: {'current_epoch':
12, 'current_step':
5760, 'step': 5761}
/content/sd-scripts/networks/lora_flux.py:1048: FutureWarning: Positional args are being deprecated, use kwargs instead. Refer to https://pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.
return super().state_dict(destination, prefix, keep_vars)
2025-05-22 05:36:43 INFO All model weights checkpointing.py:214
loaded successfully
2025-05-22 05:36:45 INFO All optimizer states checkpointing.py:222
loaded successfully
INFO All scheduler states checkpointing.py:229
loaded successfully
INFO All dataloader checkpointing.py:241
sampler states loaded
successfully
INFO Could not load random checkpointing.py:264
states
Traceback (most recent call last):
File "/content/sd-scripts/flux_train_network.py", line 559, in
trainer.train(args)
File "/content/sd-scripts/train_network.py", line 953, in train
train_util.resume_from_local_or_hf_if_specified(accelerator, args)
File "/content/sd-scripts/library/train_util.py", line 4685, in resume_from_local_or_hf_if_specified
accelerator.load_state(args.resume)
File "/usr/local/lib/python3.11/dist-packages/accelerate/accelerator.py", line 3156, in load_state
self.step = override_attributes["step"]
~~~~~~~~~~~~~~~~~~~^^^^^^^^
KeyError: 'step'
Traceback (most recent call last):
File "/usr/local/bin/accelerate", line 8, in
sys.exit(main())
^^^^^^
File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/accelerate_cli.py", line 48, in main
args.func(args)
File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 1106, in launch_command
simple_launcher(args)
File "/usr/local/lib/python3.11/dist-packages/accelerate/commands/launch.py", line 704, in simple_launcher
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with library/train_util.py at resume_from_local_or_hf_if_specified, then follow train_network.py and flux_train_network.py into accelerate's Accelerator.load_state. Inspect the recorded train_state.json and the local state-loading path; done means resuming the reported Flux training state without the shown KeyError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.