import os import logging import gradio as gr from config import config from processor import process_batch_sync, expand_tags_sync logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) os.makedirs(config.OUTPUT_DIR, exist_ok=True) CUSTOM_CSS = """ .gradio-container { max-width: 1200px !important; font-family: 'Inter', sans-serif !important; } .title-block { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 16px; padding: 28px 32px; margin-bottom: 20px; color: white; text-align: center; box-shadow: 0 8px 32px rgba(102, 126, 234, 0.25); } .title-block h1 { font-size: 28px; font-weight: 700; margin: 0 0 8px 0; letter-spacing: -0.5px; } .title-block p { font-size: 14px; opacity: 0.9; margin: 0; } .primary-btn { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important; border: none !important; border-radius: 12px !important; color: white !important; font-weight: 600 !important; font-size: 16px !important; padding: 14px 28px !important; box-shadow: 0 4px 15px rgba(102, 126, 234, 0.35) !important; transition: all 0.3s ease !important; } .primary-btn:hover { transform: translateY(-2px) !important; box-shadow: 0 6px 20px rgba(102, 126, 234, 0.5) !important; } .nsfw-toggle { border: 1px solid rgba(255, 107, 107, 0.3) !important; border-radius: 10px !important; padding: 12px !important; background: rgba(255, 107, 107, 0.05) !important; } .mode-toggle-portrait { border: 1px solid rgba(102, 126, 234, 0.3) !important; border-radius: 10px !important; padding: 12px !important; background: rgba(102, 126, 234, 0.05) !important; margin-bottom: 10px; } .tip-text { font-size: 13px; color: rgba(156, 163, 175, 0.9); background: rgba(102, 126, 234, 0.08); border-radius: 10px; padding: 12px 16px; border-left: 3px solid rgba(102, 126, 234, 0.5); margin-bottom: 15px; } .warning-text { font-size: 12px; color: #ef4444; margin-top: 4px; font-weight: bold; } .summary-box textarea { font-family: 'JetBrains Mono', Consolas, monospace !important; font-size: 13px !important; line-height: 1.6 !important; border-radius: 10px !important; } .upload-area { border: 2px dashed rgba(102, 126, 234, 0.3) !important; border-radius: 14px !important; transition: all 0.3s ease !important; } .upload-area:hover { border-color: rgba(102, 126, 234, 0.6) !important; background: rgba(102, 126, 234, 0.05) !important; } .fetch-btn { background: linear-gradient(135deg, #43e97b 0%, #38f9d7 100%) !important; border: none !important; color: #1a1a2e !important; font-weight: 600 !important; } """ def fetch_models_from_env(): models = config.AVAILABLE_MODELS if not models: return gr.update(choices=[], value=None), "⚠️ HF 环境变量 AVAILABLE_MODELS 未配置或为空" return gr.update(choices=models, value=models[0]), f"✅ 已读取到 {len(models)} 个可用模型" def _prepare_params(api_key, base_url, model_dropdown, model_text, portrait_suffix, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style): safe_api_key = (api_key or "").strip() or None safe_base_url = (base_url or "").strip() or None safe_model = (model_text or "").strip() or (model_dropdown or "").strip() or None safe_portrait_suffix = (portrait_suffix or "").strip() enabled_dims = [] if dim_appearance: enabled_dims.append("appearance") if dim_body: enabled_dims.append("body") if dim_clothing: enabled_dims.append("clothing") if dim_pose: enabled_dims.append("pose") if dim_nsfw_detail: enabled_dims.append("nsfw_detail") if dim_composition: enabled_dims.append("composition") if dim_background: enabled_dims.append("background") if dim_style: enabled_dims.append("style") if not enabled_dims: enabled_dims = ["appearance"] return safe_api_key, safe_base_url, safe_model, safe_portrait_suffix, enabled_dims def analyze_images(files, provider, api_key, base_url, model_dropdown, model_text, images_per_request, nsfw_mode, nsfw_max_rolls, portrait_mode, portrait_suffix, custom_prompt, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style): if not files: return None, "❌ 请上传至少一张图片或一个 ZIP 压缩包", [] input_paths = [f if isinstance(f, str) else f.name for f in files] safe_api_key, safe_base_url, safe_model, safe_portrait_suffix, enabled_dims = _prepare_params( api_key, base_url, model_dropdown, model_text, portrait_suffix, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style ) try: zip_paths, summary, gallery_data = process_batch_sync( input_paths=input_paths, provider=provider, api_key=safe_api_key, base_url=safe_base_url, model=safe_model, custom_prompt=(custom_prompt or "").strip() or None, images_per_request=int(images_per_request), nsfw=nsfw_mode, nsfw_max_rolls=int(nsfw_max_rolls), enabled_dims=enabled_dims, portrait=portrait_mode, portrait_suffix=safe_portrait_suffix ) return zip_paths, summary, gallery_data except Exception as e: logger.error(f"反推失败: {e}") return None, f"❌ 处理失败: {str(e)}", [] def expand_text_tags(tags, provider, api_key, base_url, model_dropdown, model_text, nsfw_mode, nsfw_max_rolls, portrait_mode, portrait_suffix, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style): safe_api_key, safe_base_url, safe_model, safe_portrait_suffix, enabled_dims = _prepare_params( api_key, base_url, model_dropdown, model_text, portrait_suffix, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style ) try: result = expand_tags_sync( tags=tags, provider=provider, api_key=safe_api_key, base_url=safe_base_url, model=safe_model, nsfw=nsfw_mode, nsfw_max_rolls=int(nsfw_max_rolls), enabled_dims=enabled_dims, portrait=portrait_mode, portrait_suffix=safe_portrait_suffix ) return result except Exception as e: logger.error(f"扩写失败: {e}") return f"❌ 扩写失败: {str(e)}" def on_nsfw_toggle(nsfw_on): if nsfw_on: return gr.update(value=True), gr.update(visible=True) else: return gr.update(value=False), gr.update(visible=False) def create_gradio_app(): with gr.Blocks(title="🖼️ Image & Tag to Prompt") as app: gr.HTML(f"") gr.HTML("""

🖼️ Image / Tag → Prompt 智能提示词引擎

动态提取变量接口 · 标签扩写 · NSFW 防截断 Roll 系统 · 肖像数据集标注

""") # 动态合并预设的 5 个基础商和从变量里拉取出来的自定义供应商 PROVIDER_CHOICES = ["openai", "gemini", "claude", "qwen", "custom"] + list(config.EXTRA_PROVIDERS.keys()) with gr.Row(equal_height=False): # ===== 左栏：全局设置 ===== with gr.Column(scale=1, min_width=320): with gr.Accordion("⚙️ API 设置", open=False): provider = gr.Dropdown(choices=PROVIDER_CHOICES, value="openai", label="AI 提供商", info="如选中额外添加的供应商，下方配置可全留空") api_key = gr.Textbox(label="API Key", type="password", placeholder="留空→使用环境变量") base_url = gr.Textbox(label="API Base URL", placeholder="留空→使用环境变量") gr.Markdown("##### 模型选择") with gr.Row(): model_dropdown = gr.Dropdown(choices=[], value=None, label="从环境变量选择", allow_custom_value=True, scale=3) fetch_btn = gr.Button("🔍 获取模型", scale=1, elem_classes="fetch-btn") fetch_status = gr.Textbox(label="", interactive=False, max_lines=1, show_label=False) model_text = gr.Textbox(label="或手动输入", placeholder="如 gpt-4o（优先读取此处）") with gr.Group(elem_classes="mode-toggle-portrait"): gr.Markdown("### 📸 肖像标注模式") portrait_mode = gr.Checkbox(label="启用客观肖像标注 (训练 Lora 推荐)", value=False, info="覆盖下方维度开关，只输出客观短词。") portrait_suffix = gr.Textbox(label="自定义人物后缀【XX】", placeholder="例：章鱼 → 「一个年轻女性章鱼」") with gr.Group(elem_classes="nsfw-toggle"): gr.Markdown("### 🔓 NSFW 破限与维度控制") nsfw_mode = gr.Checkbox(label="🔞 启用 NSFW 无限制模式", value=False, info="多层破限 + 防截断 Roll 重试") nsfw_max_rolls = gr.Slider(minimum=0, maximum=10, value=3, step=1, label="🔄 触发截断自动重跑 (Roll) 次数", visible=False) with gr.Accordion("🎛️ 描述维度开关（点击展开）", open=False): gr.Markdown("*关闭维度会强硬禁止AI描述该项，【反推】和【扩写】均生效。*") with gr.Row(): with gr.Column(min_width=140): dim_appearance = gr.Checkbox(label="👤 人物外观", value=True) dim_body = gr.Checkbox(label="🫧 身体细节", value=True) dim_clothing = gr.Checkbox(label="👗 服装状态", value=True) dim_pose = gr.Checkbox(label="🤸 姿势动作", value=True) with gr.Column(min_width=140): dim_nsfw_detail = gr.Checkbox(label="🔞 性相关细节", value=False) dim_composition = gr.Checkbox(label="📐 构图镜头", value=True) dim_background = gr.Checkbox(label="🏞️ 背景环境", value=True) dim_style = gr.Checkbox(label="🎨 画风质量", value=True) # ===== 右栏：双功能 Tab ===== with gr.Column(scale=2): with gr.Tabs(): # --- Tab 1: 图像反推 --- with gr.TabItem("🖼️ 图像反推"): images_per_request = gr.Slider(minimum=1, maximum=15, value=5, step=1, label="单次请求合并图片数", info="建议维持在 3~5 以防漏图。") with gr.Accordion("📝 完全自定义模板 (点击展开填写)", open=False): custom_prompt = gr.Textbox(label="", lines=3, placeholder="填写后将完全覆盖上方所有模式和开关（仅在图像反推生效）") gr.HTML("
⚠️ 警告：填写此项将完全无视上方的模式和开关！
") file_input = gr.Files(label="上传图片或 .zip 压缩包", file_types=["image", ".zip"], file_count="multiple", elem_classes="upload-area") submit_btn = gr.Button("🚀 开始极速反推", variant="primary", size="lg", elem_classes="primary-btn") output_file = gr.File(label="📦 下载区域 (生成原图+TXT包 及 纯TXT轻量包)", file_count="multiple") with gr.Accordion("🖼️ 图文对照画廊预览 (点击展开)", open=False): output_gallery = gr.Gallery(label="处理结果预览", columns=4, rows=2, height="auto", object_fit="contain", preview=True) output_summary = gr.Textbox(label="所有图片描述合并结果", lines=15, interactive=False, elem_classes="summary-box") # --- Tab 2: 标签扩写 --- with gr.TabItem("✍️ 标签元素扩写"): gr.HTML("""
💡 扩写说明：输入简单的元素标签，AI将根据左侧的【维度开关】和【NSFW/肖像模式】为你扩写成细节饱满的中文自然语言长段落。
""") tags_input = gr.Textbox( label="输入基础标签或元素", lines=4, placeholder="例如：1girl, 森林, 魔法师, 夜晚...\n\n(AI 会自动补全环境、光影、服装、姿势等细节)" ) expand_btn = gr.Button("✨ 开始扩写提示词", variant="primary", size="lg", elem_classes="primary-btn") expand_output = gr.Textbox(label="扩写结果", lines=12, interactive=False, elem_classes="summary-box") # ===== 逻辑绑定 ===== nsfw_mode.change(fn=on_nsfw_toggle, inputs=[nsfw_mode], outputs=[dim_nsfw_detail, nsfw_max_rolls]) fetch_btn.click(fn=fetch_models_from_env, inputs=[], outputs=[model_dropdown, fetch_status]) submit_btn.click( fn=analyze_images, inputs=[ file_input, provider, api_key, base_url, model_dropdown, model_text, images_per_request, nsfw_mode, nsfw_max_rolls, portrait_mode, portrait_suffix, custom_prompt, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style ], outputs=[output_file, output_summary, output_gallery], ) expand_btn.click( fn=expand_text_tags, inputs=[ tags_input, provider, api_key, base_url, model_dropdown, model_text, nsfw_mode, nsfw_max_rolls, portrait_mode, portrait_suffix, dim_appearance, dim_body, dim_clothing, dim_pose, dim_nsfw_detail, dim_composition, dim_background, dim_style ], outputs=[expand_output] ) return app if __name__ == "__main__": app = create_gradio_app() auth_users = config.AUTH_USERS if auth_users: def auth_fn(username, password): for u, p in auth_users: if username == u and password == p: return True return False app.launch(server_name="0.0.0.0", server_port=7860, share=False, auth=auth_fn, auth_message="🔒 请输入账号密码登录后使用") else: app.launch(server_name="0.0.0.0", server_port=7860, share=False)