mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-29 05:02:09 +03:00
44 lines
1.6 KiB
Python
Executable File
44 lines
1.6 KiB
Python
Executable File
#!/bin/env python
|
|
"""
|
|
Detect model type
|
|
|
|
Works for v1 and v2-base (EPS models), both standard inference and inpainting
|
|
|
|
But looking at model dumps between EPS and V type models, its only about parametrization, there are no differences in actual model (its just weighted differently without any structural difference)
|
|
So i don't see easy way to auto-detect if model should be run in `EPS` or `V` mode
|
|
Only difference are some calculations in `ldm/models/diffusion/ddpm.py` and by then we already need to know which code-path to trigger
|
|
(maaaybe there could be a way by looking at some cherry-picked base tensors min/max range, but I dont see that as reliable)
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
|
|
def signature(model):
|
|
if model is None:
|
|
return None
|
|
try:
|
|
size = model['state_dict']['model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight'].shape[1]
|
|
unet = model['state_dict']['model.diffusion_model.input_blocks.0.0.weight'].shape[1]
|
|
except:
|
|
return 'unknown'
|
|
guess = 'v1' if size == 768 else 'v2' # 768 for v1 and 1024 for v2
|
|
guess += '-inference' if unet == 4 else '-inpainting' # inference models have shorter inputs, 4 for inference 9 for inpainting
|
|
return guess
|
|
|
|
def load(file: str):
|
|
try:
|
|
model = torch.load(file, map_location='cpu')
|
|
return model
|
|
except Exception as err:
|
|
print(f"Error loading {f}: {err}")
|
|
|
|
if __name__ == "__main__":
|
|
sys.argv.pop(0)
|
|
for f in sys.argv:
|
|
if os.path.isfile(f):
|
|
print(f"Model {f} is of type {signature(load(f))}")
|
|
else:
|
|
print(f"{f} is not a file")
|