Froddan commited on
Commit
ee22f78
1 Parent(s): dea8cb8

Update generate_model_grid.py

Browse files

now supports .safetensors as well

Files changed (1) hide show
  1. 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')