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>
Showing
9 changed files
with
383 additions
and
2 deletions
core/__init__.py
0 → 100644
core/database.py
0 → 100644
| 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 |
core/generation.py
0 → 100644
| 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) |
core/history.py
0 → 100644
This diff is collapsed.
Click to expand it.
core/jewelry.py
0 → 100644
This diff is collapsed.
Click to expand it.
core/paths.py
0 → 100644
| 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 |
This diff is collapsed.
Click to expand it.
| ... | @@ -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( | ... | ... |
-
Please register or sign in to post a comment