kohya-ss / kohya-ss/sd-scripts
Why are model names hardcoded in analyze_checkpoint_state?
- 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