diff --git a/modules/caption/deepbooru.py b/modules/caption/deepbooru.py
index 19ff58afc..6ccb848b1 100644
--- a/modules/caption/deepbooru.py
+++ b/modules/caption/deepbooru.py
@@ -205,6 +205,7 @@ def batch(
Returns:
Combined tag results
"""
+ import os
import time
from pathlib import Path
import rich.progress as rp
@@ -214,55 +215,15 @@ def batch(
# Collect image files
image_files = []
- image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
-
- # From file picker
- if batch_files:
- for f in batch_files:
- if isinstance(f, dict):
- image_files.append(Path(f['name']))
- elif hasattr(f, 'name'):
- image_files.append(Path(f.name))
- else:
- image_files.append(Path(f))
-
- # From folder picker
- if batch_folder:
- folder_path = None
- if isinstance(batch_folder, list) and len(batch_folder) > 0:
- f = batch_folder[0]
- if isinstance(f, dict):
- folder_path = Path(f['name']).parent
- elif hasattr(f, 'name'):
- folder_path = Path(f.name).parent
- if folder_path and folder_path.is_dir():
- if recursive:
- for ext in image_extensions:
- image_files.extend(folder_path.rglob(f'*{ext}'))
- else:
- for ext in image_extensions:
- image_files.extend(folder_path.glob(f'*{ext}'))
-
- # From string path
- if batch_str and batch_str.strip():
+ if batch_files is not None:
+ image_files += [f.name for f in batch_files]
+ if batch_folder is not None:
+ image_files += [f.name for f in batch_folder]
+ if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
+ image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
folder_path = Path(batch_str.strip())
- if folder_path.is_dir():
- if recursive:
- for ext in image_extensions:
- image_files.extend(folder_path.rglob(f'*{ext}'))
- else:
- for ext in image_extensions:
- image_files.extend(folder_path.glob(f'*{ext}'))
-
- # Remove duplicates while preserving order
- seen = set()
- unique_files = []
- for f in image_files:
- f_resolved = f.resolve()
- if f_resolved not in seen:
- seen.add(f_resolved)
- unique_files.append(f)
- image_files = unique_files
+ for ext in image_extensions:
+ image_files.extend(str(p) for p in (folder_path.rglob(f'*{ext}') if recursive else folder_path.glob(f'*{ext}')))
if not image_files:
log.warning('DeepBooru batch: no images found')
@@ -280,25 +241,26 @@ def batch(
with pbar:
task = pbar.add_task(total=len(image_files), description='starting...')
- for img_path in image_files:
- pbar.update(task, advance=1, description=str(img_path.name))
+ for file in image_files:
+ file_name = os.path.basename(file)
+ pbar.update(task, advance=1, description=file_name)
try:
if shared.state.interrupted:
log.info('DeepBooru batch: interrupted')
break
- image = Image.open(img_path)
+ image = Image.open(file)
tags_str = model.tag_multi(image, **kwargs)
if save_output:
from modules.caption import tagger
- tagger.save_tags_to_file(img_path, tags_str, save_append)
+ tagger.save_tags_to_file(Path(file), tags_str, save_append)
- results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}')
+ results.append(f'{file_name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{file_name}: {tags_str}')
except Exception as e:
- log.error(f'DeepBooru batch: file="{img_path}" error={e}')
- results.append(f'{img_path.name}: ERROR - {e}')
+ log.error(f'DeepBooru batch: file="{file}" error={e}')
+ results.append(f'{file_name}: ERROR - {e}')
model.stop()
elapsed = time.time() - t0
diff --git a/modules/caption/waifudiffusion.py b/modules/caption/waifudiffusion.py
index ba7077dac..3e284db75 100644
--- a/modules/caption/waifudiffusion.py
+++ b/modules/caption/waifudiffusion.py
@@ -419,6 +419,7 @@ def batch(
Returns:
Combined tag results
"""
+ import os
from pathlib import Path
# Load model
@@ -429,55 +430,15 @@ def batch(
# Collect image files
image_files = []
- image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
-
- # From file picker
- if batch_files:
- for f in batch_files:
- if isinstance(f, dict):
- image_files.append(Path(f['name']))
- elif hasattr(f, 'name'):
- image_files.append(Path(f.name))
- else:
- image_files.append(Path(f))
-
- # From folder picker
- if batch_folder:
- folder_path = None
- if isinstance(batch_folder, list) and len(batch_folder) > 0:
- f = batch_folder[0]
- if isinstance(f, dict):
- folder_path = Path(f['name']).parent
- elif hasattr(f, 'name'):
- folder_path = Path(f.name).parent
- if folder_path and folder_path.is_dir():
- if recursive:
- for ext in image_extensions:
- image_files.extend(folder_path.rglob(f'*{ext}'))
- else:
- for ext in image_extensions:
- image_files.extend(folder_path.glob(f'*{ext}'))
-
- # From string path
- if batch_str and batch_str.strip():
+ if batch_files is not None:
+ image_files += [f.name for f in batch_files]
+ if batch_folder is not None:
+ image_files += [f.name for f in batch_folder]
+ if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
+ image_extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
folder_path = Path(batch_str.strip())
- if folder_path.is_dir():
- if recursive:
- for ext in image_extensions:
- image_files.extend(folder_path.rglob(f'*{ext}'))
- else:
- for ext in image_extensions:
- image_files.extend(folder_path.glob(f'*{ext}'))
-
- # Remove duplicates while preserving order
- seen = set()
- unique_files = []
- for f in image_files:
- f_resolved = f.resolve()
- if f_resolved not in seen:
- seen.add(f_resolved)
- unique_files.append(f)
- image_files = unique_files
+ for ext in image_extensions:
+ image_files.extend(str(p) for p in (folder_path.rglob(f'*{ext}') if recursive else folder_path.glob(f'*{ext}')))
if not image_files:
log.warning('WaifuDiffusion batch: no images found')
@@ -496,25 +457,26 @@ def batch(
with pbar:
task = pbar.add_task(total=len(image_files), description='starting...')
- for img_path in image_files:
- pbar.update(task, advance=1, description=str(img_path.name))
+ for file in image_files:
+ file_name = os.path.basename(file)
+ pbar.update(task, advance=1, description=file_name)
try:
if shared.state.interrupted:
log.info('WaifuDiffusion batch: interrupted')
break
- image = Image.open(img_path)
+ image = Image.open(file)
tags_str = tagger.predict(image, **kwargs)
if save_output:
from modules.caption import tagger as tagger_module
- tagger_module.save_tags_to_file(img_path, tags_str, save_append)
+ tagger_module.save_tags_to_file(Path(file), tags_str, save_append)
- results.append(f'{img_path.name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{img_path.name}: {tags_str}')
+ results.append(f'{file_name}: {tags_str[:100]}...' if len(tags_str) > 100 else f'{file_name}: {tags_str}')
except Exception as e:
- log.error(f'WaifuDiffusion batch: file="{img_path}" error={e}')
- results.append(f'{img_path.name}: ERROR - {e}')
+ log.error(f'WaifuDiffusion batch: file="{file}" error={e}')
+ results.append(f'{file_name}: ERROR - {e}')
elapsed = time.time() - t0
log.info(f'WaifuDiffusion batch: complete images={len(results)} time={elapsed:.1f}s')
diff --git a/modules/ui_caption.py b/modules/ui_caption.py
index 33ee8132f..ceb2c2d31 100644
--- a/modules/ui_caption.py
+++ b/modules/ui_caption.py
@@ -270,6 +270,7 @@ def create_ui():
with gr.Row():
wd_show_scores = gr.Checkbox(label='Show confidence scores', value=shared.opts.tagger_show_scores, elem_id='wd_show_scores')
gr.HTML('')
+ gr.HTML('')
with gr.Accordion(label='Caption: Batch', open=False, visible=True):
with gr.Row():
wd_batch_files = gr.File(label="Files", show_label=True, file_count='multiple', file_types=['image'], interactive=True, height=100, elem_id='wd_batch_files')