kohya-ss / kohya-ss/sd-scripts

Why are model names hardcoded in analyze_checkpoint_state?

Open
#1,968 2 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

My checkpoint is named like this, `diffusion_pytorch_model-00001-of-00002.safetensors diffusion_pytorch_model-00002-of-00002.safetensors`, it won't load because the name indices 00003 is hardcoded in your code,

```python
def analyze_checkpoint_state(ckpt_path: str) -> Tuple[bool, bool, Tuple[int, int], List[str]]:
"""
チェックポイントの状態を分析し、DiffusersかBFLか、devかschnellか、ブロック数を計算して返す。

Args:
ckpt_path (str): チェックポイントファイルまたはディレクトリのパス。

Returns:
Tuple[bool, bool, Tuple[int, int], List[str]]:
- bool: Diffusersかどうかを示すフラグ。
- bool: Schnellかどうかを示すフラグ。
- Tuple[int, int]: ダブルブロックとシングルブロックの数。
- List[str]: チェックポイントに含まれるキーのリスト。
"""
# check the state dict: Diffusers or BFL, dev or schnell, number of blocks
logger.info(f"Checking the state dict: Diffusers or BFL, dev or schnell")

if os.path.isdir(ckpt_path): # if ckpt_path is a directory, it is Diffusers
ckpt_path = os.path.join(ckpt_path, "transformer", "diffusion_pytorch_model-00001-of-00003.safetensors")
if "00001-of-00003" in ckpt_path:
ckpt_paths = [ckpt_path.replace("00001-of-00003", f"0000{i}-of-00003") for i in range(1, 4)]
else:
ckpt_paths = [ckpt_path]
```

Why is that?

Contributor guide

No contributing guide indexed for this repository

Research direction

Locate analyze_checkpoint_state and inspect how directory paths and shard names are detected. Reproduce the failure with the two diffusion_pytorch_model-00001-of-00002.safetensors files, then confirm that checkpoint loading handles the reported shard count without assuming three files.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
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.