Update app.py

app.py

@@ -1,67 +1,66 @@
-import gradio as gr
-import os
-import random
-import httpx
+# 导入区域
 import asyncio
 import gc
-import weakref
+import httpx
+import logging
+import os
 import psutil
+import random
 import threading
 import time
+import weakref
 from dataclasses import dataclass, field
-from typing import Any, Optional
-import contextlib
-import logging
-
-# 配置日志
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    handlers=[
-        logging.StreamHandler(),
-        logging.FileHandler("app.log", mode="a", encoding="utf-8"),
-    ],
-)
-logger = logging.getLogger(__name__)
+from typing import Any, Optional, Tuple, List, Dict
 
-# 内存监控配置
-MEMORY_CHECK_INTERVAL = 300  # 5分钟检查一次内存
-MEMORY_THRESHOLD_PERCENT = 80  # 内存使用百分比阈值
+import gradio as gr
 
-# 常量定义
-HTTP_STATUS_CENSORED = 451
+# 常量定义区域
+# HTTP 状态码
 HTTP_STATUS_OK = 200
+HTTP_STATUS_CENSORED = 451
+
+# 系统配置
 MAX_SEED = 2147483647  # (2**31 - 1)
 MAX_IMAGE_SIZE = 2048
-MIN_IMAGE_SIZE = 256  # Smallest dimension for SDXL like models often 512, but API might support smaller. Adjusted to API's limits.
-
-# 调试模式
+MIN_IMAGE_SIZE = 256
+MEMORY_CHECK_INTERVAL = 300  # 5分钟检查一次内存
+MEMORY_THRESHOLD_PERCENT = 80  # 内存使用百分比阈值
 DEBUG_MODE = os.environ.get("DEBUG_MODE", "false").lower() == "true"
 
-# 模型配置映射
+# API配置
+API_TOKEN = os.environ.get("API_TOKEN", "")
+if not API_TOKEN:
+    raise ValueError("环境变量中未设置API_TOKEN")
+
+# 外部链接
+DISCORD_LINK = os.environ.get("DISCORD_LINK", "https://discord.com/invite/AtRtbe9W8w")
+APP_INDEX_LINK = os.environ.get("APP_INDEX_LINK", "https://app.nieta.art/")
+APP_INDEX_ICON = "https://cdn-avatars.huggingface.co/v1/production/uploads/62be651a1e22ec8427aa7096/XQEUF5niIZXQbiOOxn8rQ.jpeg"
+
+# 模型配置
 MODEL_CONFIGS = {
-    # "ep3": "ep3.pth",
-    # "ep3latest": "ep3latest.pth",
     "ep6": "0622.pth",
 }
 
-# 连接池配置
+# HTTP 客户端配置
 HTTP_LIMITS = httpx.Limits(
     max_keepalive_connections=5, max_connections=10, keepalive_expiry=30.0
 )
-
 HTTP_TIMEOUT = httpx.Timeout(connect=30.0, read=300.0, write=30.0, pool=10.0)
 
