12d6934e by 柴进

refactor: 业务核心从 image_generator.py 拆到 core/

5082 行的怪兽塞了 12 个类(数据库/历史/Worker/词库/UI),现在按职责拆开:

  core/paths.py       get_app_data_path + save_png_with_validation + 启动迁移
  core/database.py    DatabaseManager + hash_password
  core/history.py     HistoryItem + HistoryListModel + HistoryManager
  core/generation.py  ImageGenerationWorker + Gemini 模型常量
  core/jewelry.py     DEFAULT_JEWELRY_LIBRARY + JewelryLibraryManager + PromptAssembler

image_generator.py 5082 → 3781 行,剩下全是 QWidget UI 类
(LoginDialog / DraggableThumbnail / DragDropScrollArea / ImageGeneratorWindow /
StyleDesignerTab + utils + main),task #19 QML 全量切换后整体删除。

外部消费者改 import:
  task_queue.py / temp_clean.py: from image_generator → from core.generation
  image_generator.py 顶部 from core.* 引入,LoginDialog 等内部代码无感知

冒烟测试:image_generator/task_queue/core 各自 import 通过,类身份正确。

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 1fd96b05
1 """业务核心模块。
2
3 从 image_generator.py 拆出,保留 QWidget UI 类在原文件供 task #19 整删。
4 QML 桥层和外部模块直接 import core.* 而不是 image_generator。
5 """
1 """数据库连接 + 用户认证。
2
3 DatabaseManager 只负责 MySQL 连接和 user table 的 SHA256 密码核验。
4 db_config 由 config_util 加载并传入;本模块不读 config 文件本身。
5 """
6 import hashlib
7 import logging
8
9 import pymysql
10
11
12 def hash_password(password: str) -> str:
13 """使用 SHA256 哈希密码"""
14 return hashlib.sha256(password.encode('utf-8')).hexdigest()
15
16
17 class DatabaseManager:
18 """数据库连接管理类"""
19
20 def __init__(self, db_config):
21 self.config = db_config
22 self.logger = logging.getLogger(__name__)
23
24 def authenticate(self, username, password):
25 """
26 验证用户凭证
27 返回: (success: bool, message: str)
28 """
29 try:
30 self.logger.info(f"开始用户认证: {username}")
31
32 password_hash = hash_password(password)
33
34 self.logger.debug(f"连接数据库: {self.config['host']}:{self.config.get('port', 3306)}")
35 conn = pymysql.connect(
36 host=self.config['host'],
37 port=self.config.get('port', 3306),
38 user=self.config['user'],
39 password=self.config['password'],
40 database=self.config['database'],
41 connect_timeout=5
42 )
43
44 try:
45 with conn.cursor() as cursor:
46 sql = f"SELECT * FROM {self.config['table']} WHERE user_name=%s AND passwd=%s AND status='active'"
47 cursor.execute(sql, (username, password_hash))
48 result = cursor.fetchone()
49
50 if result:
51 self.logger.info(f"用户认证成功: {username}")
52 return True, "认证成功"
53 else:
54 self.logger.warning(f"用户认证失败: {username} - 用户名或密码错误")
55 return False, "用户名或密码错误"
56 finally:
57 conn.close()
58
59 except pymysql.OperationalError as e:
60 error_msg = "无法连接到服务器,请检查网络连接"
61 self.logger.error(f"数据库连接失败: {e}")
62 return False, error_msg
63 except Exception as e:
64 error_msg = f"认证失败: {str(e)}"
65 self.logger.error(f"认证过程异常: {e}")
66 return False, error_msg
1 """图像生成 Worker + Gemini 模型常量。
2
3 ImageGenerationWorker 是 QThread,由 TaskQueueManager 拉起执行单条生成任务。
4 任务参数(prompt / 参考图 / aspect_ratio / image_size / model)从队列传入,
5 完成后通过 finished/error/progress 信号回报。
6 """
7 import base64
8 import logging
9 import os
10 from typing import Optional
11
12 from PySide6.QtCore import QThread, Signal
13 from google import genai
14 from google.genai import types
15
16
17 # 生成模式 -> Gemini 模型 ID 映射(单一真相源,消除原先两处 get_selected_model 复制粘贴)
18 # 极速模式:Nano Banana 2 (Gemini 3.1 Flash Image), 指令遵循强于 2.5-flash-image
19 # 慢速模式:Nano Banana Pro (Gemini 3 Pro Image Preview)
20 MODEL_BY_MODE = {
21 "极速模式": "gemini-3.1-flash-image-preview",
22 "慢速模式": "gemini-3-pro-image-preview",
23 }
24 MODEL_PRO = MODEL_BY_MODE["慢速模式"] # 用于 Worker 中判断是否支持 image_size 参数
25
26 # Nano Banana 2 (Flash) 独占的宽高比 —— Pro 不支持,选中这些时需提示切换到极速模式
27 FLASH_ONLY_ASPECT_RATIOS = {"1:4", "4:1", "1:8", "8:1"}
28
29
30 class ImageGenerationWorker(QThread):
31 """Worker thread for image generation"""
32 finished = Signal(bytes, str, list, str, str,
33 str) # image_bytes, prompt, reference_images, aspect_ratio, image_size, model
34 error = Signal(str)
35 progress = Signal(str)
36
37 def __init__(self, api_key, prompt, images, aspect_ratio, image_size, model=MODEL_PRO):
38 super().__init__()
39 self.logger = logging.getLogger(__name__)
40 self.api_key = api_key
41 self.prompt = prompt
42 self.images = images
43 self.aspect_ratio = aspect_ratio
44 self.image_size = image_size
45 self.model = model
46
47 # 审计元信息:供 TaskQueueManager 在信号回调中读取
48 self.finish_reason: Optional[str] = None
49
50 self.logger.info(f"图片生成任务初始化 - 模型: {model}, 尺寸: {image_size}, 宽高比: {aspect_ratio}")
51
52 def _extract_finish_reason(self, response) -> Optional[str]:
53 """从 Gemini 响应提取 finish_reason,失败返回 None(不抛异常)。"""
54 try:
55 fr = response.candidates[0].finish_reason
56 if fr is None:
57 return None
58 name = getattr(fr, "name", None)
59 return name if name else str(fr)
60 except Exception:
61 return None
62
63 def run(self):
64 """Execute image generation in background thread"""
65 try:
66 self.logger.info("开始图片生成任务")
67
68 if not self.prompt:
69 self.logger.error("图片描述为空")
70 self.error.emit("请输入图片描述!")
71 return
72
73 if not self.api_key:
74 self.logger.error("API密钥为空")
75 self.error.emit("未找到API密钥,请在config.json中配置!")
76 return
77
78 self.progress.emit("正在连接 Gemini API...")
79 self.logger.debug("正在连接 Gemini API")
80
81 client = genai.Client(api_key=self.api_key)
82
83 content_parts = [self.prompt]
84
85 for img_path in self.images:
86 with open(img_path, 'rb') as f:
87 img_data = f.read()
88
89 mime_type = "image/png"
90 if img_path.lower().endswith(('.jpg', '.jpeg')):
91 mime_type = "image/jpeg"
92
93 content_parts.append(
94 types.Part.from_bytes(
95 data=img_data,
96 mime_type=mime_type
97 )
98 )
99
100 self.progress.emit("正在生成图片...")
101
102 # 当前使用的两个模型都支持 aspect_ratio + image_size:
103 # - gemini-3.1-flash-image-preview (Nano Banana 2): 512/1K/2K/4K + 14 种 ratio
104 # - gemini-3-pro-image-preview (Nano Banana Pro): 1K/2K/4K
105 config = types.GenerateContentConfig(
106 response_modalities=["TEXT", "IMAGE"],
107 image_config=types.ImageConfig(
108 aspect_ratio=self.aspect_ratio,
109 image_size=self.image_size
110 )
111 )
112
113 response = client.models.generate_content(
114 model=self.model,
115 contents=content_parts,
116 config=config
117 )
118 self.finish_reason = self._extract_finish_reason(response)
119
120 text_fragments = []
121 parts = response.parts or []
122 for part in parts:
123 if hasattr(part, 'inline_data') and part.inline_data:
124 if isinstance(part.inline_data.data, bytes):
125 image_bytes = part.inline_data.data
126 else:
127 image_bytes = base64.b64decode(part.inline_data.data)
128
129 reference_images_bytes = []
130 for img_path in self.images:
131 if img_path and os.path.exists(img_path):
132 with open(img_path, 'rb') as f:
133 reference_images_bytes.append(f.read())
134 else:
135 reference_images_bytes.append(b'')
136
137 self.logger.info(
138 f"图片生成成功 - 模型: {self.model}, 尺寸: {self.image_size}, "
139 f"finish_reason={self.finish_reason}"
140 )
141 self.finished.emit(image_bytes, self.prompt, reference_images_bytes,
142 self.aspect_ratio, self.image_size, self.model)
143 return
144 if getattr(part, 'text', None):
145 text_fragments.append(part.text)
146
147 detail = " | ".join(t for t in text_fragments if t).strip()
148 error_msg = f"响应中没有图片数据 (finish_reason={self.finish_reason})"
149 if detail:
150 error_msg += f"\n模型说明: {detail}"
151 self.logger.error(error_msg)
152 self.error.emit(error_msg)
153
154 except Exception as e:
155 error_msg = f"图片生成异常: {e}"
156 self.logger.error(error_msg, exc_info=True)
157 self.error.emit(error_msg)
1 """应用数据路径 + 图片格式校验工具。
2
3 跨平台数据目录探测(macOS .app 包外存储 / Windows APPDATA / 开发环境同目录),
4 PNG/JPEG 格式回正(Pillow 重写防止伪装 MIME)。
5 """
6 import io
7 import logging
8 import os
9 import platform
10 import shutil
11 import sys
12 from pathlib import Path
13
14
15 def _migrate_data_from_app_bundle(target_path: Path):
16 """将 .app 内部的旧数据迁移到外部目录(仅 macOS 打包环境)"""
17 if not (getattr(sys, 'frozen', False) and platform.system() == "Darwin"):
18 return
19
20 old_path = Path(sys.executable).parent / "images"
21 if not old_path.exists() or old_path == target_path:
22 return
23
24 old_files = list(old_path.rglob("*"))
25 if not old_files:
26 return
27
28 try:
29 target_path.mkdir(parents=True, exist_ok=True)
30 migrated = 0
31 for src_file in old_path.rglob("*"):
32 if src_file.is_file():
33 rel_path = src_file.relative_to(old_path)
34 dst_file = target_path / rel_path
35 if not dst_file.exists():
36 dst_file.parent.mkdir(parents=True, exist_ok=True)
37 shutil.copy2(str(src_file), str(dst_file))
38 migrated += 1
39 print(f"已从 .app 内部迁移 {migrated} 个文件到: {target_path}")
40 except Exception as e:
41 print(f"数据迁移失败(不影响使用): {e}")
42
43
44 def get_app_data_path() -> Path:
45 """获取应用数据存储路径 - 智能选择"""
46
47 def get_candidate_paths():
48 system = platform.system()
49 candidates = []
50
51 if getattr(sys, 'frozen', False) and system == "Darwin":
52 candidates.append(Path.home() / "Library/Application Support/ZB100ImageGenerator/images")
53 elif getattr(sys, 'frozen', False):
54 candidates.append(Path(sys.executable).parent / "images")
55 else:
56 # 开发环境:保持和老路径一致 —— 项目根目录下的 images/
57 # __file__ 在 core/,需要往上一级
58 candidates.append(Path(__file__).resolve().parent.parent / "images")
59
60 if system == "Darwin":
61 candidates.append(Path.home() / "Library/Application Support/ZB100ImageGenerator/images")
62 candidates.append(Path.home() / "Documents/ZB100ImageGenerator/images")
63 elif system == "Windows":
64 candidates.append(Path(os.environ.get("APPDATA", "")) / "ZB100ImageGenerator/images")
65 candidates.append(Path.home() / "Documents/ZB100ImageGenerator/images")
66 else:
67 candidates.append(Path.home() / ".config/zb100imagegenerator/images")
68 candidates.append(Path.home() / "Documents/ZB100ImageGenerator/images")
69
70 return candidates
71
72 def test_path_write_access(path: Path) -> bool:
73 try:
74 path.mkdir(parents=True, exist_ok=True)
75 test_file = path / ".write_test"
76 test_file.write_text("test")
77 test_file.unlink()
78 return True
79 except (PermissionError, OSError) as e:
80 print(f"路径 {path} 无写入权限: {e}")
81 return False
82 except Exception as e:
83 print(f"路径 {path} 测试失败: {e}")
84 return False
85
86 candidates = get_candidate_paths()
87
88 for path in candidates:
89 if test_path_write_access(path):
90 _migrate_data_from_app_bundle(path)
91 print(f"使用图片存储路径: {path}")
92 return path
93
94 fallback_path = get_candidate_paths()[0]
95 try:
96 fallback_path.mkdir(parents=True, exist_ok=True)
97 print(f"使用备选路径: {fallback_path}")
98 return fallback_path
99 except Exception as e:
100 print(f"警告: 无法创建存储路径,将在当前目录操作: {e}")
101 return Path.cwd() / "images"
102
103
104 def save_png_with_validation(file_path: str, image_bytes: bytes) -> bool:
105 """使用 Pillow 验证并重写 PNG/JPEG,确保 MIME 与扩展名一致。
106
107 返回 True 表示 Pillow 处理成功;False 表示 Pillow 不可用或处理失败,
108 调用方应回退到原始 write_bytes。
109 """
110 try:
111 from PIL import Image
112
113 with Image.open(io.BytesIO(image_bytes)) as img:
114 file_format = img.format
115 if file_format == 'JPEG':
116 logger = logging.getLogger(__name__)
117 logger.info(f"检测到伪装PNG的JPEG文件,实际格式: {file_format}")
118
119 save_format = 'PNG' if file_path.lower().endswith('.png') else 'JPEG'
120
121 if file_format and file_format != save_format:
122 logger = logging.getLogger(__name__)
123 logger.info(f"执行格式转换: {file_format} -> {save_format}")
124
125 if save_format == 'PNG':
126 if img.mode not in ['RGBA', 'RGB', 'L']:
127 if img.mode == 'P':
128 img = img.convert('RGBA')
129 elif img.mode == 'LA':
130 img = img.convert('RGBA')
131 else:
132 img = img.convert('RGBA')
133 elif save_format == 'JPEG':
134 if img.mode in ['RGBA', 'P']:
135 img = img.convert('RGB')
136 elif img.mode == 'L':
137 img = img.convert('RGB')
138
139 img.save(file_path, save_format, optimize=True)
140
141 logger = logging.getLogger(__name__)
142 logger.info(f"图片格式验证成功: {file_path}, 保存格式: {save_format}")
143 return True
144
145 except ImportError:
146 logger = logging.getLogger(__name__)
147 logger.warning("Pillow库不可用,使用原始保存方法")
148 return False
149
150 except Exception as e:
151 logger = logging.getLogger(__name__)
152 logger.warning(f"Pillow处理失败,使用原始保存方法: {e}")
153 return False
...@@ -192,7 +192,7 @@ class TaskQueueManager(QObject): ...@@ -192,7 +192,7 @@ class TaskQueueManager(QObject):
192 self.logger.info(f"开始处理任务: {task_id[:8]}") 192 self.logger.info(f"开始处理任务: {task_id[:8]}")
193 193
194 # 导入 ImageGenerationWorker 194 # 导入 ImageGenerationWorker
195 from image_generator import ImageGenerationWorker 195 from core.generation import ImageGenerationWorker
196 196
197 # 创建 worker 197 # 创建 worker
198 self._current_worker = ImageGenerationWorker( 198 self._current_worker = ImageGenerationWorker(
......
...@@ -177,7 +177,7 @@ class TaskQueueManager(QObject): ...@@ -177,7 +177,7 @@ class TaskQueueManager(QObject):
177 self.logger.info(f"开始处理任务: {task_id[:8]}") 177 self.logger.info(f"开始处理任务: {task_id[:8]}")
178 178
179 # 导入 ImageGenerationWorker 179 # 导入 ImageGenerationWorker
180 from image_generator import ImageGenerationWorker 180 from core.generation import ImageGenerationWorker
181 181
182 # 创建 worker 182 # 创建 worker
183 self._current_worker = ImageGenerationWorker( 183 self._current_worker = ImageGenerationWorker(
......