diff --git a/locales/de.json b/locales/de.json
index d96cba50..a19e7047 100644
--- a/locales/de.json
+++ b/locales/de.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "Fehler beim Pausieren des Downloads: {error}",
"downloadResumed": "Download fortgesetzt",
"resumeFailed": "Fehler beim Fortsetzen des Downloads: {error}",
+ "downloadStopped": "Download abgebrochen",
+ "stopFailed": "Download konnte nicht abgebrochen werden: {error}",
"deleted": "Beispielbild gelöscht",
"deleteFailed": "Fehler beim Löschen des Beispielbilds",
"setPreviewFailed": "Fehler beim Setzen des Vorschaubilds"
diff --git a/locales/en.json b/locales/en.json
index 3cb82402..3785588d 100644
--- a/locales/en.json
+++ b/locales/en.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "Failed to pause download: {error}",
"downloadResumed": "Download resumed",
"resumeFailed": "Failed to resume download: {error}",
+ "downloadStopped": "Download cancelled",
+ "stopFailed": "Failed to cancel download: {error}",
"deleted": "Example image deleted",
"deleteFailed": "Failed to delete example image",
"setPreviewFailed": "Failed to set preview image"
diff --git a/locales/es.json b/locales/es.json
index bf706254..a47d171e 100644
--- a/locales/es.json
+++ b/locales/es.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "Error al pausar descarga: {error}",
"downloadResumed": "Descarga reanudada",
"resumeFailed": "Error al reanudar descarga: {error}",
+ "downloadStopped": "Descarga cancelada",
+ "stopFailed": "Error al cancelar descarga: {error}",
"deleted": "Imagen de ejemplo eliminada",
"deleteFailed": "Error al eliminar imagen de ejemplo",
"setPreviewFailed": "Error al establecer imagen de vista previa"
diff --git a/locales/fr.json b/locales/fr.json
index c64de794..93556d9f 100644
--- a/locales/fr.json
+++ b/locales/fr.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "Échec de la mise en pause du téléchargement : {error}",
"downloadResumed": "Téléchargement repris",
"resumeFailed": "Échec de la reprise du téléchargement : {error}",
+ "downloadStopped": "Téléchargement annulé",
+ "stopFailed": "Échec de l'annulation du téléchargement : {error}",
"deleted": "Image d'exemple supprimée",
"deleteFailed": "Échec de la suppression de l'image d'exemple",
"setPreviewFailed": "Échec de la définition de l'image d'aperçu"
diff --git a/locales/he.json b/locales/he.json
index ecc2ade6..79fdd7e9 100644
--- a/locales/he.json
+++ b/locales/he.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "השהיית ההורדה נכשלה: {error}",
"downloadResumed": "ההורדה חודשה",
"resumeFailed": "חידוש ההורדה נכשל: {error}",
+ "downloadStopped": "ההורדה בוטלה",
+ "stopFailed": "נכשל בביטול ההורדה: {error}",
"deleted": "תמונת הדוגמה נמחקה",
"deleteFailed": "מחיקת תמונת הדוגמה נכשלה",
"setPreviewFailed": "הגדרת תמונת התצוגה המקדימה נכשלה"
diff --git a/locales/ja.json b/locales/ja.json
index 76a71313..d39fccaa 100644
--- a/locales/ja.json
+++ b/locales/ja.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "ダウンロードの一時停止に失敗しました：{error}",
"downloadResumed": "ダウンロードが再開されました",
"resumeFailed": "ダウンロードの再開に失敗しました：{error}",
+ "downloadStopped": "ダウンロードをキャンセルしました",
+ "stopFailed": "ダウンロードのキャンセルに失敗しました：{error}",
"deleted": "例画像が削除されました",
"deleteFailed": "例画像の削除に失敗しました",
"setPreviewFailed": "プレビュー画像の設定に失敗しました"
diff --git a/locales/ko.json b/locales/ko.json
index 7fb88db8..a6565c10 100644
--- a/locales/ko.json
+++ b/locales/ko.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "다운로드 일시정지 실패: {error}",
"downloadResumed": "다운로드가 재개되었습니다",
"resumeFailed": "다운로드 재개 실패: {error}",
+ "downloadStopped": "다운로드가 취소되었습니다",
+ "stopFailed": "다운로드 취소 실패: {error}",
"deleted": "예시 이미지가 삭제되었습니다",
"deleteFailed": "예시 이미지 삭제 실패",
"setPreviewFailed": "미리보기 이미지 설정 실패"
diff --git a/locales/ru.json b/locales/ru.json
index 4b014df9..5a3b496e 100644
--- a/locales/ru.json
+++ b/locales/ru.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "Не удалось приостановить загрузку: {error}",
"downloadResumed": "Загрузка возобновлена",
"resumeFailed": "Не удалось возобновить загрузку: {error}",
+ "downloadStopped": "Загрузка отменена",
+ "stopFailed": "Не удалось отменить загрузку: {error}",
"deleted": "Пример изображения удален",
"deleteFailed": "Не удалось удалить пример изображения",
"setPreviewFailed": "Не удалось установить превью изображение"
diff --git a/locales/zh-CN.json b/locales/zh-CN.json
index a1454a81..474b61ca 100644
--- a/locales/zh-CN.json
+++ b/locales/zh-CN.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "暂停下载失败：{error}",
"downloadResumed": "下载已恢复",
"resumeFailed": "恢复下载失败：{error}",
+ "downloadStopped": "下载已取消",
+ "stopFailed": "取消下载失败：{error}",
"deleted": "示例图片已删除",
"deleteFailed": "删除示例图片失败",
"setPreviewFailed": "设置预览图片失败"
diff --git a/locales/zh-TW.json b/locales/zh-TW.json
index 1a7f98fa..9e1385fd 100644
--- a/locales/zh-TW.json
+++ b/locales/zh-TW.json
@@ -1245,6 +1245,8 @@
"pauseFailed": "暫停下載失敗：{error}",
"downloadResumed": "下載已恢復",
"resumeFailed": "恢復下載失敗：{error}",
+ "downloadStopped": "下載已取消",
+ "stopFailed": "取消下載失敗：{error}",
"deleted": "範例圖片已刪除",
"deleteFailed": "刪除範例圖片失敗",
"setPreviewFailed": "設定預覽圖片失敗"
diff --git a/py/routes/example_images_route_registrar.py b/py/routes/example_images_route_registrar.py
index aa12c3b1..63bac34b 100644
--- a/py/routes/example_images_route_registrar.py
+++ b/py/routes/example_images_route_registrar.py
@@ -22,6 +22,7 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
RouteDefinition("GET", "/api/lm/example-images-status", "get_example_images_status"),
RouteDefinition("POST", "/api/lm/pause-example-images", "pause_example_images"),
RouteDefinition("POST", "/api/lm/resume-example-images", "resume_example_images"),
+ RouteDefinition("POST", "/api/lm/stop-example-images", "stop_example_images"),
RouteDefinition("POST", "/api/lm/open-example-images-folder", "open_example_images_folder"),
RouteDefinition("GET", "/api/lm/example-image-files", "get_example_image_files"),
RouteDefinition("GET", "/api/lm/has-example-images", "has_example_images"),
diff --git a/py/routes/handlers/example_images_handlers.py b/py/routes/handlers/example_images_handlers.py
index 111013e2..f7acf5ef 100644
--- a/py/routes/handlers/example_images_handlers.py
+++ b/py/routes/handlers/example_images_handlers.py
@@ -68,6 +68,13 @@ class ExampleImagesDownloadHandler:
except DownloadNotRunningError as exc:
return web.json_response({'success': False, 'error': str(exc)}, status=400)
+ async def stop_example_images(self, request: web.Request) -> web.StreamResponse:
+ try:
+ result = await self._download_manager.stop_download(request)
+ return web.json_response(result)
+ except DownloadNotRunningError as exc:
+ return web.json_response({'success': False, 'error': str(exc)}, status=400)
+
async def force_download_example_images(self, request: web.Request) -> web.StreamResponse:
try:
payload = await request.json()
@@ -149,6 +156,7 @@ class ExampleImagesHandlerSet:
"get_example_images_status": self.download.get_example_images_status,
"pause_example_images": self.download.pause_example_images,
"resume_example_images": self.download.resume_example_images,
+ "stop_example_images": self.download.stop_example_images,
"force_download_example_images": self.download.force_download_example_images,
"import_example_images": self.management.import_example_images,
"delete_example_image": self.management.delete_example_image,
diff --git a/py/utils/example_images_download_manager.py b/py/utils/example_images_download_manager.py
index 4726e81b..85a70b6c 100644
--- a/py/utils/example_images_download_manager.py
+++ b/py/utils/example_images_download_manager.py
@@ -105,6 +105,7 @@ class DownloadManager:
self._progress = _DownloadProgress()
self._ws_manager = ws_manager
self._state_lock = state_lock or asyncio.Lock()
+ self._stop_requested = False
def _resolve_output_dir(self, library_name: str | None = None) -> str:
base_path = get_settings_manager().get('example_images_path')
@@ -145,6 +146,7 @@ class DownloadManager:
raise DownloadConfigurationError('Example images path not configured in settings')
self._progress.reset()
+ self._stop_requested = False
self._progress['status'] = 'running'
self._progress['start_time'] = time.time()
self._progress['end_time'] = None
@@ -267,6 +269,27 @@ class DownloadManager:
'success': True,
'message': 'Download resumed'
}
+
+ async def stop_download(self, request):
+ """Stop the example images download after the current model completes."""
+
+ async with self._state_lock:
+ if not self._is_downloading:
+ raise DownloadNotRunningError()
+
+ if self._progress['status'] in {'completed', 'error', 'stopped'}:
+ raise DownloadNotRunningError()
+
+ if self._progress['status'] != 'stopping':
+ self._stop_requested = True
+ self._progress['status'] = 'stopping'
+
+ await self._broadcast_progress(status='stopping')
+
+ return {
+ 'success': True,
+ 'message': 'Download stopping'
+ }
async def _download_all_example_images(
self,
@@ -311,6 +334,12 @@ class DownloadManager:
# Process each model
for i, (scanner_type, model, scanner) in enumerate(all_models):
+ async with self._state_lock:
+ current_status = self._progress['status']
+
+ if current_status not in {'running', 'paused', 'stopping'}:
+ break
+
# Main logic for processing model is here, but actual operations are delegated to other classes
was_remote_download = await self._process_model(
scanner_type,
@@ -321,24 +350,59 @@ class DownloadManager:
downloader,
library_name,
)
-
+
# Update progress
self._progress['completed'] += 1
- await self._broadcast_progress(status='running')
-
+
+ async with self._state_lock:
+ current_status = self._progress['status']
+ should_stop = self._stop_requested and current_status == 'stopping'
+
+ broadcast_status = 'running' if current_status == 'running' else current_status
+ await self._broadcast_progress(status=broadcast_status)
+
+ if should_stop:
+ break
+
# Only add delay after remote download of models, and not after processing the last model
- if was_remote_download and i < len(all_models) - 1 and self._progress['status'] == 'running':
+ if (
+ was_remote_download
+ and i < len(all_models) - 1
+ and current_status == 'running'
+ ):
await asyncio.sleep(delay)
-
- # Mark as completed
- self._progress['status'] = 'completed'
- self._progress['end_time'] = time.time()
- logger.debug(
- "Example images download completed: %s/%s models processed",
- self._progress['completed'],
- self._progress['total'],
- )
- await self._broadcast_progress(status='completed')
+
+ async with self._state_lock:
+ if self._stop_requested and self._progress['status'] == 'stopping':
+ self._progress['status'] = 'stopped'
+ self._progress['end_time'] = time.time()
+ self._stop_requested = False
+ final_status = 'stopped'
+ elif self._progress['status'] not in {'error', 'stopped'}:
+ self._progress['status'] = 'completed'
+ self._progress['end_time'] = time.time()
+ self._stop_requested = False
+ final_status = 'completed'
+ else:
+ final_status = self._progress['status']
+ self._stop_requested = False
+ if self._progress['end_time'] is None:
+ self._progress['end_time'] = time.time()
+
+ if final_status == 'completed':
+ logger.debug(
+ "Example images download completed: %s/%s models processed",
+ self._progress['completed'],
+ self._progress['total'],
+ )
+ elif final_status == 'stopped':
+ logger.debug(
+ "Example images download stopped: %s/%s models processed",
+ self._progress['completed'],
+ self._progress['total'],
+ )
+
+ await self._broadcast_progress(status=final_status)
except Exception as e:
error_msg = f"Error during example images download: {str(e)}"
@@ -360,6 +424,7 @@ class DownloadManager:
async with self._state_lock:
self._is_downloading = False
self._download_task = None
+ self._stop_requested = False
async def _process_model(
self,
@@ -378,7 +443,7 @@ class DownloadManager:
await asyncio.sleep(1)
# Check if download should continue
- if self._progress['status'] != 'running':
+ if self._progress['status'] not in {'running', 'stopping'}:
logger.info(f"Download stopped: {self._progress['status']}")
return False # Return False to indicate no remote download happened
@@ -567,6 +632,7 @@ class DownloadManager:
raise DownloadConfigurationError('Example images path not configured in settings')
self._progress.reset()
+ self._stop_requested = False
self._progress['total'] = len(model_hashes)
self._progress['status'] = 'running'
self._progress['start_time'] = time.time()
@@ -588,10 +654,15 @@ class DownloadManager:
async with self._state_lock:
self._is_downloading = False
+ final_status = self._progress['status']
+
+ message = 'Force download completed'
+ if final_status == 'stopped':
+ message = 'Force download stopped'
return {
'success': True,
- 'message': 'Force download completed',
+ 'message': message,
'result': result
}
@@ -649,6 +720,12 @@ class DownloadManager:
# Process each model
success_count = 0
for i, (scanner_type, model, scanner) in enumerate(models_to_process):
+ async with self._state_lock:
+ current_status = self._progress['status']
+
+ if current_status not in {'running', 'paused', 'stopping'}:
+ break
+
# Force process this model regardless of previous status
was_successful = await self._process_specific_model(
scanner_type,
@@ -659,32 +736,65 @@ class DownloadManager:
downloader,
library_name,
)
-
+
if was_successful:
success_count += 1
-
+
# Update progress
self._progress['completed'] += 1
+ async with self._state_lock:
+ current_status = self._progress['status']
+ should_stop = self._stop_requested and current_status == 'stopping'
+
+ broadcast_status = 'running' if current_status == 'running' else current_status
# Send progress update via WebSocket
- await self._broadcast_progress(status='running')
-
+ await self._broadcast_progress(status=broadcast_status)
+
+ if should_stop:
+ break
+
# Only add delay after remote download, and not after processing the last model
- if was_successful and i < len(models_to_process) - 1 and self._progress['status'] == 'running':
+ if (
+ was_successful
+ and i < len(models_to_process) - 1
+ and current_status == 'running'
+ ):
await asyncio.sleep(delay)
-
- # Mark as completed
- self._progress['status'] = 'completed'
- self._progress['end_time'] = time.time()
- logger.debug(
- "Forced example images download completed: %s/%s models processed",
- self._progress['completed'],
- self._progress['total'],
- )
+
+ async with self._state_lock:
+ if self._stop_requested and self._progress['status'] == 'stopping':
+ self._progress['status'] = 'stopped'
+ self._progress['end_time'] = time.time()
+ self._stop_requested = False
+ final_status = 'stopped'
+ elif self._progress['status'] not in {'error', 'stopped'}:
+ self._progress['status'] = 'completed'
+ self._progress['end_time'] = time.time()
+ self._stop_requested = False
+ final_status = 'completed'
+ else:
+ final_status = self._progress['status']
+ self._stop_requested = False
+ if self._progress['end_time'] is None:
+ self._progress['end_time'] = time.time()
+
+ if final_status == 'completed':
+ logger.debug(
+ "Forced example images download completed: %s/%s models processed",
+ self._progress['completed'],
+ self._progress['total'],
+ )
+ elif final_status == 'stopped':
+ logger.debug(
+ "Forced example images download stopped: %s/%s models processed",
+ self._progress['completed'],
+ self._progress['total'],
+ )
# Send final progress via WebSocket
- await self._broadcast_progress(status='completed')
-
+ await self._broadcast_progress(status=final_status)
+
return {
'total': self._progress['total'],
'processed': self._progress['completed'],
@@ -726,7 +836,7 @@ class DownloadManager:
await asyncio.sleep(1)
# Check if download should continue
- if self._progress['status'] != 'running':
+ if self._progress['status'] not in {'running', 'stopping'}:
logger.info(f"Download stopped: {self._progress['status']}")
return False
diff --git a/static/js/managers/ExampleImagesManager.js b/static/js/managers/ExampleImagesManager.js
index d2a05fef..162f0f4f 100644
--- a/static/js/managers/ExampleImagesManager.js
+++ b/static/js/managers/ExampleImagesManager.js
@@ -13,8 +13,10 @@ export class ExampleImagesManager {
this.progressPanel = null;
this.isProgressPanelCollapsed = false;
this.pauseButton = null; // Store reference to the pause button
+ this.stopButton = null;
this.isMigrating = false; // Track migration state separately from downloading
this.hasShownCompletionToast = false; // Flag to track if completion toast has been shown
+ this.isStopping = false;
// Auto download properties
this.autoDownloadInterval = null;
@@ -52,11 +54,16 @@ export class ExampleImagesManager {
// Initialize progress panel button handlers
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
+ this.stopButton = document.getElementById('stopExampleDownloadBtn');
const collapseBtn = document.getElementById('collapseProgressBtn');
-
+
if (this.pauseButton) {
this.pauseButton.onclick = () => this.pauseDownload();
}
+
+ if (this.stopButton) {
+ this.stopButton.onclick = () => this.stopDownload();
+ }
if (collapseBtn) {
collapseBtn.onclick = () => this.toggleProgressPanel();
@@ -210,10 +217,14 @@ export class ExampleImagesManager {
updateDownloadButtonText() {
const btnTextElement = document.getElementById('exampleDownloadBtnText');
if (btnTextElement) {
- if (this.isDownloading && this.isPaused) {
+ if (this.isStopping) {
+ btnTextElement.textContent = "Stopping...";
+ } else if (this.isDownloading && this.isPaused) {
btnTextElement.textContent = "Resume";
} else if (!this.isDownloading) {
btnTextElement.textContent = "Download";
+ } else {
+ btnTextElement.textContent = "Download";
}
}
}
@@ -239,18 +250,22 @@ export class ExampleImagesManager {
});
const data = await response.json();
-
+
if (data.success) {
this.isDownloading = true;
this.isPaused = false;
+ this.isStopping = false;
this.hasShownCompletionToast = false; // Reset toast flag when starting new download
this.startTime = new Date();
this.updateUI(data.status);
this.showProgressPanel();
this.startProgressUpdates();
this.updateDownloadButtonText();
+ if (this.stopButton) {
+ this.stopButton.disabled = false;
+ }
showToast('toast.exampleImages.downloadStarted', {}, 'success');
-
+
// Close settings modal
modalManager.closeModal('settingsModal');
} else {
@@ -263,7 +278,7 @@ export class ExampleImagesManager {
}
async pauseDownload() {
- if (!this.isDownloading || this.isPaused) {
+ if (!this.isDownloading || this.isPaused || this.isStopping) {
return;
}
@@ -299,21 +314,21 @@ export class ExampleImagesManager {
}
async resumeDownload() {
- if (!this.isDownloading || !this.isPaused) {
+ if (!this.isDownloading || !this.isPaused || this.isStopping) {
return;
}
-
+
try {
const response = await fetch('/api/lm/resume-example-images', {
method: 'POST'
});
-
+
const data = await response.json();
-
+
if (data.success) {
this.isPaused = false;
document.getElementById('downloadStatusText').textContent = 'Downloading';
-
+
// Only update the icon element, not the entire innerHTML
if (this.pauseButton) {
const iconElement = this.pauseButton.querySelector('i');
@@ -322,7 +337,7 @@ export class ExampleImagesManager {
}
this.pauseButton.onclick = () => this.pauseDownload();
}
-
+
this.updateDownloadButtonText();
showToast('toast.exampleImages.downloadResumed', {}, 'success');
} else {
@@ -333,6 +348,60 @@ export class ExampleImagesManager {
showToast('toast.exampleImages.resumeFailed', {}, 'error');
}
}
+
+ async stopDownload() {
+ if (this.isStopping) {
+ return;
+ }
+
+ if (!this.isDownloading) {
+ this.hideProgressPanel();
+ return;
+ }
+
+ this.isStopping = true;
+ this.isPaused = false;
+ this.updateDownloadButtonText();
+
+ if (this.stopButton) {
+ this.stopButton.disabled = true;
+ }
+
+ try {
+ const response = await fetch('/api/lm/stop-example-images', {
+ method: 'POST'
+ });
+
+ let data;
+ try {
+ data = await response.json();
+ } catch (parseError) {
+ data = { success: false, error: 'Invalid server response' };
+ }
+
+ if (response.ok && data.success) {
+ showToast('toast.exampleImages.downloadStopped', {}, 'info');
+ this.hideProgressPanel();
+ } else {
+ this.isStopping = false;
+ if (this.stopButton) {
+ this.stopButton.disabled = false;
+ }
+ const errorMessage = data && data.error ? data.error : 'Unknown error';
+ showToast('toast.exampleImages.stopFailed', { error: errorMessage }, 'error');
+ }
+ } catch (error) {
+ console.error('Failed to stop download:', error);
+ this.isStopping = false;
+ if (this.stopButton) {
+ this.stopButton.disabled = false;
+ }
+ const errorMessage = error && error.message ? error.message : 'Unknown error';
+ showToast('toast.exampleImages.stopFailed', { error: errorMessage }, 'error');
+ } finally {
+ this.updateDownloadButtonText();
+ }
+ }
startProgressUpdates() {
// Clear any existing interval
@@ -352,21 +421,36 @@ export class ExampleImagesManager {
const data = await response.json();
if (data.success) {
+ const currentStatus = data.status.status;
this.isDownloading = data.is_downloading;
- this.isPaused = data.status.status === 'paused';
+ this.isPaused = currentStatus === 'paused';
this.isMigrating = data.is_migrating || false;
-
+
+ if (currentStatus === 'stopping') {
+ this.isStopping = true;
+ } else if (
+ !data.is_downloading ||
+ currentStatus === 'stopped' ||
+ currentStatus === 'completed' ||
+ currentStatus === 'error'
+ ) {
+ this.isStopping = false;
+ }
+
// Update download button text
this.updateDownloadButtonText();
-
+
if (this.isDownloading) {
this.updateUI(data.status);
} else {
// Download completed or failed
clearInterval(this.progressUpdateInterval);
this.progressUpdateInterval = null;
-
- if (data.status.status === 'completed' && !this.hasShownCompletionToast) {
+ if (this.stopButton) {
+ this.stopButton.disabled = true;
+ }
+
+ if (currentStatus === 'completed' && !this.hasShownCompletionToast) {
const actionType = this.isMigrating ? 'migration' : 'download';
showToast('toast.downloads.imagesCompleted', { action: actionType }, 'success');
// Mark as shown to prevent duplicate toasts
@@ -375,10 +459,13 @@ export class ExampleImagesManager {
this.isMigrating = false;
// Hide the panel after a delay
setTimeout(() => this.hideProgressPanel(), 5000);
- } else if (data.status.status === 'error') {
+ } else if (currentStatus === 'error') {
const actionType = this.isMigrating ? 'migration' : 'download';
showToast('toast.downloads.imagesFailed', { action: actionType }, 'error');
this.isMigrating = false;
+ } else if (currentStatus === 'stopped') {
+ this.hideProgressPanel();
+ this.isMigrating = false;
}
}
}
@@ -434,7 +521,11 @@ export class ExampleImagesManager {
if (!this.pauseButton) {
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
}
-
+
+ if (!this.stopButton) {
+ this.stopButton = document.getElementById('stopExampleDownloadBtn');
+ }
+
if (this.pauseButton) {
// Check if the button already has the SVG elements
let hasProgressElements = !!this.pauseButton.querySelector('.mini-progress-circle');
@@ -456,12 +547,14 @@ export class ExampleImagesManager {
iconElement.className = status.status === 'paused' ? 'fas fa-play' : 'fas fa-pause';
}
}
-
+
// Update click handler
- this.pauseButton.onclick = status.status === 'paused'
- ? () => this.resumeDownload()
+ this.pauseButton.onclick = status.status === 'paused'
+ ? () => this.resumeDownload()
: () => this.pauseDownload();
-
+
+ this.pauseButton.disabled = ['completed', 'error', 'stopped'].includes(status.status) || status.status === 'stopping';
+
// Update progress immediately
const progressBar = document.getElementById('downloadProgressBar');
if (progressBar) {
@@ -469,6 +562,15 @@ export class ExampleImagesManager {
this.updateMiniProgress(progressPercent);
}
}
+
+ if (this.stopButton) {
+ if (status.status === 'stopping' || this.isStopping) {
+ this.stopButton.disabled = true;
+ } else {
+ const canStop = ['running', 'paused'].includes(status.status);
+ this.stopButton.disabled = !canStop;
+ }
+ }
// Update title text
const titleElement = document.querySelector('.progress-panel-title');
@@ -584,6 +686,8 @@ export class ExampleImagesManager {
case 'paused': return 'Paused';
case 'completed': return 'Completed';
case 'error': return 'Error';
+ case 'stopping': return 'Stopping';
+ case 'stopped': return 'Stopped';
default: return 'Initializing';
}
}
diff --git a/templates/components/progress_panel.html b/templates/components/progress_panel.html
index eaee35f9..ad51a621 100644
--- a/templates/components/progress_panel.html
+++ b/templates/components/progress_panel.html
@@ -13,6 +13,9 @@
+
diff --git a/tests/routes/test_example_images_route_registrar_handlers.py b/tests/routes/test_example_images_route_registrar_handlers.py
index b4441f0b..1dceffc3 100644
--- a/tests/routes/test_example_images_route_registrar_handlers.py
+++ b/tests/routes/test_example_images_route_registrar_handlers.py
@@ -41,9 +41,11 @@ class StubDownloadManager:
def __init__(self) -> None:
self.pause_calls = 0
self.resume_calls = 0
+ self.stop_calls = 0
self.force_payloads: list[dict[str, Any]] = []
self.pause_error: Exception | None = None
self.resume_error: Exception | None = None
+ self.stop_error: Exception | None = None
self.force_error: Exception | None = None
async def get_status(self, request: web.Request) -> dict[str, Any]:
@@ -61,6 +63,12 @@ class StubDownloadManager:
raise self.resume_error
return {"success": True, "message": "resumed"}
+ async def stop_download(self, request: web.Request) -> dict[str, Any]:
+ self.stop_calls += 1
+ if self.stop_error:
+ raise self.stop_error
+ return {"success": True, "message": "stopping"}
+
async def start_force_download(self, payload: dict[str, Any]) -> dict[str, Any]:
self.force_payloads.append(payload)
if self.force_error:
@@ -193,17 +201,22 @@ async def test_pause_and_resume_return_client_errors_when_not_running():
async with registrar_app() as harness:
harness.download_manager.pause_error = DownloadNotRunningError()
harness.download_manager.resume_error = DownloadNotRunningError("Stopped")
+ harness.download_manager.stop_error = DownloadNotRunningError("Not running")
pause_response = await harness.client.post("/api/lm/pause-example-images")
resume_response = await harness.client.post("/api/lm/resume-example-images")
+ stop_response = await harness.client.post("/api/lm/stop-example-images")
assert pause_response.status == 400
assert resume_response.status == 400
+ assert stop_response.status == 400
pause_body = await _json(pause_response)
resume_body = await _json(resume_response)
+ stop_body = await _json(stop_response)
assert pause_body == {"success": False, "error": "No download in progress"}
assert resume_body == {"success": False, "error": "Stopped"}
+ assert stop_body == {"success": False, "error": "Not running"}
async def test_import_route_returns_validation_errors():
diff --git a/tests/routes/test_example_images_routes.py b/tests/routes/test_example_images_routes.py
index 1cc57eb0..10b305c1 100644
--- a/tests/routes/test_example_images_routes.py
+++ b/tests/routes/test_example_images_routes.py
@@ -51,6 +51,10 @@ class StubDownloadManager:
self.calls.append(("resume_download", None))
return {"operation": "resume_download"}
+ async def stop_download(self, request: web.Request) -> dict:
+ self.calls.append(("stop_download", None))
+ return {"operation": "stop_download"}
+
async def start_force_download(self, payload: Any) -> dict:
self.calls.append(("start_force_download", payload))
return {"operation": "start_force_download", "payload": payload}
@@ -195,19 +199,23 @@ async def test_status_route_returns_manager_payload():
assert harness.download_manager.calls == [("get_status", {"detail": "true"})]
-async def test_pause_and_resume_routes_delegate():
+async def test_pause_resume_and_stop_routes_delegate():
async with example_images_app() as harness:
pause_response = await harness.client.post("/api/lm/pause-example-images")
resume_response = await harness.client.post("/api/lm/resume-example-images")
+ stop_response = await harness.client.post("/api/lm/stop-example-images")
assert pause_response.status == 200
assert await pause_response.json() == {"operation": "pause_download"}
assert resume_response.status == 200
assert await resume_response.json() == {"operation": "resume_download"}
+ assert stop_response.status == 200
+ assert await stop_response.json() == {"operation": "stop_download"}
- assert harness.download_manager.calls[-2:] == [
+ assert harness.download_manager.calls[-3:] == [
("pause_download", None),
("resume_download", None),
+ ("stop_download", None),
]
@@ -309,6 +317,10 @@ async def test_download_handler_methods_delegate() -> None:
self.calls.append(("resume_download", request))
return {"status": "running"}
+ async def stop_download(self, request) -> dict:
+ self.calls.append(("stop_download", request))
+ return {"status": "stopping"}
+
async def start_force_download(self, payload) -> dict:
self.calls.append(("start_force_download", payload))
return {"status": "force", "payload": payload}
@@ -342,6 +354,8 @@ async def test_download_handler_methods_delegate() -> None:
assert json.loads(pause_response.text) == {"status": "paused"}
resume_response = await handler.resume_example_images(request)
assert json.loads(resume_response.text) == {"status": "running"}
+ stop_response = await handler.stop_example_images(request)
+ assert json.loads(stop_response.text) == {"status": "stopping"}
force_response = await handler.force_download_example_images(request)
assert json.loads(force_response.text) == {"status": "force", "payload": {"foo": "bar"}}
@@ -350,6 +364,7 @@ async def test_download_handler_methods_delegate() -> None:
("get_status", request),
("pause_download", request),
("resume_download", request),
+ ("stop_download", request),
("start_force_download", {"foo": "bar"}),
]
@@ -460,6 +475,7 @@ def test_handler_set_route_mapping_includes_all_handlers() -> None:
"get_example_images_status",
"pause_example_images",
"resume_example_images",
+ "stop_example_images",
"force_download_example_images",
"import_example_images",
"delete_example_image",
diff --git a/tests/utils/test_example_images_download_manager_unit.py b/tests/utils/test_example_images_download_manager_unit.py
index 42d88bf8..f6617ba6 100644
--- a/tests/utils/test_example_images_download_manager_unit.py
+++ b/tests/utils/test_example_images_download_manager_unit.py
@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
+import time
from typing import Any, Dict
import pytest
@@ -128,6 +129,59 @@ async def test_pause_and_resume_flow(monkeypatch: pytest.MonkeyPatch, tmp_path)
await asyncio.wait_for(task, timeout=1)
+async def test_stop_download_transitions_to_stopped(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None:
+ settings_manager = get_settings_manager()
+ settings_manager.settings["example_images_path"] = str(tmp_path)
+ settings_manager.settings["libraries"] = {"default": {}}
+ settings_manager.settings["active_library"] = "default"
+
+ ws_manager = RecordingWebSocketManager()
+ manager = download_module.DownloadManager(ws_manager=ws_manager)
+
+ started = asyncio.Event()
+ release = asyncio.Event()
+
+ async def fake_download(self, *_args):
+ started.set()
+ await release.wait()
+ async with self._state_lock:
+ if self._stop_requested and self._progress['status'] == 'stopping':
+ self._progress['status'] = 'stopped'
+ else:
+ self._progress['status'] = 'completed'
+ self._progress['end_time'] = time.time()
+ self._stop_requested = False
+ await self._broadcast_progress(status=self._progress['status'])
+ async with self._state_lock:
+ self._is_downloading = False
+ self._download_task = None
+
+ monkeypatch.setattr(
+ download_module.DownloadManager,
+ "_download_all_example_images",
+ fake_download,
+ )
+
+ await manager.start_download({})
+ await asyncio.wait_for(started.wait(), timeout=1)
+
+ stop_response = await manager.stop_download(object())
+ assert stop_response == {"success": True, "message": "Download stopping"}
+ assert manager._progress["status"] == "stopping"
+
+ task = manager._download_task
+ assert task is not None
+ release.set()
+ await asyncio.wait_for(task, timeout=1)
+
+ assert manager._progress["status"] == "stopped"
+ assert manager._is_downloading is False
+ assert manager._stop_requested is False
+ statuses = [payload["status"] for payload in ws_manager.payloads]
+ assert "stopping" in statuses
+ assert "stopped" in statuses
+
+
async def test_pause_or_resume_without_running_download(monkeypatch: pytest.MonkeyPatch) -> None:
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
@@ -136,3 +190,6 @@ async def test_pause_or_resume_without_running_download(monkeypatch: pytest.Monk
with pytest.raises(download_module.DownloadNotRunningError):
await manager.resume_download(object())
+
+ with pytest.raises(download_module.DownloadNotRunningError):
+ await manager.stop_download(object())