Update generate_model_grid.py
Browse filesnow supports .safetensors as well
- generate_model_grid.py +2 -2
generate_model_grid.py
CHANGED
@@ -244,12 +244,12 @@ class Script(scripts.Script):
|
|
244 |
z_valuez = ''
|
245 |
dirpath = os.path.join(models_path, 'Stable-diffusion')
|
246 |
for path in os.listdir(dirpath):
|
247 |
-
if path.endswith('.ckpt'):
|
248 |
filenames.append(path)
|
249 |
else:
|
250 |
if os.path.isdir(os.path.join(dirpath,path)):
|
251 |
for subpath in os.listdir(os.path.join(dirpath,path)):
|
252 |
-
if subpath.endswith('.ckpt'):
|
253 |
filenames.append(subpath)
|
254 |
|
255 |
filenames.append('model.ckpt')
|
|
|
244 |
z_valuez = ''
|
245 |
dirpath = os.path.join(models_path, 'Stable-diffusion')
|
246 |
for path in os.listdir(dirpath):
|
247 |
+
if path.endswith('.ckpt') or path.endswith('.safetensors'):
|
248 |
filenames.append(path)
|
249 |
else:
|
250 |
if os.path.isdir(os.path.join(dirpath,path)):
|
251 |
for subpath in os.listdir(os.path.join(dirpath,path)):
|
252 |
+
if subpath.endswith('.ckpt') or path.endswith('.safetensors'):
|
253 |
filenames.append(subpath)
|
254 |
|
255 |
filenames.append('model.ckpt')
|