mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
feat(example-images): add stop control for download panel
This commit is contained in:
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "Fehler beim Pausieren des Downloads: {error}",
|
"pauseFailed": "Fehler beim Pausieren des Downloads: {error}",
|
||||||
"downloadResumed": "Download fortgesetzt",
|
"downloadResumed": "Download fortgesetzt",
|
||||||
"resumeFailed": "Fehler beim Fortsetzen des Downloads: {error}",
|
"resumeFailed": "Fehler beim Fortsetzen des Downloads: {error}",
|
||||||
|
"downloadStopped": "Download abgebrochen",
|
||||||
|
"stopFailed": "Download konnte nicht abgebrochen werden: {error}",
|
||||||
"deleted": "Beispielbild gelöscht",
|
"deleted": "Beispielbild gelöscht",
|
||||||
"deleteFailed": "Fehler beim Löschen des Beispielbilds",
|
"deleteFailed": "Fehler beim Löschen des Beispielbilds",
|
||||||
"setPreviewFailed": "Fehler beim Setzen des Vorschaubilds"
|
"setPreviewFailed": "Fehler beim Setzen des Vorschaubilds"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "Failed to pause download: {error}",
|
"pauseFailed": "Failed to pause download: {error}",
|
||||||
"downloadResumed": "Download resumed",
|
"downloadResumed": "Download resumed",
|
||||||
"resumeFailed": "Failed to resume download: {error}",
|
"resumeFailed": "Failed to resume download: {error}",
|
||||||
|
"downloadStopped": "Download cancelled",
|
||||||
|
"stopFailed": "Failed to cancel download: {error}",
|
||||||
"deleted": "Example image deleted",
|
"deleted": "Example image deleted",
|
||||||
"deleteFailed": "Failed to delete example image",
|
"deleteFailed": "Failed to delete example image",
|
||||||
"setPreviewFailed": "Failed to set preview image"
|
"setPreviewFailed": "Failed to set preview image"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "Error al pausar descarga: {error}",
|
"pauseFailed": "Error al pausar descarga: {error}",
|
||||||
"downloadResumed": "Descarga reanudada",
|
"downloadResumed": "Descarga reanudada",
|
||||||
"resumeFailed": "Error al reanudar descarga: {error}",
|
"resumeFailed": "Error al reanudar descarga: {error}",
|
||||||
|
"downloadStopped": "Descarga cancelada",
|
||||||
|
"stopFailed": "Error al cancelar descarga: {error}",
|
||||||
"deleted": "Imagen de ejemplo eliminada",
|
"deleted": "Imagen de ejemplo eliminada",
|
||||||
"deleteFailed": "Error al eliminar imagen de ejemplo",
|
"deleteFailed": "Error al eliminar imagen de ejemplo",
|
||||||
"setPreviewFailed": "Error al establecer imagen de vista previa"
|
"setPreviewFailed": "Error al establecer imagen de vista previa"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "Échec de la mise en pause du téléchargement : {error}",
|
"pauseFailed": "Échec de la mise en pause du téléchargement : {error}",
|
||||||
"downloadResumed": "Téléchargement repris",
|
"downloadResumed": "Téléchargement repris",
|
||||||
"resumeFailed": "Échec de la reprise du téléchargement : {error}",
|
"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",
|
"deleted": "Image d'exemple supprimée",
|
||||||
"deleteFailed": "Échec de la suppression de l'image d'exemple",
|
"deleteFailed": "Échec de la suppression de l'image d'exemple",
|
||||||
"setPreviewFailed": "Échec de la définition de l'image d'aperçu"
|
"setPreviewFailed": "Échec de la définition de l'image d'aperçu"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "השהיית ההורדה נכשלה: {error}",
|
"pauseFailed": "השהיית ההורדה נכשלה: {error}",
|
||||||
"downloadResumed": "ההורדה חודשה",
|
"downloadResumed": "ההורדה חודשה",
|
||||||
"resumeFailed": "חידוש ההורדה נכשל: {error}",
|
"resumeFailed": "חידוש ההורדה נכשל: {error}",
|
||||||
|
"downloadStopped": "ההורדה בוטלה",
|
||||||
|
"stopFailed": "נכשל בביטול ההורדה: {error}",
|
||||||
"deleted": "תמונת הדוגמה נמחקה",
|
"deleted": "תמונת הדוגמה נמחקה",
|
||||||
"deleteFailed": "מחיקת תמונת הדוגמה נכשלה",
|
"deleteFailed": "מחיקת תמונת הדוגמה נכשלה",
|
||||||
"setPreviewFailed": "הגדרת תמונת התצוגה המקדימה נכשלה"
|
"setPreviewFailed": "הגדרת תמונת התצוגה המקדימה נכשלה"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "ダウンロードの一時停止に失敗しました:{error}",
|
"pauseFailed": "ダウンロードの一時停止に失敗しました:{error}",
|
||||||
"downloadResumed": "ダウンロードが再開されました",
|
"downloadResumed": "ダウンロードが再開されました",
|
||||||
"resumeFailed": "ダウンロードの再開に失敗しました:{error}",
|
"resumeFailed": "ダウンロードの再開に失敗しました:{error}",
|
||||||
|
"downloadStopped": "ダウンロードをキャンセルしました",
|
||||||
|
"stopFailed": "ダウンロードのキャンセルに失敗しました:{error}",
|
||||||
"deleted": "例画像が削除されました",
|
"deleted": "例画像が削除されました",
|
||||||
"deleteFailed": "例画像の削除に失敗しました",
|
"deleteFailed": "例画像の削除に失敗しました",
|
||||||
"setPreviewFailed": "プレビュー画像の設定に失敗しました"
|
"setPreviewFailed": "プレビュー画像の設定に失敗しました"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "다운로드 일시정지 실패: {error}",
|
"pauseFailed": "다운로드 일시정지 실패: {error}",
|
||||||
"downloadResumed": "다운로드가 재개되었습니다",
|
"downloadResumed": "다운로드가 재개되었습니다",
|
||||||
"resumeFailed": "다운로드 재개 실패: {error}",
|
"resumeFailed": "다운로드 재개 실패: {error}",
|
||||||
|
"downloadStopped": "다운로드가 취소되었습니다",
|
||||||
|
"stopFailed": "다운로드 취소 실패: {error}",
|
||||||
"deleted": "예시 이미지가 삭제되었습니다",
|
"deleted": "예시 이미지가 삭제되었습니다",
|
||||||
"deleteFailed": "예시 이미지 삭제 실패",
|
"deleteFailed": "예시 이미지 삭제 실패",
|
||||||
"setPreviewFailed": "미리보기 이미지 설정 실패"
|
"setPreviewFailed": "미리보기 이미지 설정 실패"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "Не удалось приостановить загрузку: {error}",
|
"pauseFailed": "Не удалось приостановить загрузку: {error}",
|
||||||
"downloadResumed": "Загрузка возобновлена",
|
"downloadResumed": "Загрузка возобновлена",
|
||||||
"resumeFailed": "Не удалось возобновить загрузку: {error}",
|
"resumeFailed": "Не удалось возобновить загрузку: {error}",
|
||||||
|
"downloadStopped": "Загрузка отменена",
|
||||||
|
"stopFailed": "Не удалось отменить загрузку: {error}",
|
||||||
"deleted": "Пример изображения удален",
|
"deleted": "Пример изображения удален",
|
||||||
"deleteFailed": "Не удалось удалить пример изображения",
|
"deleteFailed": "Не удалось удалить пример изображения",
|
||||||
"setPreviewFailed": "Не удалось установить превью изображение"
|
"setPreviewFailed": "Не удалось установить превью изображение"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "暂停下载失败:{error}",
|
"pauseFailed": "暂停下载失败:{error}",
|
||||||
"downloadResumed": "下载已恢复",
|
"downloadResumed": "下载已恢复",
|
||||||
"resumeFailed": "恢复下载失败:{error}",
|
"resumeFailed": "恢复下载失败:{error}",
|
||||||
|
"downloadStopped": "下载已取消",
|
||||||
|
"stopFailed": "取消下载失败:{error}",
|
||||||
"deleted": "示例图片已删除",
|
"deleted": "示例图片已删除",
|
||||||
"deleteFailed": "删除示例图片失败",
|
"deleteFailed": "删除示例图片失败",
|
||||||
"setPreviewFailed": "设置预览图片失败"
|
"setPreviewFailed": "设置预览图片失败"
|
||||||
|
|||||||
@@ -1245,6 +1245,8 @@
|
|||||||
"pauseFailed": "暫停下載失敗:{error}",
|
"pauseFailed": "暫停下載失敗:{error}",
|
||||||
"downloadResumed": "下載已恢復",
|
"downloadResumed": "下載已恢復",
|
||||||
"resumeFailed": "恢復下載失敗:{error}",
|
"resumeFailed": "恢復下載失敗:{error}",
|
||||||
|
"downloadStopped": "下載已取消",
|
||||||
|
"stopFailed": "取消下載失敗:{error}",
|
||||||
"deleted": "範例圖片已刪除",
|
"deleted": "範例圖片已刪除",
|
||||||
"deleteFailed": "刪除範例圖片失敗",
|
"deleteFailed": "刪除範例圖片失敗",
|
||||||
"setPreviewFailed": "設定預覽圖片失敗"
|
"setPreviewFailed": "設定預覽圖片失敗"
|
||||||
|
|||||||
@@ -22,6 +22,7 @@ ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
|
|||||||
RouteDefinition("GET", "/api/lm/example-images-status", "get_example_images_status"),
|
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/pause-example-images", "pause_example_images"),
|
||||||
RouteDefinition("POST", "/api/lm/resume-example-images", "resume_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("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/example-image-files", "get_example_image_files"),
|
||||||
RouteDefinition("GET", "/api/lm/has-example-images", "has_example_images"),
|
RouteDefinition("GET", "/api/lm/has-example-images", "has_example_images"),
|
||||||
|
|||||||
@@ -68,6 +68,13 @@ class ExampleImagesDownloadHandler:
|
|||||||
except DownloadNotRunningError as exc:
|
except DownloadNotRunningError as exc:
|
||||||
return web.json_response({'success': False, 'error': str(exc)}, status=400)
|
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:
|
async def force_download_example_images(self, request: web.Request) -> web.StreamResponse:
|
||||||
try:
|
try:
|
||||||
payload = await request.json()
|
payload = await request.json()
|
||||||
@@ -149,6 +156,7 @@ class ExampleImagesHandlerSet:
|
|||||||
"get_example_images_status": self.download.get_example_images_status,
|
"get_example_images_status": self.download.get_example_images_status,
|
||||||
"pause_example_images": self.download.pause_example_images,
|
"pause_example_images": self.download.pause_example_images,
|
||||||
"resume_example_images": self.download.resume_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,
|
"force_download_example_images": self.download.force_download_example_images,
|
||||||
"import_example_images": self.management.import_example_images,
|
"import_example_images": self.management.import_example_images,
|
||||||
"delete_example_image": self.management.delete_example_image,
|
"delete_example_image": self.management.delete_example_image,
|
||||||
|
|||||||
@@ -105,6 +105,7 @@ class DownloadManager:
|
|||||||
self._progress = _DownloadProgress()
|
self._progress = _DownloadProgress()
|
||||||
self._ws_manager = ws_manager
|
self._ws_manager = ws_manager
|
||||||
self._state_lock = state_lock or asyncio.Lock()
|
self._state_lock = state_lock or asyncio.Lock()
|
||||||
|
self._stop_requested = False
|
||||||
|
|
||||||
def _resolve_output_dir(self, library_name: str | None = None) -> str:
|
def _resolve_output_dir(self, library_name: str | None = None) -> str:
|
||||||
base_path = get_settings_manager().get('example_images_path')
|
base_path = get_settings_manager().get('example_images_path')
|
||||||
@@ -145,6 +146,7 @@ class DownloadManager:
|
|||||||
raise DownloadConfigurationError('Example images path not configured in settings')
|
raise DownloadConfigurationError('Example images path not configured in settings')
|
||||||
|
|
||||||
self._progress.reset()
|
self._progress.reset()
|
||||||
|
self._stop_requested = False
|
||||||
self._progress['status'] = 'running'
|
self._progress['status'] = 'running'
|
||||||
self._progress['start_time'] = time.time()
|
self._progress['start_time'] = time.time()
|
||||||
self._progress['end_time'] = None
|
self._progress['end_time'] = None
|
||||||
@@ -267,6 +269,27 @@ class DownloadManager:
|
|||||||
'success': True,
|
'success': True,
|
||||||
'message': 'Download resumed'
|
'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(
|
async def _download_all_example_images(
|
||||||
self,
|
self,
|
||||||
@@ -311,6 +334,12 @@ class DownloadManager:
|
|||||||
|
|
||||||
# Process each model
|
# Process each model
|
||||||
for i, (scanner_type, model, scanner) in enumerate(all_models):
|
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
|
# Main logic for processing model is here, but actual operations are delegated to other classes
|
||||||
was_remote_download = await self._process_model(
|
was_remote_download = await self._process_model(
|
||||||
scanner_type,
|
scanner_type,
|
||||||
@@ -321,24 +350,59 @@ class DownloadManager:
|
|||||||
downloader,
|
downloader,
|
||||||
library_name,
|
library_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Update progress
|
# Update progress
|
||||||
self._progress['completed'] += 1
|
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
|
# 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)
|
await asyncio.sleep(delay)
|
||||||
|
|
||||||
# Mark as completed
|
async with self._state_lock:
|
||||||
self._progress['status'] = 'completed'
|
if self._stop_requested and self._progress['status'] == 'stopping':
|
||||||
self._progress['end_time'] = time.time()
|
self._progress['status'] = 'stopped'
|
||||||
logger.debug(
|
self._progress['end_time'] = time.time()
|
||||||
"Example images download completed: %s/%s models processed",
|
self._stop_requested = False
|
||||||
self._progress['completed'],
|
final_status = 'stopped'
|
||||||
self._progress['total'],
|
elif self._progress['status'] not in {'error', 'stopped'}:
|
||||||
)
|
self._progress['status'] = 'completed'
|
||||||
await self._broadcast_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:
|
except Exception as e:
|
||||||
error_msg = f"Error during example images download: {str(e)}"
|
error_msg = f"Error during example images download: {str(e)}"
|
||||||
@@ -360,6 +424,7 @@ class DownloadManager:
|
|||||||
async with self._state_lock:
|
async with self._state_lock:
|
||||||
self._is_downloading = False
|
self._is_downloading = False
|
||||||
self._download_task = None
|
self._download_task = None
|
||||||
|
self._stop_requested = False
|
||||||
|
|
||||||
async def _process_model(
|
async def _process_model(
|
||||||
self,
|
self,
|
||||||
@@ -378,7 +443,7 @@ class DownloadManager:
|
|||||||
await asyncio.sleep(1)
|
await asyncio.sleep(1)
|
||||||
|
|
||||||
# Check if download should continue
|
# 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']}")
|
logger.info(f"Download stopped: {self._progress['status']}")
|
||||||
return False # Return False to indicate no remote download happened
|
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')
|
raise DownloadConfigurationError('Example images path not configured in settings')
|
||||||
|
|
||||||
self._progress.reset()
|
self._progress.reset()
|
||||||
|
self._stop_requested = False
|
||||||
self._progress['total'] = len(model_hashes)
|
self._progress['total'] = len(model_hashes)
|
||||||
self._progress['status'] = 'running'
|
self._progress['status'] = 'running'
|
||||||
self._progress['start_time'] = time.time()
|
self._progress['start_time'] = time.time()
|
||||||
@@ -588,10 +654,15 @@ class DownloadManager:
|
|||||||
|
|
||||||
async with self._state_lock:
|
async with self._state_lock:
|
||||||
self._is_downloading = False
|
self._is_downloading = False
|
||||||
|
final_status = self._progress['status']
|
||||||
|
|
||||||
|
message = 'Force download completed'
|
||||||
|
if final_status == 'stopped':
|
||||||
|
message = 'Force download stopped'
|
||||||
|
|
||||||
return {
|
return {
|
||||||
'success': True,
|
'success': True,
|
||||||
'message': 'Force download completed',
|
'message': message,
|
||||||
'result': result
|
'result': result
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -649,6 +720,12 @@ class DownloadManager:
|
|||||||
# Process each model
|
# Process each model
|
||||||
success_count = 0
|
success_count = 0
|
||||||
for i, (scanner_type, model, scanner) in enumerate(models_to_process):
|
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
|
# Force process this model regardless of previous status
|
||||||
was_successful = await self._process_specific_model(
|
was_successful = await self._process_specific_model(
|
||||||
scanner_type,
|
scanner_type,
|
||||||
@@ -659,32 +736,65 @@ class DownloadManager:
|
|||||||
downloader,
|
downloader,
|
||||||
library_name,
|
library_name,
|
||||||
)
|
)
|
||||||
|
|
||||||
if was_successful:
|
if was_successful:
|
||||||
success_count += 1
|
success_count += 1
|
||||||
|
|
||||||
# Update progress
|
# Update progress
|
||||||
self._progress['completed'] += 1
|
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
|
# 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
|
# 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)
|
await asyncio.sleep(delay)
|
||||||
|
|
||||||
# Mark as completed
|
async with self._state_lock:
|
||||||
self._progress['status'] = 'completed'
|
if self._stop_requested and self._progress['status'] == 'stopping':
|
||||||
self._progress['end_time'] = time.time()
|
self._progress['status'] = 'stopped'
|
||||||
logger.debug(
|
self._progress['end_time'] = time.time()
|
||||||
"Forced example images download completed: %s/%s models processed",
|
self._stop_requested = False
|
||||||
self._progress['completed'],
|
final_status = 'stopped'
|
||||||
self._progress['total'],
|
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
|
# Send final progress via WebSocket
|
||||||
await self._broadcast_progress(status='completed')
|
await self._broadcast_progress(status=final_status)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
'total': self._progress['total'],
|
'total': self._progress['total'],
|
||||||
'processed': self._progress['completed'],
|
'processed': self._progress['completed'],
|
||||||
@@ -726,7 +836,7 @@ class DownloadManager:
|
|||||||
await asyncio.sleep(1)
|
await asyncio.sleep(1)
|
||||||
|
|
||||||
# Check if download should continue
|
# 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']}")
|
logger.info(f"Download stopped: {self._progress['status']}")
|
||||||
return False
|
return False
|
||||||
|
|
||||||
|
|||||||
@@ -13,8 +13,10 @@ export class ExampleImagesManager {
|
|||||||
this.progressPanel = null;
|
this.progressPanel = null;
|
||||||
this.isProgressPanelCollapsed = false;
|
this.isProgressPanelCollapsed = false;
|
||||||
this.pauseButton = null; // Store reference to the pause button
|
this.pauseButton = null; // Store reference to the pause button
|
||||||
|
this.stopButton = null;
|
||||||
this.isMigrating = false; // Track migration state separately from downloading
|
this.isMigrating = false; // Track migration state separately from downloading
|
||||||
this.hasShownCompletionToast = false; // Flag to track if completion toast has been shown
|
this.hasShownCompletionToast = false; // Flag to track if completion toast has been shown
|
||||||
|
this.isStopping = false;
|
||||||
|
|
||||||
// Auto download properties
|
// Auto download properties
|
||||||
this.autoDownloadInterval = null;
|
this.autoDownloadInterval = null;
|
||||||
@@ -52,11 +54,16 @@ export class ExampleImagesManager {
|
|||||||
|
|
||||||
// Initialize progress panel button handlers
|
// Initialize progress panel button handlers
|
||||||
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
|
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
|
||||||
|
this.stopButton = document.getElementById('stopExampleDownloadBtn');
|
||||||
const collapseBtn = document.getElementById('collapseProgressBtn');
|
const collapseBtn = document.getElementById('collapseProgressBtn');
|
||||||
|
|
||||||
if (this.pauseButton) {
|
if (this.pauseButton) {
|
||||||
this.pauseButton.onclick = () => this.pauseDownload();
|
this.pauseButton.onclick = () => this.pauseDownload();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (this.stopButton) {
|
||||||
|
this.stopButton.onclick = () => this.stopDownload();
|
||||||
|
}
|
||||||
|
|
||||||
if (collapseBtn) {
|
if (collapseBtn) {
|
||||||
collapseBtn.onclick = () => this.toggleProgressPanel();
|
collapseBtn.onclick = () => this.toggleProgressPanel();
|
||||||
@@ -210,10 +217,14 @@ export class ExampleImagesManager {
|
|||||||
updateDownloadButtonText() {
|
updateDownloadButtonText() {
|
||||||
const btnTextElement = document.getElementById('exampleDownloadBtnText');
|
const btnTextElement = document.getElementById('exampleDownloadBtnText');
|
||||||
if (btnTextElement) {
|
if (btnTextElement) {
|
||||||
if (this.isDownloading && this.isPaused) {
|
if (this.isStopping) {
|
||||||
|
btnTextElement.textContent = "Stopping...";
|
||||||
|
} else if (this.isDownloading && this.isPaused) {
|
||||||
btnTextElement.textContent = "Resume";
|
btnTextElement.textContent = "Resume";
|
||||||
} else if (!this.isDownloading) {
|
} else if (!this.isDownloading) {
|
||||||
btnTextElement.textContent = "Download";
|
btnTextElement.textContent = "Download";
|
||||||
|
} else {
|
||||||
|
btnTextElement.textContent = "Download";
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -239,18 +250,22 @@ export class ExampleImagesManager {
|
|||||||
});
|
});
|
||||||
|
|
||||||
const data = await response.json();
|
const data = await response.json();
|
||||||
|
|
||||||
if (data.success) {
|
if (data.success) {
|
||||||
this.isDownloading = true;
|
this.isDownloading = true;
|
||||||
this.isPaused = false;
|
this.isPaused = false;
|
||||||
|
this.isStopping = false;
|
||||||
this.hasShownCompletionToast = false; // Reset toast flag when starting new download
|
this.hasShownCompletionToast = false; // Reset toast flag when starting new download
|
||||||
this.startTime = new Date();
|
this.startTime = new Date();
|
||||||
this.updateUI(data.status);
|
this.updateUI(data.status);
|
||||||
this.showProgressPanel();
|
this.showProgressPanel();
|
||||||
this.startProgressUpdates();
|
this.startProgressUpdates();
|
||||||
this.updateDownloadButtonText();
|
this.updateDownloadButtonText();
|
||||||
|
if (this.stopButton) {
|
||||||
|
this.stopButton.disabled = false;
|
||||||
|
}
|
||||||
showToast('toast.exampleImages.downloadStarted', {}, 'success');
|
showToast('toast.exampleImages.downloadStarted', {}, 'success');
|
||||||
|
|
||||||
// Close settings modal
|
// Close settings modal
|
||||||
modalManager.closeModal('settingsModal');
|
modalManager.closeModal('settingsModal');
|
||||||
} else {
|
} else {
|
||||||
@@ -263,7 +278,7 @@ export class ExampleImagesManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async pauseDownload() {
|
async pauseDownload() {
|
||||||
if (!this.isDownloading || this.isPaused) {
|
if (!this.isDownloading || this.isPaused || this.isStopping) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -299,21 +314,21 @@ export class ExampleImagesManager {
|
|||||||
}
|
}
|
||||||
|
|
||||||
async resumeDownload() {
|
async resumeDownload() {
|
||||||
if (!this.isDownloading || !this.isPaused) {
|
if (!this.isDownloading || !this.isPaused || this.isStopping) {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const response = await fetch('/api/lm/resume-example-images', {
|
const response = await fetch('/api/lm/resume-example-images', {
|
||||||
method: 'POST'
|
method: 'POST'
|
||||||
});
|
});
|
||||||
|
|
||||||
const data = await response.json();
|
const data = await response.json();
|
||||||
|
|
||||||
if (data.success) {
|
if (data.success) {
|
||||||
this.isPaused = false;
|
this.isPaused = false;
|
||||||
document.getElementById('downloadStatusText').textContent = 'Downloading';
|
document.getElementById('downloadStatusText').textContent = 'Downloading';
|
||||||
|
|
||||||
// Only update the icon element, not the entire innerHTML
|
// Only update the icon element, not the entire innerHTML
|
||||||
if (this.pauseButton) {
|
if (this.pauseButton) {
|
||||||
const iconElement = this.pauseButton.querySelector('i');
|
const iconElement = this.pauseButton.querySelector('i');
|
||||||
@@ -322,7 +337,7 @@ export class ExampleImagesManager {
|
|||||||
}
|
}
|
||||||
this.pauseButton.onclick = () => this.pauseDownload();
|
this.pauseButton.onclick = () => this.pauseDownload();
|
||||||
}
|
}
|
||||||
|
|
||||||
this.updateDownloadButtonText();
|
this.updateDownloadButtonText();
|
||||||
showToast('toast.exampleImages.downloadResumed', {}, 'success');
|
showToast('toast.exampleImages.downloadResumed', {}, 'success');
|
||||||
} else {
|
} else {
|
||||||
@@ -333,6 +348,60 @@ export class ExampleImagesManager {
|
|||||||
showToast('toast.exampleImages.resumeFailed', {}, 'error');
|
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() {
|
startProgressUpdates() {
|
||||||
// Clear any existing interval
|
// Clear any existing interval
|
||||||
@@ -352,21 +421,36 @@ export class ExampleImagesManager {
|
|||||||
const data = await response.json();
|
const data = await response.json();
|
||||||
|
|
||||||
if (data.success) {
|
if (data.success) {
|
||||||
|
const currentStatus = data.status.status;
|
||||||
this.isDownloading = data.is_downloading;
|
this.isDownloading = data.is_downloading;
|
||||||
this.isPaused = data.status.status === 'paused';
|
this.isPaused = currentStatus === 'paused';
|
||||||
this.isMigrating = data.is_migrating || false;
|
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
|
// Update download button text
|
||||||
this.updateDownloadButtonText();
|
this.updateDownloadButtonText();
|
||||||
|
|
||||||
if (this.isDownloading) {
|
if (this.isDownloading) {
|
||||||
this.updateUI(data.status);
|
this.updateUI(data.status);
|
||||||
} else {
|
} else {
|
||||||
// Download completed or failed
|
// Download completed or failed
|
||||||
clearInterval(this.progressUpdateInterval);
|
clearInterval(this.progressUpdateInterval);
|
||||||
this.progressUpdateInterval = null;
|
this.progressUpdateInterval = null;
|
||||||
|
if (this.stopButton) {
|
||||||
if (data.status.status === 'completed' && !this.hasShownCompletionToast) {
|
this.stopButton.disabled = true;
|
||||||
|
}
|
||||||
|
|
||||||
|
if (currentStatus === 'completed' && !this.hasShownCompletionToast) {
|
||||||
const actionType = this.isMigrating ? 'migration' : 'download';
|
const actionType = this.isMigrating ? 'migration' : 'download';
|
||||||
showToast('toast.downloads.imagesCompleted', { action: actionType }, 'success');
|
showToast('toast.downloads.imagesCompleted', { action: actionType }, 'success');
|
||||||
// Mark as shown to prevent duplicate toasts
|
// Mark as shown to prevent duplicate toasts
|
||||||
@@ -375,10 +459,13 @@ export class ExampleImagesManager {
|
|||||||
this.isMigrating = false;
|
this.isMigrating = false;
|
||||||
// Hide the panel after a delay
|
// Hide the panel after a delay
|
||||||
setTimeout(() => this.hideProgressPanel(), 5000);
|
setTimeout(() => this.hideProgressPanel(), 5000);
|
||||||
} else if (data.status.status === 'error') {
|
} else if (currentStatus === 'error') {
|
||||||
const actionType = this.isMigrating ? 'migration' : 'download';
|
const actionType = this.isMigrating ? 'migration' : 'download';
|
||||||
showToast('toast.downloads.imagesFailed', { action: actionType }, 'error');
|
showToast('toast.downloads.imagesFailed', { action: actionType }, 'error');
|
||||||
this.isMigrating = false;
|
this.isMigrating = false;
|
||||||
|
} else if (currentStatus === 'stopped') {
|
||||||
|
this.hideProgressPanel();
|
||||||
|
this.isMigrating = false;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -434,7 +521,11 @@ export class ExampleImagesManager {
|
|||||||
if (!this.pauseButton) {
|
if (!this.pauseButton) {
|
||||||
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
|
this.pauseButton = document.getElementById('pauseExampleDownloadBtn');
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (!this.stopButton) {
|
||||||
|
this.stopButton = document.getElementById('stopExampleDownloadBtn');
|
||||||
|
}
|
||||||
|
|
||||||
if (this.pauseButton) {
|
if (this.pauseButton) {
|
||||||
// Check if the button already has the SVG elements
|
// Check if the button already has the SVG elements
|
||||||
let hasProgressElements = !!this.pauseButton.querySelector('.mini-progress-circle');
|
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';
|
iconElement.className = status.status === 'paused' ? 'fas fa-play' : 'fas fa-pause';
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// Update click handler
|
// Update click handler
|
||||||
this.pauseButton.onclick = status.status === 'paused'
|
this.pauseButton.onclick = status.status === 'paused'
|
||||||
? () => this.resumeDownload()
|
? () => this.resumeDownload()
|
||||||
: () => this.pauseDownload();
|
: () => this.pauseDownload();
|
||||||
|
|
||||||
|
this.pauseButton.disabled = ['completed', 'error', 'stopped'].includes(status.status) || status.status === 'stopping';
|
||||||
|
|
||||||
// Update progress immediately
|
// Update progress immediately
|
||||||
const progressBar = document.getElementById('downloadProgressBar');
|
const progressBar = document.getElementById('downloadProgressBar');
|
||||||
if (progressBar) {
|
if (progressBar) {
|
||||||
@@ -469,6 +562,15 @@ export class ExampleImagesManager {
|
|||||||
this.updateMiniProgress(progressPercent);
|
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
|
// Update title text
|
||||||
const titleElement = document.querySelector('.progress-panel-title');
|
const titleElement = document.querySelector('.progress-panel-title');
|
||||||
@@ -584,6 +686,8 @@ export class ExampleImagesManager {
|
|||||||
case 'paused': return 'Paused';
|
case 'paused': return 'Paused';
|
||||||
case 'completed': return 'Completed';
|
case 'completed': return 'Completed';
|
||||||
case 'error': return 'Error';
|
case 'error': return 'Error';
|
||||||
|
case 'stopping': return 'Stopping';
|
||||||
|
case 'stopped': return 'Stopped';
|
||||||
default: return 'Initializing';
|
default: return 'Initializing';
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -13,6 +13,9 @@
|
|||||||
</svg>
|
</svg>
|
||||||
<span class="progress-percent"></span>
|
<span class="progress-percent"></span>
|
||||||
</button>
|
</button>
|
||||||
|
<button id="stopExampleDownloadBtn" class="icon-button">
|
||||||
|
<i class="fas fa-stop"></i>
|
||||||
|
</button>
|
||||||
<button id="collapseProgressBtn" class="icon-button">
|
<button id="collapseProgressBtn" class="icon-button">
|
||||||
<i class="fas fa-chevron-down"></i>
|
<i class="fas fa-chevron-down"></i>
|
||||||
</button>
|
</button>
|
||||||
|
|||||||
@@ -41,9 +41,11 @@ class StubDownloadManager:
|
|||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.pause_calls = 0
|
self.pause_calls = 0
|
||||||
self.resume_calls = 0
|
self.resume_calls = 0
|
||||||
|
self.stop_calls = 0
|
||||||
self.force_payloads: list[dict[str, Any]] = []
|
self.force_payloads: list[dict[str, Any]] = []
|
||||||
self.pause_error: Exception | None = None
|
self.pause_error: Exception | None = None
|
||||||
self.resume_error: Exception | None = None
|
self.resume_error: Exception | None = None
|
||||||
|
self.stop_error: Exception | None = None
|
||||||
self.force_error: Exception | None = None
|
self.force_error: Exception | None = None
|
||||||
|
|
||||||
async def get_status(self, request: web.Request) -> dict[str, Any]:
|
async def get_status(self, request: web.Request) -> dict[str, Any]:
|
||||||
@@ -61,6 +63,12 @@ class StubDownloadManager:
|
|||||||
raise self.resume_error
|
raise self.resume_error
|
||||||
return {"success": True, "message": "resumed"}
|
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]:
|
async def start_force_download(self, payload: dict[str, Any]) -> dict[str, Any]:
|
||||||
self.force_payloads.append(payload)
|
self.force_payloads.append(payload)
|
||||||
if self.force_error:
|
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:
|
async with registrar_app() as harness:
|
||||||
harness.download_manager.pause_error = DownloadNotRunningError()
|
harness.download_manager.pause_error = DownloadNotRunningError()
|
||||||
harness.download_manager.resume_error = DownloadNotRunningError("Stopped")
|
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")
|
pause_response = await harness.client.post("/api/lm/pause-example-images")
|
||||||
resume_response = await harness.client.post("/api/lm/resume-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 pause_response.status == 400
|
||||||
assert resume_response.status == 400
|
assert resume_response.status == 400
|
||||||
|
assert stop_response.status == 400
|
||||||
|
|
||||||
pause_body = await _json(pause_response)
|
pause_body = await _json(pause_response)
|
||||||
resume_body = await _json(resume_response)
|
resume_body = await _json(resume_response)
|
||||||
|
stop_body = await _json(stop_response)
|
||||||
assert pause_body == {"success": False, "error": "No download in progress"}
|
assert pause_body == {"success": False, "error": "No download in progress"}
|
||||||
assert resume_body == {"success": False, "error": "Stopped"}
|
assert resume_body == {"success": False, "error": "Stopped"}
|
||||||
|
assert stop_body == {"success": False, "error": "Not running"}
|
||||||
|
|
||||||
|
|
||||||
async def test_import_route_returns_validation_errors():
|
async def test_import_route_returns_validation_errors():
|
||||||
|
|||||||
@@ -51,6 +51,10 @@ class StubDownloadManager:
|
|||||||
self.calls.append(("resume_download", None))
|
self.calls.append(("resume_download", None))
|
||||||
return {"operation": "resume_download"}
|
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:
|
async def start_force_download(self, payload: Any) -> dict:
|
||||||
self.calls.append(("start_force_download", payload))
|
self.calls.append(("start_force_download", payload))
|
||||||
return {"operation": "start_force_download", "payload": 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"})]
|
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:
|
async with example_images_app() as harness:
|
||||||
pause_response = await harness.client.post("/api/lm/pause-example-images")
|
pause_response = await harness.client.post("/api/lm/pause-example-images")
|
||||||
resume_response = await harness.client.post("/api/lm/resume-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 pause_response.status == 200
|
||||||
assert await pause_response.json() == {"operation": "pause_download"}
|
assert await pause_response.json() == {"operation": "pause_download"}
|
||||||
assert resume_response.status == 200
|
assert resume_response.status == 200
|
||||||
assert await resume_response.json() == {"operation": "resume_download"}
|
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),
|
("pause_download", None),
|
||||||
("resume_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))
|
self.calls.append(("resume_download", request))
|
||||||
return {"status": "running"}
|
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:
|
async def start_force_download(self, payload) -> dict:
|
||||||
self.calls.append(("start_force_download", payload))
|
self.calls.append(("start_force_download", payload))
|
||||||
return {"status": "force", "payload": 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"}
|
assert json.loads(pause_response.text) == {"status": "paused"}
|
||||||
resume_response = await handler.resume_example_images(request)
|
resume_response = await handler.resume_example_images(request)
|
||||||
assert json.loads(resume_response.text) == {"status": "running"}
|
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)
|
force_response = await handler.force_download_example_images(request)
|
||||||
assert json.loads(force_response.text) == {"status": "force", "payload": {"foo": "bar"}}
|
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),
|
("get_status", request),
|
||||||
("pause_download", request),
|
("pause_download", request),
|
||||||
("resume_download", request),
|
("resume_download", request),
|
||||||
|
("stop_download", request),
|
||||||
("start_force_download", {"foo": "bar"}),
|
("start_force_download", {"foo": "bar"}),
|
||||||
]
|
]
|
||||||
|
|
||||||
@@ -460,6 +475,7 @@ def test_handler_set_route_mapping_includes_all_handlers() -> None:
|
|||||||
"get_example_images_status",
|
"get_example_images_status",
|
||||||
"pause_example_images",
|
"pause_example_images",
|
||||||
"resume_example_images",
|
"resume_example_images",
|
||||||
|
"stop_example_images",
|
||||||
"force_download_example_images",
|
"force_download_example_images",
|
||||||
"import_example_images",
|
"import_example_images",
|
||||||
"delete_example_image",
|
"delete_example_image",
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import asyncio
|
import asyncio
|
||||||
|
import time
|
||||||
from typing import Any, Dict
|
from typing import Any, Dict
|
||||||
|
|
||||||
import pytest
|
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)
|
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:
|
async def test_pause_or_resume_without_running_download(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
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):
|
with pytest.raises(download_module.DownloadNotRunningError):
|
||||||
await manager.resume_download(object())
|
await manager.resume_download(object())
|
||||||
|
|
||||||
|
with pytest.raises(download_module.DownloadNotRunningError):
|
||||||
|
await manager.stop_download(object())
|
||||||
|
|||||||
Reference in New Issue
Block a user