Skip to content

柴进 / GoogleNanoBananaApp

Go to a project
Toggle navigation
Toggle navigation pinning
f29b806b by shady
Options

Merge remote-tracking branch 'origin/master'

2 parents f4ceb240 68b19c6f
Inline Side-by-side
Showing 3 changed files with 980 additions and 1 deletions
...@@ -14,7 +14,7 @@ ...@@ -14,7 +14,7 @@
14 "database": "saas_user", 14 "database": "saas_user",
15 "table": "nano_banana_users" 15 "table": "nano_banana_users"
16 }, 16 },
17 "last_user": "testuser", 17 "last_user": "chaijin",
18 "saved_password_hash": "50630320e4a550f2dba371820dad9d9301d456d101aca4d5ad8f4f3bcc9c1ed9", 18 "saved_password_hash": "50630320e4a550f2dba371820dad9d9301d456d101aca4d5ad8f4f3bcc9c1ed9",
19 "logging_config": { 19 "logging_config": {
20 "enabled": true, 20 "enabled": true,
......
1 """
2 任务队列系统
3 提供异步图像生成任务的队列管理和 UI 组件
4 """
5
6 from dataclasses import dataclass, field
7 from datetime import datetime
8 from enum import Enum
9 from typing import Optional, List, Dict
10 from queue import Queue
11 from threading import Lock
12 import uuid
13 import logging
14 import io
15
16 from PySide6.QtCore import QObject, Signal, QTimer, Qt
17 from PySide6.QtWidgets import (
18 QWidget, QVBoxLayout, QHBoxLayout, QLabel,
19 QPushButton, QListWidget, QListWidgetItem, QDialog, QScrollArea, QFrame
20 )
21 from PySide6.QtGui import QPixmap, QMouseEvent
22 from PIL import Image
23
24
25 class TaskType(Enum):
26 """任务类型"""
27 IMAGE_GENERATION = "image_gen"
28 STYLE_DESIGN = "style_design"
29
30
31 class TaskStatus(Enum):
32 """任务状态"""
33 PENDING = "pending"
34 RUNNING = "running"
35 COMPLETED = "completed"
36 FAILED = "failed"
37 CANCELLED = "cancelled"
38
39
40 @dataclass
41 class Task:
42 """任务数据模型"""
43 # 标识
44 id: str
45 type: TaskType
46 status: TaskStatus
47
48 # 输入参数
49 prompt: str
50 api_key: str
51 reference_images: List[str]
52 aspect_ratio: str
53 image_size: str
54 model: str
55
56 # 时间戳
57 created_at: datetime
58 started_at: Optional[datetime] = None
59 completed_at: Optional[datetime] = None
60
61 # 结果
62 result_bytes: Optional[bytes] = None
63 error_message: Optional[str] = None
64
65 # UI 相关
66 thumbnail: Optional[bytes] = None
67 progress: float = 0.0
68
69
70 class TaskQueueManager(QObject):
71 """
72 单例任务队列管理器
73 管理所有图像生成任务的生命周期
74 """
75 # Signals
76 task_added = Signal(Task)
77 task_started = Signal(str) # task_id
78 task_completed = Signal(str, bytes, str, list, str, str, str) # task_id, image_bytes, prompt, ref_images, aspect_ratio, image_size, model
79 task_failed = Signal(str, str) # task_id, error_message
80 task_progress = Signal(str, float, str) # task_id, progress, status_text
81
82 _instance = None
83 _lock = Lock()
84
85 def __new__(cls):
86 """单例模式"""
87 if cls._instance is None:
88 with cls._lock:
89 if cls._instance is None:
90 cls._instance = super().__new__(cls)
91 return cls._instance
92
93 def __init__(self):
94 if hasattr(self, '_initialized'):
95 return
96 super().__init__()
97
98 self.logger = logging.getLogger(__name__)
99 self._tasks: Dict[str, Task] = {}
100 self._queue = Queue()
101 self._current_worker = None
102 self._max_queue_size = 10
103 self._max_history_size = 10 # 只保留最近10条完成任务
104
105 self._initialized = True
106 self.logger.info("TaskQueueManager 初始化完成")
107
108 def submit_task(
109 self,
110 task_type: TaskType,
111 prompt: str,
112 api_key: str,
113 reference_images: List[str],
114 aspect_ratio: str,
115 image_size: str,
116 model: str
117 ) -> str:
118 """
119 提交新任务到队列
120
121 Args:
122 task_type: 任务类型
123 prompt: 图片描述
124 api_key: API 密钥
125 reference_images: 参考图片路径列表
126 aspect_ratio: 宽高比
127 image_size: 图片尺寸
128 model: 模型名称
129
130 Returns:
131 task_id: 任务唯一标识
132
133 Raises:
134 RuntimeError: 队列已满
135 """
136 # 检查队列容量
137 if self._queue.qsize() >= self._max_queue_size:
138 raise RuntimeError(f"任务队列已满 (最大 {self._max_queue_size} 个)")
139
140 # 创建任务
141 task = Task(
142 id=str(uuid.uuid4()),
143 type=task_type,
144 status=TaskStatus.PENDING,
145 prompt=prompt,
146 api_key=api_key,
147 reference_images=reference_images.copy() if reference_images else [],
148 aspect_ratio=aspect_ratio,
149 image_size=image_size,
150 model=model,
151 created_at=datetime.now()
152 )
153
154 self._tasks[task.id] = task
155 self._queue.put(task.id)
156
157 self.logger.info(f"任务已提交: {task.id[:8]} - {prompt[:30]}")
158 self.task_added.emit(task)
159
160 # 如果没有正在运行的任务，启动处理
161 if self._current_worker is None or not self._current_worker.isRunning():
162 self._process_next()
163
164 return task.id
165
166 def _process_next(self):
167 """处理队列中的下一个任务"""
168 if self._queue.empty():
169 self.logger.debug("队列为空，无任务处理")
170 return
171
172 task_id = self._queue.get()
173 task = self._tasks[task_id]
174 task.status = TaskStatus.RUNNING
175 task.started_at = datetime.now()
176
177 self.logger.info(f"开始处理任务: {task_id[:8]}")
178
179 # 导入 ImageGenerationWorker
180 from image_generator import ImageGenerationWorker
181
182 # 创建 worker
183 self._current_worker = ImageGenerationWorker(
184 task.api_key,
185 task.prompt,
186 task.reference_images,
187 task.aspect_ratio,
188 task.image_size,
189 task.model
190 )
191
192 # 绑定信号
193 self._current_worker.finished.connect(
194 lambda img_bytes, prompt, ref_imgs, ar, size, model:
195 self._on_task_completed(task_id, img_bytes, prompt, ref_imgs, ar, size, model)
196 )
197 self._current_worker.error.connect(
198 lambda error: self._on_task_failed(task_id, error)
199 )
200 self._current_worker.progress.connect(
201 lambda status: self.task_progress.emit(task_id, 0.5, status)
202 )
203
204 self.task_started.emit(task_id)
205 self._current_worker.start()
206
207 def _on_task_completed(self, task_id: str, image_bytes: bytes, prompt: str,
208 reference_images: list, aspect_ratio: str, image_size: str, model: str):
209 """任务完成回调"""
210 task = self._tasks.get(task_id)
211 if not task:
212 self.logger.error(f"任务 {task_id[:8]} 不存在")
213 return
214
215 task.status = TaskStatus.COMPLETED
216 task.completed_at = datetime.now()
217 task.result_bytes = image_bytes
218
219 # 生成缩略图
220 try:
221 task.thumbnail = self._create_thumbnail(image_bytes)
222 except Exception as e:
223 self.logger.warning(f"生成缩略图失败: {e}")
224
225 elapsed = (task.completed_at - task.started_at).total_seconds()
226 self.logger.info(f"任务完成: {task_id[:8]} - 耗时 {elapsed:.1f}s")
227
228 self.task_completed.emit(task_id, image_bytes, prompt, reference_images,
229 aspect_ratio, image_size, model)
230
231 # 清理旧任务历史,只保留最近的完成任务
232 self._cleanup_old_tasks()
233
234 # 处理下一个任务
235 self._process_next()
236
237 def _on_task_failed(self, task_id: str, error: str):
238 """任务失败回调"""
239 task = self._tasks.get(task_id)
240 if not task:
241 self.logger.error(f"任务 {task_id[:8]} 不存在")
242 return
243
244 task.status = TaskStatus.FAILED
245 task.completed_at = datetime.now()
246 task.error_message = error
247
248 self.logger.error(f"任务失败: {task_id[:8]} - {error}")
249
250 self.task_failed.emit(task_id, error)
251
252 # 清理旧任务历史
253 self._cleanup_old_tasks()
254
255 # 处理下一个任务
256 self._process_next()
257
258 def _cleanup_old_tasks(self):
259 """清理旧任务,只保留最近的完成/失败任务"""
260 # 获取所有已完成和失败的任务,按完成时间排序
261 finished_tasks = [
262 t for t in self._tasks.values()
263 if t.status in [TaskStatus.COMPLETED, TaskStatus.FAILED] and t.completed_at
264 ]
265 finished_tasks.sort(key=lambda t: t.completed_at, reverse=True)
266
267 # 只保留最近的 N 条
268 if len(finished_tasks) > self._max_history_size:
269 tasks_to_remove = finished_tasks[self._max_history_size:]
270 for task in tasks_to_remove:
271 del self._tasks[task.id]
272 self.logger.debug(f"清理旧任务: {task.id[:8]}")
273
274 def _create_thumbnail(self, image_bytes: bytes) -> bytes:
275 """
276 创建缩略图 (50x50)
277
278 Args:
279 image_bytes: 原始图片字节
280
281 Returns:
282 缩略图字节
283 """
284 img = Image.open(io.BytesIO(image_bytes))
285 img.thumbnail((50, 50))
286
287 thumb_io = io.BytesIO()
288 img.save(thumb_io, format='PNG')
289 return thumb_io.getvalue()
290
291 def get_task(self, task_id: str) -> Optional[Task]:
292 """获取任务详情"""
293 return self._tasks.get(task_id)
294
295 def get_all_tasks(self) -> List[Task]:
296 """获取所有任务"""
297 return list(self._tasks.values())
298
299 def get_pending_count(self) -> int:
300 """获取等待中任务数"""
301 return sum(1 for t in self._tasks.values() if t.status == TaskStatus.PENDING)
302
303 def get_running_count(self) -> int:
304 """获取运行中任务数"""
305 return sum(1 for t in self._tasks.values() if t.status == TaskStatus.RUNNING)
306
307 def cancel_task(self, task_id: str):
308 """取消任务 (仅等待中任务)"""
309 task = self._tasks.get(task_id)
310 if task and task.status == TaskStatus.PENDING:
311 task.status = TaskStatus.CANCELLED
312 self.logger.info(f"任务已取消: {task_id[:8]}")
313
314
315 class TaskQueueWidget(QWidget):
316 """
317 右侧极窄任务列表
318 按设计文档显示任务状态文字列表
319 """
320
321 def __init__(self, manager: TaskQueueManager, parent=None):
322 super().__init__(parent)
323 self.logger = logging.getLogger(__name__)
324 self.manager = manager
325 self.parent_window = parent # 用于数据回填
326
327 self._setup_ui()
328 self._connect_signals()
329 self._update_summary()
330
331 def _setup_ui(self):
332 """构建右侧任务列表 UI"""
333 layout = QVBoxLayout()
334 layout.setContentsMargins(8, 8, 8, 8)
335 layout.setSpacing(2)
336
337 # 标题
338 title = QLabel("任务队列")
339 title.setStyleSheet("QLabel { font-weight: bold; font-size: 10px; color: #666; }")
340 title.setAlignment(Qt.AlignCenter)
341 layout.addWidget(title)
342
343 # 分隔线
344 line = QLabel()
345 line.setFrameStyle(QFrame.HLine | QFrame.Sunken)
346 layout.addWidget(line)
347
348 # 任务状态列表 - 可点击的列表项
349 self.task_list = QListWidget()
350 self.task_list.setStyleSheet("""
351 QListWidget {
352 border: none;
353 font-size: 11px;
354 padding: 2px;
355 }
356 QListWidget::item {
357 padding: 4px 2px;
358 border-bottom: 1px solid #eee;
359 min-height: 20px;
360 }
361 QListWidget::item:hover {
362 background-color: #e3f2fd;
363 cursor: pointer;
364 }
365 QListWidget::item:selected {
366 background-color: #bbdefb;
367 }
368 """)
369 self.task_list.itemClicked.connect(self._on_task_item_clicked)
370 layout.addWidget(self.task_list)
371
372 layout.addStretch()
373 self.setLayout(layout)
374
375 # 设置极窄宽度
376 self.setMaximumWidth(120)
377 self.setMinimumWidth(80)
378
379 # 样式
380 self.setStyleSheet("""
381 TaskQueueWidget {
382 background-color: #f5f5f5;
383 border-left: 1px solid #ddd;
384 }
385 """)
386
387 def _connect_signals(self):
388 """绑定信号"""
389 self.manager.task_added.connect(self._on_task_added)
390 self.manager.task_started.connect(self._on_task_started)
391 self.manager.task_completed.connect(self._on_task_completed)
392 self.manager.task_failed.connect(self._on_task_failed)
393 self.manager.task_progress.connect(self._on_task_progress)
394
395
396 def _update_summary(self):
397 """更新任务列表 - 显示可点击的状态项"""
398 self.task_list.clear()
399 tasks = self.manager.get_all_tasks()
400
401 # 按状态分类并转换为文字
402 for task in tasks:
403 # 状态文字和颜色
404 if task.status == TaskStatus.RUNNING:
405 status_text = "执行中"
406 color = "#FF9500" # 橙色
407 elif task.status == TaskStatus.PENDING:
408 status_text = "等待中"
409 color = "#007AFF" # 蓝色
410 elif task.status == TaskStatus.COMPLETED:
411 status_text = "已完成"
412 color = "#34C759" # 绿色
413 elif task.status == TaskStatus.FAILED:
414 status_text = "失败"
415 color = "#FF3B30" # 红色
416 else:
417 status_text = "未知"
418 color = "#666666" # 灰色
419
420 # 创建列表项
421 item = QListWidgetItem(status_text)
422 item.setData(Qt.UserRole, task.id) # 存储任务ID
423
424 # 设置颜色
425 if hasattr(item, 'setForeground'):
426 from PySide6.QtGui import QBrush, QColor
427 item.setForeground(QBrush(QColor(color)))
428
429 # 添加到列表
430 self.task_list.addItem(item)
431
432 # 最多显示最近10个任务
433 if self.task_list.count() > 10:
434 for i in range(self.task_list.count() - 10):
435 self.task_list.takeItem(0)
436
437
438 def _on_task_added(self, task: Task):
439 """任务添加回调"""
440 self._update_summary()
441
442 def _on_task_item_clicked(self, item: QListWidgetItem):
443 """点击任务项 - 回填数据到主窗口"""
444 task_id = item.data(Qt.UserRole)
445 if not task_id:
446 return
447
448 task = self.manager.get_task(task_id)
449 if not task or not self.parent_window:
450 return
451
452 # 切换到对应的标签页
453 if task.type == TaskType.STYLE_DESIGN:
454 self.parent_window.tab_widget.setCurrentIndex(1) # 款式设计标签
455 style_tab = self.parent_window.tab_widget.currentWidget()
456 else:
457 self.parent_window.tab_widget.setCurrentIndex(0) # 图片生成标签
458 gen_tab = self.parent_window.tab_widget.currentWidget()
459
460 # 如果是已完成任务，直接在主窗口显示结果
461 if task.status == TaskStatus.COMPLETED and task.result_bytes:
462 if task.type == TaskType.STYLE_DESIGN and hasattr(style_tab, '_display_generated_image_from_bytes'):
463 # 款式设计：将图片数据存储到样式标签并显示
464 style_tab.generated_image_bytes = task.result_bytes
465 style_tab._display_generated_image_from_bytes()
466 elif hasattr(gen_tab, '_display_generated_image_from_bytes'):
467 # 图片生成：将图片数据存储到主窗口并显示
468 self.parent_window.generated_image_bytes = task.result_bytes
469 gen_tab._display_generated_image_from_bytes()
470
471 # 回填参数到主窗口
472 self._load_task_to_main_window(task)
473
474 def _load_task_to_main_window(self, task: Task):
475 """将任务数据回填到主窗口"""
476 try:
477 if task.type == TaskType.STYLE_DESIGN:
478 # 款式设计标签页 - 回填prompt到预览框
479 self.parent_window.tab_widget.setCurrentIndex(1) # 款式设计标签
480 style_tab = self.parent_window.tab_widget.currentWidget()
481
482 # 回填prompt到预览框
483 if hasattr(style_tab, 'prompt_preview'):
484 style_tab.prompt_preview.setPlainText(task.prompt)
485
486 # 回填设置
487 if hasattr(style_tab, 'aspect_ratio') and task.aspect_ratio:
488 index = style_tab.aspect_ratio.findText(task.aspect_ratio)
489 if index >= 0:
490 style_tab.aspect_ratio.setCurrentIndex(index)
491 if hasattr(style_tab, 'image_size') and task.image_size:
492 index = style_tab.image_size.findText(task.image_size)
493 if index >= 0:
494 style_tab.image_size.setCurrentIndex(index)
495
496 else:
497 # 图片生成标签页
498 self.parent_window.tab_widget.setCurrentIndex(0) # 图片生成标签
499 gen_tab = self.parent_window.tab_widget.currentWidget()
500
501 # 回填prompt
502 if hasattr(gen_tab, 'prompt_text'):
503 gen_tab.prompt_text.setPlainText(task.prompt)
504
505 # 回填参考图片
506 if task.reference_images and hasattr(gen_tab, 'add_reference_image'):
507 for ref_path in task.reference_images:
508 gen_tab.add_reference_image(ref_path)
509
510 # 回填设置
511 if hasattr(gen_tab, 'aspect_ratio') and task.aspect_ratio:
512 index = gen_tab.aspect_ratio.findText(task.aspect_ratio)
513 if index >= 0:
514 gen_tab.aspect_ratio.setCurrentIndex(index)
515 if hasattr(gen_tab, 'image_size') and task.image_size:
516 index = gen_tab.image_size.findText(task.image_size)
517 if index >= 0:
518 gen_tab.image_size.setCurrentIndex(index)
519
520 except Exception as e:
521 self.logger.error(f"回填数据到主窗口失败: {e}")
522
523 def _on_task_started(self, task_id: str):
524 """任务开始回调"""
525 self._update_summary()
526
527 def _on_task_completed(self, task_id: str, *args):
528 """任务完成回调"""
529 self._update_summary()
530
531 def _on_task_failed(self, task_id: str, error: str):
532 """任务失败回调"""
533 self._update_summary()
534
535 def _show_task_detail_dialog(self):
536 """显示任务详情弹窗"""
537 if self.detail_dialog and self.detail_dialog.isVisible():
538 self.detail_dialog.raise_()
539 return
540
541 dialog = QDialog(self)
542 dialog.setWindowTitle("任务详情")
543 dialog.resize(400, 500)
544
545 layout = QVBoxLayout()
546
547 # 标题
548 title = QLabel("📋 任务队列")
549 title.setStyleSheet("QLabel { font-weight: bold; font-size: 14px; }")
550 layout.addWidget(title)
551
552 # 任务列表
553 task_list = QListWidget()
554 task_list.itemClicked.connect(lambda item: self._on_detail_task_clicked(item, dialog))
555
556 # 添加任务项
557 for task in self.manager.get_all_tasks():
558 item = QListWidgetItem()
559 item.setData(Qt.UserRole, task.id)
560
561 status_map = {
562 TaskStatus.RUNNING: ("●", "#FF9500"),
563 TaskStatus.PENDING: ("○", "#8E8E93"),
564 TaskStatus.COMPLETED: ("✓", "#34C759"),
565 TaskStatus.FAILED: ("✗", "#FF3B30"),
566 }
567 icon, color = status_map.get(task.status, ("?", "#000"))
568
569 prompt_preview = task.prompt[:30] + "..." if len(task.prompt) > 30 else task.prompt
570 display_text = f"{icon} {prompt_preview}"
571
572 item.setText(display_text)
573 if hasattr(item, 'setForeground'):
574 from PySide6.QtGui import QBrush, QColor
575 item.setForeground(QBrush(QColor(color)))
576
577 task_list.addItem(item)
578
579 layout.addWidget(task_list)
580
581 # 关闭按钮
582 close_btn = QPushButton("关闭")
583 close_btn.clicked.connect(dialog.close)
584 layout.addWidget(close_btn)
585
586 dialog.setLayout(layout)
587 self.detail_dialog = dialog
588 dialog.exec()
589
590 def _on_detail_task_clicked(self, item: QListWidgetItem, dialog: QDialog):
591 """详情弹窗中点击任务项"""
592 dialog.close()
593 self._on_task_item_clicked(item)
594
595 def _on_task_progress(self, task_id: str, progress: float, status_text: str):
596 """任务进度回调"""
597 # 侧边栏模式下进度信息通过状态标签显示
598 # 不需要额外更新，由_update_summary统一处理
599 pass
600
601
602 def _on_task_item_clicked(self, item: QListWidgetItem):
603 """单击任务 - 回填数据到主窗口或显示结果"""
604 task_id = item.data(Qt.UserRole)
605 task = self.manager.get_task(task_id)
606
607 if not task or not self.parent_window:
608 return
609
610 # 切换到对应的标签页
611 if task.type == TaskType.STYLE_DESIGN:
612 self.parent_window.tab_widget.setCurrentIndex(1) # 款式设计标签
613 style_tab = self.parent_window.tab_widget.currentWidget()
614 else:
615 self.parent_window.tab_widget.setCurrentIndex(0) # 图片生成标签
616 gen_tab = self.parent_window.tab_widget.currentWidget()
617
618 # 如果是已完成任务，直接在主窗口显示结果
619 if task.status == TaskStatus.COMPLETED and task.result_bytes:
620 self.parent_window.generated_image_bytes = task.result_bytes
621 # 显示生成的图片
622 if hasattr(self.parent_window, 'display_generated_image'):
623 self.parent_window.display_generated_image()
624 elif task.type == TaskType.STYLE_DESIGN and hasattr(style_tab, '_display_generated_image_from_bytes'):
625 style_tab._display_generated_image_from_bytes()
626 elif hasattr(gen_tab, '_display_generated_image_from_bytes'):
627 gen_tab._display_generated_image_from_bytes()
628
629 # 回填参数
630 self._load_task_to_main_window(task)
631
632 def _load_task_to_main_window(self, task: Task):
633 """将任务数据回填到主窗口"""
634 try:
635 if task.type == TaskType.STYLE_DESIGN:
636 # 切换到款式设计标签
637 self.parent_window.tab_widget.setCurrentIndex(1)
638 style_tab = self.parent_window.tab_widget.currentWidget()
639 if hasattr(style_tab, 'library_manager'):
640 # 款式设计不需要回填prompt，因为是参数组合
641 pass
642 else:
643 # 切换到图片生成标签
644 self.parent_window.tab_widget.setCurrentIndex(0)
645 gen_tab = self.parent_window.tab_widget.currentWidget()
646 if hasattr(gen_tab, 'prompt_input'):
647 # 回填prompt
648 gen_tab.prompt_input.setPlainText(task.prompt)
649 # 回填参考图片
650 if task.reference_images:
651 for ref_path in task.reference_images:
652 if hasattr(gen_tab, 'add_reference_image'):
653 gen_tab.add_reference_image(ref_path)
654 # 回填设置
655 if hasattr(gen_tab, 'aspect_ratio') and task.aspect_ratio:
656 index = gen_tab.aspect_ratio.findText(task.aspect_ratio)
657 if index >= 0:
658 gen_tab.aspect_ratio.setCurrentIndex(index)
659 if hasattr(gen_tab, 'image_size') and task.image_size:
660 index = gen_tab.image_size.findText(task.image_size)
661 if index >= 0:
662 gen_tab.image_size.setCurrentIndex(index)
663
664 except Exception as e:
665 self.logger.error(f"回填数据到主窗口失败: {e}")
1 """
2 任务队列系统
3 提供异步图像生成任务的队列管理和 UI 组件
4 """
5
6 from dataclasses import dataclass, field
7 from datetime import datetime
8 from enum import Enum
9 from typing import Optional, List, Dict
10 from queue import Queue
11 from threading import Lock
12 import uuid
13 import logging
14 import io
15
16 from PySide6.QtCore import QObject, Signal, QTimer, Qt
17 from PySide6.QtWidgets import (
18 QWidget, QVBoxLayout, QHBoxLayout, QLabel,
19 QPushButton, QListWidget, QListWidgetItem, QDialog, QScrollArea, QFrame
20 )
21 from PySide6.QtGui import QPixmap, QMouseEvent
22 from PIL import Image
23
24
25 class TaskType(Enum):
26 """任务类型"""
27 IMAGE_GENERATION = "image_gen"
28 STYLE_DESIGN = "style_design"
29
30
31 class TaskStatus(Enum):
32 """任务状态"""
33 PENDING = "pending"
34 RUNNING = "running"
35 COMPLETED = "completed"
36 FAILED = "failed"
37 CANCELLED = "cancelled"
38
39
40 @dataclass
41 class Task:
42 """任务数据模型"""
43 # 标识
44 id: str
45 type: TaskType
46 status: TaskStatus
47
48 # 输入参数
49 prompt: str
50 api_key: str
51 reference_images: List[str]
52 aspect_ratio: str
53 image_size: str
54 model: str
55
56 # 时间戳
57 created_at: datetime
58 started_at: Optional[datetime] = None
59 completed_at: Optional[datetime] = None
60
61 # 结果
62 result_bytes: Optional[bytes] = None
63 error_message: Optional[str] = None
64
65 # UI 相关
66 thumbnail: Optional[bytes] = None
67 progress: float = 0.0
68
69
70 class TaskQueueManager(QObject):
71 """
72 单例任务队列管理器
73 管理所有图像生成任务的生命周期
74 """
75 # Signals
76 task_added = Signal(Task)
77 task_started = Signal(str) # task_id
78 task_completed = Signal(str, bytes, str, list, str, str, str) # task_id, image_bytes, prompt, ref_images, aspect_ratio, image_size, model
79 task_failed = Signal(str, str) # task_id, error_message
80 task_progress = Signal(str, float, str) # task_id, progress, status_text
81
82 _instance = None
83 _lock = Lock()
84
85 def __new__(cls):
86 """单例模式"""
87 if cls._instance is None:
88 with cls._lock:
89 if cls._instance is None:
90 cls._instance = super().__new__(cls)
91 return cls._instance
92
93 def __init__(self):
94 if hasattr(self, '_initialized'):
95 return
96 super().__init__()
97
98 self.logger = logging.getLogger(__name__)
99 self._tasks: Dict[str, Task] = {}
100 self._queue = Queue()
101 self._current_worker = None
102 self._max_queue_size = 10
103 self._max_history_size = 10 # 只保留最近10条完成任务
104
105 self._initialized = True
106 self.logger.info("TaskQueueManager 初始化完成")
107
108 def submit_task(
109 self,
110 task_type: TaskType,
111 prompt: str,
112 api_key: str,
113 reference_images: List[str],
114 aspect_ratio: str,
115 image_size: str,
116 model: str
117 ) -> str:
118 """
119 提交新任务到队列
120
121 Args:
122 task_type: 任务类型
123 prompt: 图片描述
124 api_key: API 密钥
125 reference_images: 参考图片路径列表
126 aspect_ratio: 宽高比
127 image_size: 图片尺寸
128 model: 模型名称
129
130 Returns:
131 task_id: 任务唯一标识
132
133 Raises:
134 RuntimeError: 队列已满
135 """
136 # 检查队列容量
137 if self._queue.qsize() >= self._max_queue_size:
138 raise RuntimeError(f"任务队列已满 (最大 {self._max_queue_size} 个)")
139
140 # 创建任务
141 task = Task(
142 id=str(uuid.uuid4()),
143 type=task_type,
144 status=TaskStatus.PENDING,
145 prompt=prompt,
146 api_key=api_key,
147 reference_images=reference_images.copy() if reference_images else [],
148 aspect_ratio=aspect_ratio,
149 image_size=image_size,
150 model=model,
151 created_at=datetime.now()
152 )
153
154 self._tasks[task.id] = task
155 self._queue.put(task.id)
156
157 self.logger.info(f"任务已提交: {task.id[:8]} - {prompt[:30]}")
158 self.task_added.emit(task)
159
160 # 如果没有正在运行的任务，启动处理
161 if self._current_worker is None or not self._current_worker.isRunning():
162 self._process_next()
163
164 return task.id
165
166 def _process_next(self):
167 """处理队列中的下一个任务"""
168 if self._queue.empty():
169 self.logger.debug("队列为空，无任务处理")
170 return
171
172 task_id = self._queue.get()
173 task = self._tasks[task_id]
174 task.status = TaskStatus.RUNNING
175 task.started_at = datetime.now()
176
177 self.logger.info(f"开始处理任务: {task_id[:8]}")
178
179 # 导入 ImageGenerationWorker
180 from image_generator import ImageGenerationWorker
181
182 # 创建 worker
183 self._current_worker = ImageGenerationWorker(
184 task.api_key,
185 task.prompt,
186 task.reference_images,
187 task.aspect_ratio,
188 task.image_size,
189 task.model
190 )
191
192 # 绑定信号
193 self._current_worker.finished.connect(
194 lambda img_bytes, prompt, ref_imgs, ar, size, model:
195 self._on_task_completed(task_id, img_bytes, prompt, ref_imgs, ar, size, model)
196 )
197 self._current_worker.error.connect(
198 lambda error: self._on_task_failed(task_id, error)
199 )
200 self._current_worker.progress.connect(
201 lambda status: self.task_progress.emit(task_id, 0.5, status)
202 )
203
204 self.task_started.emit(task_id)
205 self._current_worker.start()
206
207 def _on_task_completed(self, task_id: str, image_bytes: bytes, prompt: str,
208 reference_images: list, aspect_ratio: str, image_size: str, model: str):
209 """任务完成回调"""
210 task = self._tasks.get(task_id)
211 if not task:
212 self.logger.error(f"任务 {task_id[:8]} 不存在")
213 return
214
215 task.status = TaskStatus.COMPLETED
216 task.completed_at = datetime.now()
217 task.result_bytes = image_bytes
218
219 # 生成缩略图
220 try:
221 task.thumbnail = self._create_thumbnail(image_bytes)
222 except Exception as e:
223 self.logger.warning(f"生成缩略图失败: {e}")
224
225 elapsed = (task.completed_at - task.started_at).total_seconds()
226 self.logger.info(f"任务完成: {task_id[:8]} - 耗时 {elapsed:.1f}s")
227
228 self.task_completed.emit(task_id, image_bytes, prompt, reference_images,
229 aspect_ratio, image_size, model)
230
231 # 清理旧任务历史,只保留最近的完成任务
232 self._cleanup_old_tasks()
233
234 # 处理下一个任务
235 self._process_next()
236
237 def _on_task_failed(self, task_id: str, error: str):
238 """任务失败回调"""
239 task = self._tasks.get(task_id)
240 if not task:
241 self.logger.error(f"任务 {task_id[:8]} 不存在")
242 return
243
244 task.status = TaskStatus.FAILED
245 task.completed_at = datetime.now()
246 task.error_message = error
247
248 self.logger.error(f"任务失败: {task_id[:8]} - {error}")
249
250 self.task_failed.emit(task_id, error)
251
252 # 清理旧任务历史
253 self._cleanup_old_tasks()
254
255 # 处理下一个任务
256 self._process_next()
257
258 def _cleanup_old_tasks(self):
259 """清理旧任务,只保留最近的完成/失败任务"""
260 # 获取所有已完成和失败的任务,按完成时间排序
261 finished_tasks = [
262 t for t in self._tasks.values()
263 if t.status in [TaskStatus.COMPLETED, TaskStatus.FAILED] and t.completed_at
264 ]
265 finished_tasks.sort(key=lambda t: t.completed_at, reverse=True)
266
267 # 只保留最近的 N 条
268 if len(finished_tasks) > self._max_history_size:
269 tasks_to_remove = finished_tasks[self._max_history_size:]
270 for task in tasks_to_remove:
271 del self._tasks[task.id]
272 self.logger.debug(f"清理旧任务: {task.id[:8]}")
273
274 def _create_thumbnail(self, image_bytes: bytes) -> bytes:
275 """
276 创建缩略图 (50x50)
277
278 Args:
279 image_bytes: 原始图片字节
280
281 Returns:
282 缩略图字节
283 """
284 img = Image.open(io.BytesIO(image_bytes))
285 img.thumbnail((50, 50))
286
287 thumb_io = io.BytesIO()
288 img.save(thumb_io, format='PNG')
289 return thumb_io.getvalue()
290
291 def get_task(self, task_id: str) -> Optional[Task]:
292 """获取任务详情"""
293 return self._tasks.get(task_id)
294
295 def get_all_tasks(self) -> List[Task]:
296 """获取所有任务"""
297 return list(self._tasks.values())
298
299 def get_pending_count(self) -> int:
300 """获取等待中任务数"""
301 return sum(1 for t in self._tasks.values() if t.status == TaskStatus.PENDING)
302
303 def get_running_count(self) -> int:
304 """获取运行中任务数"""
305 return sum(1 for t in self._tasks.values() if t.status == TaskStatus.RUNNING)
306
307 def cancel_task(self, task_id: str):
308 """取消任务 (仅等待中任务)"""
309 task = self._tasks.get(task_id)
310 if task and task.status == TaskStatus.PENDING:
311 task.status = TaskStatus.CANCELLED
312 self.logger.info(f"任务已取消: {task_id[:8]}")
313
314