-
-def validate_dimensions(width: int, height: int) -> tuple[int, int]:
-    """验证并调整图片尺寸"""
-    width = max(MIN_IMAGE_SIZE, min(int(width), MAX_IMAGE_SIZE))
-    height = max(MIN_IMAGE_SIZE, min(int(height), MAX_IMAGE_SIZE))
-    width = (width // 32) * 32
-    height = (height // 32) * 32
-    return width, height
+# 日志配置
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    handlers=[
+        logging.StreamHandler(),
+        logging.FileHandler("app.log", mode="a", encoding="utf-8"),
+    ],
+)
+logger = logging.getLogger(__name__)
 
 
+# 数据模型定义
 @dataclass
 class LuminaConfig:
     """Lumina模型配置"""
@@ -79,16 +78,111 @@ class ImageGenerationConfig:
     width: int = 1024
     height: int = 1024
     seed: int | None = None
-    use_polish: bool = False  # This wasn't exposed in UI, assuming false
+    use_polish: bool = False
     is_lumina: bool = True
     lumina_config: LuminaConfig = field(default_factory=LuminaConfig)
 
 
+# 全局变量
+image_client = None
+polish_client = None
+
+# 示例提示词
+example_titles = [
+    "Emotionless figure stares upward in heavy rain, drenched and defeated, set against a monochrome grey background.",
+    "Glamorous woman in vintage cheongsam on dim staircase, exuding 1930s Shanghai elegance and ethereal light",
+    "薄明のサイバーパンク都市で、銀灰獣人少女が近未来装備を纏い振り返るクールなアニメイラスト",
+    "都会の孤独な銀髪女性のアニメ風イラスト",
+    "仰视视角下，翘腿而坐的西装大佬俯视着你，气场冷峻危险",
+    "黄昏天台，白裙女子独望天空，惆怅静谧，风中飘扬",
+]
+
+full_prompts = {
+    example_titles[
+        0
+    ]: "looking_up, rain, heavy_rain, emotionless, high_angle, perspective, face_focus, from_above, from_side, grey_background, empty_eyes, defeated_look, limited_palette, partially_colored, monochrome_background, wet_hair, wet_clothes, messy_hair A figure stands motionless in torrential rain, their face captured from a high-angle perspective with vacant eyes staring upward into the downpour. Sodden strands of hair cling to pale skin, soaked clothing blending into the monochrome grey void surrounding them. The limited color palette isolates faint blue tones in their parted lips against the washed-out background, their expression of utter defeat frozen as raindrops streak vertically through the frame. Sharp focus on their emotionless features contrasts with the blurred chaos of falling water, the composition's oppressive atmosphere amplified by the overhead view's psychological distance.",
+    example_titles[
+        1
+    ]: "Elegant anime-style illustration of a glamorous woman standing on a dimly lit staircase, bathed in soft, ethereal light. Medium full-body shot. She has short, wavy black hair adorned with ornate golden hairpieces and a delicate feather. Her makeup is refined, with deep red lipstick and sharp eyeliner accentuating her poised expression. She wears a dark green high-collared cheongsam embroidered with a vividly colored crane and landscape motif, blending traditional Chinese elements with a touch of Art Deco. Her black lace gloves, beaded pearl chains draped across her body, and intricate earrings enhance her luxurious, vintage aura. A fur stole rests over her arms, adding texture and sophistication. The setting is subtly blurred, with falling sparkles and filtered light through windowpanes creating a dreamy, nostalgic atmosphere. The composition evokes 1930s Shanghai elegance, rendered in a painterly digital style with high fashion sensibility and dramatic lighting contrast.",
+    example_titles[
+        2
+    ]: "薄明のサイバーパンク都市を背景にした、ハイブリッドな獣人少女のアニメイラスト。構図はバストアップ〜ウエストまでの斜めアングルで、キャラクターが観覧者にやや振り返るような姿勢をとっている。彼女の顔は精巧な狼のマズルと人間的な表情が融合しており、瞳はグリッチ風の蛍光ブルーに輝くオッドアイ。毛皮は柔らかな銀灰色で、頬や耳の先端には淡い紫のグラデーションが入り、毛並みは丁寧に一本ずつ描写されている。衣装は近未来的なデザインのボディスーツで、胸元には蛍光回路のような文様が走り、左肩には機械義肢風のメタルアーマーが装着されている。耳にはピアスとデータ受信装置が取り付けられており、首には布製のスカーフと識別タグが結ばれている。背景にはネオンに照らされた廃墟ビルとホログラム広告が揺らめき、空にはサイバー・オーロラが流れている。カラーリングはダークトーンに蛍光アクセントが効いたネオ東京風配色（黒、紫、青、サイバーグリーン）。アートスタイルはセルルックアニメ＋ハードエッジグラデーション、ディテール豊かで表情と毛の質感に特に重点を置く。キャラクターの雰囲気はクールで孤独だが、どこか繊細な感情を秘めた印象。",
+    example_titles[
+        3
+    ]: "銀白色の長い髪を持つ若い女性が黒いコートを羽織り、都会の街中に佇む、メランコリックな雰囲気のアニメ風イラスト。背景にはグレースケールで描かれた、顔のない人々の群衆がぼかされて描かれ、彼女だけがシャープな焦点で際立っている。彼女は長いダークカラーのオーバーコートを着用し、物憂げで遠くを見つめるような表情を浮かべている。シーンは現代の都市の街角で設定されており、背景には建築的な要素が描かれている。配色はモノクロームで、黒・白・灰色を基調としたフィルム・ノワール風の雰囲気を醸し出している。アートスタイルはリアルな都市描写とアニメ的キャラクターデザインを融合させており、主人公の細部まで描き込まれた描写と、幽霊のように描かれた背景の人々とのコントラストによって、強い孤独感と疎外感を表現している。どんよりとした曇り空のライティングと都市背景が、重苦しい空気感を一層引き立てている。キャラクターデザインは、暗い服装に映える銀髪により、どこか神秘的または超自然的な要素を感じさせる。プロフェッショナルなデジタルイラストで、映画のような質感と強い感情的トーンを備え、現代社会における孤独や断絶のテーマを描いている。雰囲気としては、心理的または超自然的な物語の一場面を想起させる構成となっている。",
+    example_titles[
+        4
+    ]: "戏剧感十足的数字插画，描绘一位黑社会风格的长发美男子大佬坐姿翘着二郎腿，从大仰视视角俯视观者，气场强烈。他身穿剪裁锋利的深色西装，带有低调的细条纹图案，外套敞开，内衬深色马甲与整洁的衬衫，显得既优雅又危险。他的长发黑亮顺滑，自肩膀垂落，几缕发丝在环境光中微微反光。他一只手漫不经心地夹着一根燃着的雪茄，袅袅烟雾在昏暗中盘旋上升。他俯视下方，眉头微蹙，唇角紧绷，脸上带有一丝明显的不悦，神情冷峻。仰视视角放大了他身上的压迫感与支配力，整个人散发出沉稳却不可侵犯的气势。另一只手或搭在椅扶或搁在膝上，动作自信从容。背景为昏暗豪华的私人会客室或办公室，可见金属饰条、皮革沙发或昏黄灯光等元素。画面风格精致、电影感强烈，通过明暗对比、表情细节与强光源制造紧张与权威氛围。",
+    example_titles[
+        5
+    ]: "惆怅氛围的数字插画，描绘一位身穿白裙的女子独自站在天台上，在黄昏时分凝望天空。她的长发与白裙的裙摆在风中轻轻飘扬，画面充满静谧与动感的交织。她仰望着染上橙粉色霞光的天空，脸上露出若有所思、带着忧郁的神情，仿佛沉浸在回忆或思念之中。夕阳的余晖洒落在她身上与天台上，为画面增添一层柔和的金色光晕。她的站姿略带沉重，肩膀微微下垂，双手自然垂落或轻轻拽住裙摆，展现出内敛的情绪张力。天台背景简洁，设有低矮护栏、通风设备等结构，远处是模糊的城市天际线，进一步强化孤独氛围。风吹起她的发丝与裙摆，为静止的场景注入一丝淡淡的动态诗意。画风柔和写意，色彩渐变细腻，注重光影、空气感与人物情绪的融合，营造出宁静而感伤的黄昏场景。",
+}
+
+
+# 工具函数
+def validate_dimensions(width: int, height: int) -> tuple[int, int]:
+    """验证并调整图片尺寸"""
+    width = max(MIN_IMAGE_SIZE, min(int(width), MAX_IMAGE_SIZE))
+    height = max(MIN_IMAGE_SIZE, min(int(height), MAX_IMAGE_SIZE))
+    width = (width // 32) * 32
+    height = (height // 32) * 32
+    return width, height
+
+
+class PolishClient:
+    """提示词润色客户端"""
+
+    def __init__(self):
+        self.x_token = API_TOKEN
+        if not self.x_token:
+            raise ValueError("环境变量中未设置API_TOKEN")
+
+        self.url = "https://api.talesofai.cn/v3/gpt/dify/text-complete"
+        self.headers = {
+            "x-token": self.x_token,
+            "Content-Type": "application/json",
+            "x-nieta-app-version": "5.14.0",
+            "x-platform": "nieta-app/web",
+        }
+
+    async def polish_text(self, input_text: str) -> str:
+        """润色文本"""
+        payload = {
+            "query": "",
+            "response_mode": "blocking",
+            "preset_key": "latitude://28|live|running",
+            "inputs": {"query": input_text},
+        }
+
+        async with httpx.AsyncClient(timeout=HTTP_TIMEOUT) as client:
+            response = await client.post(self.url, headers=self.headers, json=payload)
+            if response.status_code == HTTP_STATUS_OK:
+                response_data = response.json()
+                polished_text = response_data.get("answer", input_text)
+                return polished_text.strip()
+            else:
+                logger.warning(f"润色API调用失败: {response.status_code}")
+                return input_text
+
+
+async def polish_prompt(prompt: str) -> str:
+    """提示词润色函数 - 使用外部API润色"""
+    global polish_client
+    if polish_client is None:
+        polish_client = PolishClient()
+
+    logger.info(f"润色提示词: {prompt[:50]}...")
+    polished_prompt = await polish_client.polish_text(prompt)
+    logger.info("提示词润色完成")
+    return polished_prompt
+
+
+# 核心类定义
 class ImageClient:
     """图像生成客户端"""
 
     def __init__(self) -> None:
-        self.x_token = os.environ.get("API_TOKEN", "")
+        self.x_token = API_TOKEN
         if not self.x_token:
             raise ValueError("环境变量中未设置API_TOKEN")
 
@@ -100,18 +194,16 @@ class ImageClient:
         self.polling_interval = 3.0
         self.default_headers = {
             "Content-Type": "application/json",
-            "x-platform": "nieta-app/web",  # Or a generic identifier if preferred
+            "x-platform": "nieta-app/web",
             "X-Token": self.x_token,
         }
 
-        # 预创建HTTP客户端配置，但不保持持久连接
         self._client_config = {
             "limits": HTTP_LIMITS,
             "timeout": HTTP_TIMEOUT,
             "headers": self.default_headers,
         }
 
-        # 用于跟踪活跃任务的弱引用集合
         self._active_tasks = weakref.WeakSet()
 
     def _prepare_prompt_data(
@@ -153,7 +245,7 @@ class ImageClient:
             "seed": config.seed,
             "meta": {"entrance": "PICTURE,PURE"},
             "context_model_series": None,
-            "negative_freetext": "",  # Negative handled in rawPrompt
+            "negative_freetext": "",
             "advanced_translator": config.use_polish,
         }
         if config.is_lumina:
@@ -171,7 +263,6 @@ class ImageClient:
     async def _poll_task_status(self, task_uuid: str) -> dict[str, Any]:
         """轮询任务状态 - 优化内存使用和连接管理"""
         status_url = self.lumina_task_status_url.format(task_uuid=task_uuid)
-
         poll_timeout = httpx.Timeout(connect=10.0, read=30.0, write=10.0, pool=5.0)
 
         try:
@@ -181,7 +272,6 @@ class ImageClient:
                 for attempt in range(self.max_polling_attempts):
                     try:
                         response = await client.get(status_url)
-
                         if response.status_code != HTTP_STATUS_OK:
                             logger.warning(f"轮询失败 - 状态码: {response.status_code}")
                             return {
@@ -189,7 +279,6 @@ class ImageClient:
                                 "error": f"获取任务状态失败: {response.status_code} - {response.text[:200]}",
                             }
 
-                        # 解析JSON响应
                         try:
                             result = response.json()
                         except Exception as e:
@@ -200,7 +289,6 @@ class ImageClient:
                             }
 
                         task_status = result.get("task_status")
-
                         if task_status == "SUCCESS":
                             artifacts = result.get("artifacts", [])
                             if artifacts and len(artifacts) > 0:
@@ -222,20 +310,18 @@ class ImageClient:
                             return {"success": False, "error": "任务超时"}
 
                         await asyncio.sleep(self.polling_interval)
-
-                        # 每10次轮询执行一次垃圾回收以防止内存累积
                         if attempt % 10 == 0 and attempt > 0:
                             gc.collect()
                             logger.debug(f"轮询第{attempt}次，执行内存清理")
 
                     except asyncio.TimeoutError:
                         logger.warning(f"轮询超时 - 尝试 {attempt}")
-                        if attempt >= 3:  # 连续3次超时后退出
+                        if attempt >= 3:
                             return {
                                 "success": False,
                                 "error": "连接超时，请检查网络状况",
                             }
-                        await asyncio.sleep(self.polling_interval * 2)  # 延长等待时间
+                        await asyncio.sleep(self.polling_interval * 2)
                         continue
                     except Exception as e:
                         logger.warning(f"轮询异常 - 尝试 {attempt}: {str(e)}")
@@ -265,23 +351,21 @@ class ImageClient:
         cfg: float,
         steps: int,
         model_name: str = "ep6",
+        use_polish: bool = False,
     ) -> tuple[str | None, str | None]:
         """生成图片"""
         try:
-            # 获取模型路径
             model_path = MODEL_CONFIGS.get(model_name, MODEL_CONFIGS["ep6"])
-
-            # 准备配置
             config = ImageGenerationConfig(
                 prompts=self._prepare_prompt_data(prompt, negative_prompt),
                 width=width,
                 height=height,
                 seed=seed,
+                use_polish=use_polish,
                 is_lumina=True,
                 lumina_config=LuminaConfig(model_name=model_path, cfg=cfg, step=steps),
             )
 
-            # 发送生成请求
             async with httpx.AsyncClient(**self._client_config) as client:
                 payload = self._build_payload(config)
                 if DEBUG_MODE:
@@ -302,18 +386,14 @@ class ImageClient:
 
                 if response.status_code == HTTP_STATUS_CENSORED:
                     return None, "内容不合规"
-
-                # 处理并发限制错误
                 if response.status_code == 433:
                     return None, "⏳ 服务器正忙，同时生成的图片数量已达上限，请稍后重试"
-
                 if response.status_code != HTTP_STATUS_OK:
                     return (
                         None,
                         f"API请求失败: {response.status_code} - {response.text[:200]}",
                     )
 
-                # API直接返回UUID字符串（根据model_studio的实现）
                 content = response.text.strip()
                 task_uuid = content.replace('"', "")
 
@@ -323,7 +403,6 @@ class ImageClient:
                 if not task_uuid:
                     return None, f"未获取到任务ID，API响应: {response.text[:200]}"
 
-                # 轮询任务状态
                 result = await self._poll_task_status(task_uuid)
                 if result["success"]:
                     return result["image_url"], None
@@ -334,13 +413,11 @@ class ImageClient:
             logger.error(f"生成图片异常: {str(e)}")
             return None, f"生成图片时发生错误: {str(e)}"
         finally:
-            # 确保资源清理
             gc.collect()
 
     def cleanup(self):
         """清理资源"""
         try:
-            # 清理活跃任务
             self._active_tasks.clear()
             gc.collect()
             logger.info("ImageClient资源清理完成")
@@ -348,24 +425,6 @@ class ImageClient:
             logger.error(f"资源清理异常: {str(e)}")
 
 
-# 创建图片生成客户端实例
-image_client: Optional[ImageClient] = None
-
-
-def get_image_client() -> ImageClient:
-    """获取图片生成客户端实例 - 单例模式"""
-    global image_client
-    if image_client is None:
-        try:
-            image_client = ImageClient()
-            logger.info("ImageClient初始化成功")
-        except Exception as e:
-            logger.error(f"ImageClient初始化失败: {e}")
-            raise
-    return image_client
-
-
-# 内存监控
 class MemoryMonitor:
     """内存监控器"""
 
@@ -394,21 +453,15 @@ class MemoryMonitor:
         """内存监控循环"""
         while self.running:
             try:
-                # 获取当前进程内存使用情况
                 process = psutil.Process()
                 memory_info = process.memory_info()
                 memory_mb = memory_info.rss / 1024 / 1024
-
-                # 获取系统总内存并计算使用百分比
-                system_memory = psutil.virtual_memory()
                 memory_percent = process.memory_percent()
 
-                # 记录内存使用情况（基于百分比而非固定阈值）
                 if memory_percent > MEMORY_THRESHOLD_PERCENT:
                     logger.warning(
                         f"内存使用量过高: {memory_mb:.2f} MB ({memory_percent:.1f}%)"
                     )
-                    # 执行垃圾回收
                     gc.collect()
                     logger.info("执行垃圾回收")
                 else:
@@ -422,73 +475,45 @@ class MemoryMonitor:
                 time.sleep(MEMORY_CHECK_INTERVAL)
 
 
-# 全局内存监控实例
-memory_monitor = MemoryMonitor()
+# 辅助函数
+def get_image_client() -> ImageClient:
+    """获取图片生成客户端实例 - 单例模式"""
+    global image_client
+    if image_client is None:
+        try:
+            image_client = ImageClient()
+            logger.info("ImageClient初始化成功")
+        except Exception as e:
+            logger.error(f"ImageClient初始化失败: {e}")
+            raise
+    return image_client
 
 
-# 优雅关闭处理
 async def cleanup_resources():
     """清理应用资源"""
-    global image_client, memory_monitor
-
+    global image_client, memory_monitor, polish_client
     try:
-        # 停止内存监控
         memory_monitor.stop_monitoring()
-
-        # 清理图片客户端
         if image_client is not None:
             image_client.cleanup()
             image_client = None
-
-        # 强制垃圾回收
+        if polish_client is not None:
+            polish_client = None
         gc.collect()
-
-        # 记录最终内存使用情况
         try:
             process = psutil.Process()
             memory_mb = process.memory_info().rss / 1024 / 1024
             logger.info(f"清理后内存使用: {memory_mb:.2f} MB")
         except Exception:
             pass
-
         logger.info("应用资源清理完成")
     except Exception as e:
         logger.error(f"资源清理异常: {str(e)}")
 
 
-# Example prompts
-example_titles = [
-    "Emotionless figure stares upward in heavy rain, drenched and defeated, set against a monochrome grey background.",
-    "Glamorous woman in vintage cheongsam on dim staircase, exuding 1930s Shanghai elegance and ethereal light",
-    "薄明のサイバーパンク都市で、銀灰獣人少女が近未来装備を纏い振り返るクールなアニメイラスト",
-    "都会の孤独な銀髪女性のアニメ風イラスト",
-    "仰视视角下，翘腿而坐的西装大佬俯视着你，气场冷峻危险",
-    "黄昏天台，白裙女子独望天空，惆怅静谧，风中飘扬",
-]
-full_prompts = {
-    example_titles[
-        0
-    ]: "looking_up, rain, heavy_rain, emotionless, high_angle, perspective, face_focus, from_above, from_side, grey_background, empty_eyes, defeated_look, limited_palette, partially_colored, monochrome_background, wet_hair, wet_clothes, messy_hair A figure stands motionless in torrential rain, their face captured from a high-angle perspective with vacant eyes staring upward into the downpour. Sodden strands of hair cling to pale skin, soaked clothing blending into the monochrome grey void surrounding them. The limited color palette isolates faint blue tones in their parted lips against the washed-out background, their expression of utter defeat frozen as raindrops streak vertically through the frame. Sharp focus on their emotionless features contrasts with the blurred chaos of falling water, the composition's oppressive atmosphere amplified by the overhead view's psychological distance.",
-    example_titles[
-        1
-    ]: "Elegant anime-style illustration of a glamorous woman standing on a dimly lit staircase, bathed in soft, ethereal light. Medium full-body shot. She has short, wavy black hair adorned with ornate golden hairpieces and a delicate feather. Her makeup is refined, with deep red lipstick and sharp eyeliner accentuating her poised expression. She wears a dark green high-collared cheongsam embroidered with a vividly colored crane and landscape motif, blending traditional Chinese elements with a touch of Art Deco. Her black lace gloves, beaded pearl chains draped across her body, and intricate earrings enhance her luxurious, vintage aura. A fur stole rests over her arms, adding texture and sophistication. The setting is subtly blurred, with falling sparkles and filtered light through windowpanes creating a dreamy, nostalgic atmosphere. The composition evokes 1930s Shanghai elegance, rendered in a painterly digital style with high fashion sensibility and dramatic lighting contrast.",
-    example_titles[
-        2
-    ]: "薄明のサイバーパンク都市を背景にした、ハイブリッドな獣人少女のアニメイラスト。構図はバストアップ〜ウエストまでの斜めアングルで、キャラクターが観覧者にやや振り返るような姿勢をとっている。彼女の顔は精巧な狼のマズルと人間的な表情が融合しており、瞳はグリッチ風の蛍光ブルーに輝くオッドアイ。毛皮は柔らかな銀灰色で、頬や耳の先端には淡い紫のグラデーションが入り、毛並みは丁寧に一本ずつ描写されている。衣装は近未来的なデザインのボディスーツで、胸元には蛍光回路のような文様が走り、左肩には機械義肢風のメタルアーマーが装着されている。耳にはピアスとデータ受信装置が取り付けられており、首には布製のスカーフと識別タグが結ばれている。背景にはネオンに照らされた廃墟ビルとホログラム広告が揺らめき、空にはサイバー・オーロラが流れている。カラーリングはダークトーンに蛍光アクセントが効いたネオ東京風配色（黒、紫、青、サイバーグリーン）。アートスタイルはセルルックアニメ＋ハードエッジグラデーション、ディテール豊かで表情と毛の質感に特に重点を置く。キャラクターの雰囲気はクールで孤独だが、どこか繊細な感情を秘めた印象。",
-    example_titles[
-        3
-    ]: "銀白色の長い髪を持つ若い女性が黒いコートを羽織り、都会の街中に佇む、メランコリックな雰囲気のアニメ風イラスト。背景にはグレースケールで描かれた、顔のない人々の群衆がぼかされて描かれ、彼女だけがシャープな焦点で際立っている。彼女は長いダークカラーのオーバーコートを着用し、物憂げで遠くを見つめるような表情を浮かべている。シーンは現代の都市の街角で設定されており、背景には建築的な要素が描かれている。配色はモノクロームで、黒・白・灰色を基調としたフィルム・ノワール風の雰囲気を醸し出している。アートスタイルはリアルな都市描写とアニメ的キャラクターデザインを融合させており、主人公の細部まで描き込まれた描写と、幽霊のように描かれた背景の人々とのコントラストによって、強い孤独感と疎外感を表現している。どんよりとした曇り空のライティングと都市背景が、重苦しい空気感を一層引き立てている。キャラクターデザインは、暗い服装に映える銀髪により、どこか神秘的または超自然的な要素を感じさせる。プロフェッショナルなデジタルイラストで、映画のような質感と強い感情的トーンを備え、現代社会における孤独や断絶のテーマを描いている。雰囲気としては、心理的または超自然的な物語の一場面を想起させる構成となっている。",
-    example_titles[
-        4
-    ]: "戏剧感十足的数字插画，描绘一位黑社会风格的长发美男子大佬坐姿翘着二郎腿，从大仰视视角俯视观者，气场强烈。他身穿剪裁锋利的深色西装，带有低调的细条纹图案，外套敞开，内衬深色马甲与整洁的衬衫，显得既优雅又危险。他的长发黑亮顺滑，自肩膀垂落，几缕发丝在环境光中微微反光。他一只手漫不经心地夹着一根燃着的雪茄，袅袅烟雾在昏暗中盘旋上升。他俯视下方，眉头微蹙，唇角紧绷，脸上带有一丝明显的不悦，神情冷峻。仰视视角放大了他身上的压迫感与支配力，整个人散发出沉稳却不可侵犯的气势。另一只手或搭在椅扶或搁在膝上，动作自信从容。背景为昏暗豪华的私人会客室或办公室，可见金属饰条、皮革沙发或昏黄灯光等元素。画面风格精致、电影感强烈，通过明暗对比、表情细节与强光源制造紧张与权威氛围。",
-    example_titles[
-        5
-    ]: "惆怅氛围的数字插画，描绘一位身穿白裙的女子独自站在天台上，在黄昏时分凝望天空。她的长发与白裙的裙摆在风中轻轻飘扬，画面充满静谧与动感的交织。她仰望着染上橙粉色霞光的天空，脸上露出若有所思、带着忧郁的神情，仿佛沉浸在回忆或思念之中。夕阳的余晖洒落在她身上与天台上，为画面增添一层柔和的金色光晕。她的站姿略带沉重，肩膀微微下垂，双手自然垂落或轻轻拽住裙摆，展现出内敛的情绪张力。天台背景简洁，设有低矮护栏���通风设备等结构，远处是模糊的城市天际线，进一步强化孤独氛围。风吹起她的发丝与裙摆，为静止的场景注入一丝淡淡的动态诗意。画风柔和写意，色彩渐变细腻，注重光影、空气感与人物情绪的融合，营造出宁静而感伤的黄昏场景。",
-}
-
-
 async def infer(
     prompt_text,
+    use_polish_val,
     seed_val,
     randomize_seed_val,
     width_val,
@@ -498,11 +523,8 @@ async def infer(
     model_name_val,
 ):
     """推理函数 - 优化内存管理和错误处理"""
-    client = None
     try:
         logger.info(f"开始生成图像: {prompt_text[:50]}...")
-
-        # 获取客户端实例
         try:
             client = get_image_client()
         except Exception as e:
@@ -512,21 +534,27 @@ async def infer(
         if not prompt_text.strip():
             raise gr.Error("提示词不能为空")
 
+        final_prompt = prompt_text
+        if use_polish_val:
+            final_prompt = await polish_prompt(prompt_text)
+
+        # 为生成添加系统前缀，但不在UI显示
+        system_prefix = "You are an assistant designed to generate anime images with the highest degree of image-text alignment based on danbooru tags. <Prompt Start>"
+        generation_prompt = f"{system_prefix} {final_prompt}"
+
         current_seed = int(seed_val)
         if randomize_seed_val:
             current_seed = random.randint(0, MAX_SEED)
 
-        # 验证并调整尺寸
         width_val, height_val = validate_dimensions(width_val, height_val)
 
-        # 验证其他参数
         if not (1.0 <= float(cfg_val) <= 20.0):
             raise gr.Error("CFG Scale 必须在 1.0 到 20.0 之间")
         if not (1 <= int(steps_val) <= 50):
             raise gr.Error("Steps 必须在 1 到 50 之间")
 
         image_url, error = await client.generate_image(
-            prompt=prompt_text,
+            prompt=generation_prompt,  # 使用包含系统前缀的提示词
             negative_prompt="",
             seed=current_seed,
             width=width_val,
@@ -534,6 +562,7 @@ async def infer(
             cfg=float(cfg_val),
             steps=int(steps_val),
             model_name=model_name_val,
+            use_polish=False,
         )
 
         if error:
@@ -541,153 +570,17 @@ async def infer(
             raise gr.Error(error)
 
         logger.info("图像生成成功")
-        return image_url, current_seed
+        return image_url, final_prompt, current_seed  # 返回不带系统前缀的用户提示词
 
     except gr.Error:
-        # 重新抛出Gradio错误
         raise
     except Exception as e:
         logger.error(f"推理过程异常: {str(e)}")
         raise gr.Error(f"生成图像时发生意外错误: {str(e)}")
     finally:
-        # 确保内存清理
         gc.collect()
 
 
-# Links for HTML header
-DISCORD_LINK = os.environ.get(
-    "DISCORD_LINK", "https://discord.com/invite/AtRtbe9W8w"
-)  # Example
-APP_INDEX_LINK = os.environ.get("APP_INDEX_LINK", "https://app.nieta.art/")  # Example
-APP_INDEX_ICON = "https://cdn-avatars.huggingface.co/v1/production/uploads/62be651a1e22ec8427aa7096/XQEUF5niIZXQbiOOxn8rQ.jpeg"  # Using HF logo
-
-
-with gr.Blocks(theme=gr.themes.Soft(), title="Lumina Image Playground") as demo:
-    gr.Markdown("<h1>🎨 Lumina Text-to-Image Playground | Lumina 文本生图工具</h1>")
-    gr.Markdown(
-        "Fine-tuned Lumina model specialized for anime/manga style generation! Supports Chinese, English, and Japanese prompts. Model under active development - more exciting features coming soon!\n 🌸 专为二次元风格优化的Lumina模型！支持中文、英文、日文三语提示词，让您的创意无界限！模型持续优化中，敬请期待。"
-    )
-
-    gr.HTML(f"""
-        <div style="display: flex; justify-content: flex-start; align-items: center; gap: 15px; margin-bottom: 20px; padding: 10px;">
-            <a href="{DISCORD_LINK}" target="_blank" style="text-decoration: none; color: #5865F2; font-weight: 500; display: inline-flex; align-items: center; gap: 5px;">
-                <img src="https://assets-global.website-files.com/6257adef93867e50d84d30e2/636e0a69f118df70ad7828d4_icon_clyde_blurple_RGB.svg" alt="Discord" style="height: 20px;">
-                Join Discord | 加入Discord
-            </a>
-            <a href="{APP_INDEX_LINK}" target="_blank" style="text-decoration: none; color: #333; font-weight: 500; display: inline-flex; align-items: center; gap: 5px;">
-                <img src="{APP_INDEX_ICON}" alt="App Index" style="height: 20px; border-radius: 3px;">
-                Nieta Home | 捏Ta主页
-            </a>
-        </div>
-    """)
-
-    with gr.Row(variant="panel"):
-        with gr.Column(scale=2):  # Controls Panel
-            gr.Markdown("## ⚙️ Generation Controls | 生成控制")
-            prompt = gr.Textbox(
-                label="Prompt | 提示词",
-                lines=5,
-                placeholder="e.g., A majestic dragon soaring through a cyberpunk city skyline, neon lights reflecting off its scales, intricate details. | 例如：一条威武的巨龙翱翔在赛博朋克城市天际线，霓虹灯映照在它的鳞片上，细节精美。",
-                info="Describe the image you want to create. | 描述您想要创建的图像。",
-            )
-
-            with gr.Accordion("🔧 Advanced Settings | 高级设置", open=True):
-                model_name = gr.Dropdown(
-                    label="Model Version | 模型版本",
-                    choices=list(MODEL_CONFIGS.keys()),
-                    value="ep6",
-                    info="Select the generation model. | 选择生成模型。",
-                )
-                with gr.Row():
-                    cfg = gr.Slider(
-                        label="CFG Scale | CFG缩放",
-                        minimum=1.0,
-                        maximum=20.0,
-                        step=0.1,
-                        value=5.5,
-                        info="Guidance strength. Higher values adhere more to prompt. | 引导强度，更高的值更贴近提示词。",
-                    )
-                    steps = gr.Slider(
-                        label="Sampling Steps | 采样步数",
-                        minimum=1,
-                        maximum=50,
-                        step=1,
-                        value=30,
-                        info="Number of steps. More steps can improve quality but take longer. | 步数，更多步数可提高质量但耗时更长。",
-                    )
-
-                with gr.Row():
-                    width = gr.Slider(
-                        label="Width | 宽度",
-                        minimum=MIN_IMAGE_SIZE,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=1024,
-                    )
-                    height = gr.Slider(
-                        label="Height | 高度",
-                        minimum=MIN_IMAGE_SIZE,
-                        maximum=MAX_IMAGE_SIZE,
-                        step=32,
-                        value=1024,
-                    )
-
-                with gr.Row():
-                    seed = gr.Slider(
-                        label="Seed | 种子",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=random.randint(0, MAX_SEED),
-                    )
-                    randomize_seed = gr.Checkbox(
-                        label="Randomize Seed | 随机种子",
-                        value=True,
-                        info="Use a new random seed for each generation if checked. | 勾选后每次生成使用新的随机种子。",
-                    )
-
-            run_button = gr.Button(
-                "🚀 Generate Image | 生成图像", variant="primary", scale=0
-            )
-
-            with gr.Group():
-                gr.Markdown("### ✨ Example Prompts | 示例提示词")
-                for i, title in enumerate(example_titles):
-                    btn = gr.Button(title)
-                    btn.click(lambda t=title: full_prompts[t], outputs=[prompt])
-
-        with gr.Column(scale=3):  # Output Panel
-            gr.Markdown("## 🖼️ Generated Image | 生成图像")
-            result_image = gr.Image(
-                label="Output Image | 输出图像",
-                show_label=False,
-                type="filepath",
-                height=600,  # Max display height
-                show_download_button=True,
-                interactive=False,
-                elem_id="result_image_display",  # for potential CSS targeting if needed
-            )
-            generated_seed_info = gr.Textbox(
-                label="Seed Used | 使用的种子",
-                interactive=False,
-                placeholder="The seed for the generated image will appear here. | 生成图像所使用的种子值将显示在此处。",
-            )
-
-    # Event Handlers
-    inputs_list = [prompt, seed, randomize_seed, width, height, cfg, steps, model_name]
-    outputs_list = [result_image, generated_seed_info]
-
-    run_button.click(
-        fn=infer, inputs=inputs_list, outputs=outputs_list, api_name="generate_image"
-    )
-    prompt.submit(
-        fn=infer,
-        inputs=inputs_list,
-        outputs=outputs_list,
-        api_name="generate_image_submit",
-    )
-
-
 def setup_signal_handlers():
     """设置信号处理器以实现优雅关闭"""
     import signal
@@ -702,20 +595,178 @@ def setup_signal_handlers():
             logger.error(f"清理资源时出错: {str(e)}")
         logger.info("应用已安全关闭")
 
-    # 注册信号处理器
     signal.signal(signal.SIGINT, cleanup_handler)
     signal.signal(signal.SIGTERM, cleanup_handler)
-
-    # 注册退出处理器
     atexit.register(cleanup_handler)
 
 
-if __name__ == "__main__":
+# UI 构建
+def build_ui():
+    """构建Gradio UI"""
+    with gr.Blocks(theme=gr.themes.Soft(), title="Lumina Image Playground") as demo:
+        gr.Markdown("<h1>🎨 Lumina Text-to-Image Playground | Lumina 文本生图工具</h1>")
+        gr.Markdown(
+            "Fine-tuned Lumina model specialized for anime/manga style generation! Supports Chinese, English, and Japanese prompts. Model under active development - more exciting features coming soon!\n 🌸 专为二次元风格优化的Lumina模型！支持中文、英文、日文三语提示词，让您的创意无界限！模型持续优化中，敬请期待。"
+        )
+
+        gr.HTML(f"""
+            <div style="display: flex; justify-content: flex-start; align-items: center; gap: 15px; margin-bottom: 20px; padding: 10px;">
+                <a href="{DISCORD_LINK}" target="_blank" style="text-decoration: none; color: #5865F2; font-weight: 500; display: inline-flex; align-items: center; gap: 5px;">
+                    <img src="https://assets-global.website-files.com/6257adef93867e50d84d30e2/636e0a69f118df70ad7828d4_icon_clyde_blurple_RGB.svg" alt="Discord" style="height: 20px;">
+                    Join Discord | 加入Discord
+                </a>
+                <a href="{APP_INDEX_LINK}" target="_blank" style="text-decoration: none; color: #333; font-weight: 500; display: inline-flex; align-items: center; gap: 5px;">
+                    <img src="{APP_INDEX_ICON}" alt="App Index" style="height: 20px; border-radius: 3px;">
+                    Nieta Home | 捏Ta主页
+                </a>
+            </div>
+        """)
+
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=2):  # Controls Panel
+                gr.Markdown("## ⚙️ Generation Controls | 生成控制")
+                prompt = gr.Textbox(
+                    label="Prompt | 提示词",
+                    lines=5,
+                    placeholder="e.g., A majestic dragon soaring through a cyberpunk city skyline, neon lights reflecting off its scales, intricate details. | 例如：一条威武的巨龙翱翔在赛博朋克城市天际线，霓虹灯映照在它的鳞片上，细节精美。",
+                    info="Describe the image you want to create. | 描述您想要创建的图像。",
+                )
+
+                use_polish = gr.Checkbox(
+                    label="✨ Auto Polish Prompt | 自动润色提示词",
+                    value=True,
+                    info="Automatically optimize and enhance your prompt for better results. | 自动优化和增强您的提示词以获得更好的效果。",
+                )
+
+                run_button = gr.Button(
+                    "🚀 Generate Image | 生成图像", variant="primary", scale=0
+                )
+
+                with gr.Accordion("🔧 Advanced Settings | 高级设置", open=False):
+                    model_name = gr.Dropdown(
+                        label="Model Version | 模型版本",
+                        choices=list(MODEL_CONFIGS.keys()),
+                        value="ep6",
+                        info="Select the generation model. | 选择生成模型。",
+                    )
+                    with gr.Row():
+                        cfg = gr.Slider(
+                            label="CFG Scale | CFG缩放",
+                            minimum=1.0,
+                            maximum=20.0,
+                            step=0.1,
+                            value=5.5,
+                            info="Guidance strength. Higher values adhere more to prompt. | 引导强度，更高的值更贴近提示词。",
+                        )
+                        steps = gr.Slider(
+                            label="Sampling Steps | 采样步数",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=30,
+                            info="Number of steps. More steps can improve quality but take longer. | 步数，更多步数可提高质量但耗时更长。",
+                        )
+
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width | 宽度",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=1024,
+                        )
+                        height = gr.Slider(
+                            label="Height | 高度",
+                            minimum=MIN_IMAGE_SIZE,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=1024,
+                        )
+
+                    with gr.Row():
+                        seed = gr.Slider(
+                            label="Seed | 种子",
+                            minimum=0,
+                            maximum=MAX_SEED,
+                            step=1,
+                            value=random.randint(0, MAX_SEED),
+                        )
+                        randomize_seed = gr.Checkbox(
+                            label="Randomize Seed | 随机种子",
+                            value=True,
+                            info="Use a new random seed for each generation if checked. | 勾选后每次生成使用新的随机种子。",
+                        )
+
+                with gr.Group():
+                    gr.Markdown("### ✨ Example Prompts | 示例提示词")
+                    for i, title in enumerate(example_titles):
+                        btn = gr.Button(title)
+                        btn.click(lambda t=title: full_prompts[t], outputs=[prompt])
+
+            with gr.Column(scale=3):  # Output Panel
+                gr.Markdown("## 🖼️ Generated Image | 生成图像")
+                result_image = gr.Image(
+                    label="Output Image | 输出图像",
+                    show_label=False,
+                    type="filepath",
+                    height=600,
+                    show_download_button=True,
+                    interactive=False,
+                    elem_id="result_image_display",
+                )
+                used_prompt_info = gr.Textbox(
+                    label="Used Prompt | 使用的提示词",
+                    interactive=False,
+                    lines=3,
+                    placeholder="The actual prompt used for generation will appear here. | 生成时实际使用的提示词将显示在此处。",
+                )
+                generated_seed_info = gr.Textbox(
+                    label="Seed Used | 使用的种子",
+                    interactive=False,
+                    placeholder="The seed for the generated image will appear here. | 生成图像所使用的种子值将显示在此处。",
+                )
+
+        # 事件处理
+        inputs_list = [
+            prompt,
+            use_polish,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            cfg,
+            steps,
+            model_name,
+        ]
+        outputs_list = [result_image, used_prompt_info, generated_seed_info]
+
+        run_button.click(
+            fn=infer,
+            inputs=inputs_list,
+            outputs=outputs_list,
+            api_name="generate_image",
+        )
+        prompt.submit(
+            fn=infer,
+            inputs=inputs_list,
+            outputs=outputs_list,
+            api_name="generate_image_submit",
+        )
+
+    return demo
+
+
+# 主函数
+def main():
+    """主函数"""
+    global memory_monitor
+
     try:
         # 设置信号处理器
         setup_signal_handlers()
 
         # 启动内存监控
+        memory_monitor = MemoryMonitor()
         memory_monitor.start_monitoring()
 
         if DEBUG_MODE:
@@ -749,14 +800,15 @@ if __name__ == "__main__":
             pass
 
         logger.info("启动Gradio应用...")
+        demo = build_ui()
         demo.launch(
             debug=DEBUG_MODE,
             show_error=True,
-            server_name="0.0.0.0",  # 允许外部访问
+            server_name="0.0.0.0",
             server_port=7860,
             share=False,
-            max_threads=40,  # 限制线程数
-            favicon_path=None,  # 避免额外的文件加载
+            max_threads=40,
+            favicon_path=None,
         )
 
     except KeyboardInterrupt:
@@ -771,3 +823,7 @@ if __name__ == "__main__":
         except Exception as e:
             logger.error(f"最终清理时出错: {str(e)}")
         logger.info("应用已退出")
+
+
+if __name__ == "__main__":
+    main()