feat: enhance auto-organize functionality with empty directory cleanup and progress reporting

This commit is contained in:
Will Miao
2025-08-13 10:36:31 +08:00
parent a920921570
commit 5a0b3470f1
2 changed files with 104 additions and 50 deletions

View File

@@ -852,7 +852,7 @@ class BaseModelRoutes(ABC):
current_dir = os.path.dirname(file_path) current_dir = os.path.dirname(file_path)
# Skip if already in correct location # Skip if already in correct location
if os.path.normpath(current_dir) == os.path.normpath(target_dir): if current_dir.replace(os.sep, '/') == target_dir.replace(os.sep, '/'):
skipped_count += 1 skipped_count += 1
processed += 1 processed += 1
continue continue
@@ -914,6 +914,25 @@ class BaseModelRoutes(ABC):
await asyncio.sleep(0.1) await asyncio.sleep(0.1)
# Send completion message # Send completion message
await ws_manager.broadcast({
'type': 'auto_organize_progress',
'status': 'cleaning',
'total': total_models,
'processed': processed,
'success': success_count,
'failures': failure_count,
'skipped': skipped_count,
'message': 'Cleaning up empty directories...'
})
# Clean up empty directories after organizing
from ..utils.utils import remove_empty_dirs
cleanup_counts = {}
for root in model_roots:
removed = remove_empty_dirs(root)
cleanup_counts[root] = removed
# Send cleanup completed message
await ws_manager.broadcast({ await ws_manager.broadcast({
'type': 'auto_organize_progress', 'type': 'auto_organize_progress',
'status': 'completed', 'status': 'completed',
@@ -921,7 +940,8 @@ class BaseModelRoutes(ABC):
'processed': processed, 'processed': processed,
'success': success_count, 'success': success_count,
'failures': failure_count, 'failures': failure_count,
'skipped': skipped_count 'skipped': skipped_count,
'cleanup': cleanup_counts
}) })
# Prepare response with limited details # Prepare response with limited details
@@ -933,7 +953,8 @@ class BaseModelRoutes(ABC):
'success': success_count, 'success': success_count,
'skipped': skipped_count, 'skipped': skipped_count,
'failures': failure_count, 'failures': failure_count,
'organization_type': 'flat' if is_flat_structure else 'structured' 'organization_type': 'flat' if is_flat_structure else 'structured',
'cleaned_dirs': cleanup_counts
} }
} }

View File

@@ -151,8 +151,8 @@ def calculate_relative_path_for_model(model_data: Dict) -> str:
# Get base model name from model metadata # Get base model name from model metadata
civitai_data = model_data.get('civitai', {}) civitai_data = model_data.get('civitai', {})
# For CivitAI models, prefer civitai data; for non-CivitAI models, use model_data directly # For CivitAI models, prefer civitai data only if 'id' exists; for non-CivitAI models, use model_data directly
if civitai_data: if civitai_data and civitai_data.get('id') is not None:
base_model = civitai_data.get('baseModel', '') base_model = civitai_data.get('baseModel', '')
else: else:
# Fallback to model_data fields for non-CivitAI models # Fallback to model_data fields for non-CivitAI models
@@ -184,3 +184,36 @@ def calculate_relative_path_for_model(model_data: Dict) -> str:
formatted_path = formatted_path.replace('{first_tag}', first_tag) formatted_path = formatted_path.replace('{first_tag}', first_tag)
return formatted_path return formatted_path
def remove_empty_dirs(path):
"""Recursively remove empty directories starting from the given path.
Args:
path (str): Root directory to start cleaning from
Returns:
int: Number of empty directories removed
"""
removed_count = 0
if not os.path.isdir(path):
return removed_count
# List all files in directory
files = os.listdir(path)
# Process all subdirectories first
for file in files:
full_path = os.path.join(path, file)
if os.path.isdir(full_path):
removed_count += remove_empty_dirs(full_path)
# Check if directory is now empty (after processing subdirectories)
if not os.listdir(path):
try:
os.rmdir(path)
removed_count += 1
except OSError:
pass
return removed_count