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

Update generate_model_grid.py

Browse files

Should now also look in subfolders

Files changed (1) hide show
  1. generate_model_grid.py +8 -1
generate_model_grid.py CHANGED
@@ -242,9 +242,16 @@ class Script(scripts.Script):
242
  def ui(self, is_img2img):
243
  filenames = []
244
  z_valuez = ''
245
- for path in os.listdir(os.path.join(models_path, 'Stable-diffusion')):
 
246
  if path.endswith('.ckpt'):
247
  filenames.append(path)
 
 
 
 
 
 
248
  filenames.append('model.ckpt')
249
 
250
  with gr.Row():
 
242
  def ui(self, is_img2img):
243
  filenames = []
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')
256
 
257
  with gr.Row():