""" * @file run_pipeline.py * @brief AI海报生成系统主服务入口和API服务器 * 集成多个AI模型提供统一的海报生成接口 * * @author 王秀强 (2310460@mail.nankai.edu.cn) * @date 2025.6.9 * @version v1.0.0 * * @details * 本文件主要实现： * - FastAPI服务器和RESTful API接口 * - 用户输入分析和海报生成流程编排 * - Vue组件代码生成和PSD文件合成 * - 会话管理和文件下载服务 * - 集成DeepSeek、Kimi、ComfyUI等AI服务 * * @note * - 依赖外部ComfyUI服务(101.201.50.90:8188)进行图片生成 * - 需要配置DEEPSEEK_API_KEY和MOONSHOT_API_KEY环境变量 * - PSD生成优先使用手动创建的模板文件 * - 支持CORS跨域访问，生产环境需调整安全配置 * * @usage * # API服务器模式 * python run_pipeline.py * # 选择: 2 (API服务器模式) * * # 本地测试模式 * python run_pipeline.py * # 选择: 1 (本地测试模式) * * @copyright * (c) 2025 砚生项目组 */ """ import os from dotenv import load_dotenv import yaml from prompt_analysis import llm_user_analysis from generate_layout import call_deepseek, generate_vue_code, save_code from generate_text import load_config_from_file, get_poster_content_suggestions from fastapi import FastAPI from fastapi.responses import FileResponse, JSONResponse from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from flux_con import comfyui_img_info from export_psd_from_json import create_psd_from_images as create_psd_impl from colorama import init, Fore, Style import json import shutil import uuid from datetime import datetime # 初始化colorama init(autoreset=True) # 配置路径 config_paths = { "font": "../configs/font.yaml", "output_folder": "../outputs/", } # 加载环境变量和配置 load_dotenv() app = FastAPI(title="AI海报生成系统API", version="1.0.0") # 添加CORS中间件，允许前端跨域访问 app.add_middleware( CORSMiddleware, allow_origins=["*"], # 在生产环境中应该设置具体的域名 allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 请求模型 class PosterRequest(BaseModel): user_input: str session_id: str = None # 可选的会话ID，用于跟踪同一个用户的请求 class GenerateVueRequest(BaseModel): user_input: str session_id: str = None class GeneratePSDRequest(BaseModel): user_input: str session_id: str = None use_manual_psd: bool = False # 是否使用手动创建的PSD文件 # 响应模型 class ApiResponse(BaseModel): status: str message: str data: dict = None session_id: str = None # 全局变量存储会话数据 sessions = {} # 加载字体配置 try: with open(config_paths["font"], "r", encoding="utf-8") as f: fonts_config = yaml.safe_load(f) print(f"{Fore.GREEN}✅ 字体配置加载成功{Style.RESET_ALL}") except Exception as e: print(f"{Fore.YELLOW}⚠️ 字体配置加载失败: {e}，使用默认配置{Style.RESET_ALL}") fonts_config = {} # 辅助函数 def print_step(step_num, description, status="进行中"): """打印带颜色的步骤信息""" if status == "进行中": print(f"{Fore.BLUE}📋 步骤{step_num}: {description}...{Style.RESET_ALL}") elif status == "完成": print(f"{Fore.GREEN}✅ 步骤{step_num}: {description} - 完成{Style.RESET_ALL}") elif status == "错误": print(f"{Fore.RED}❌ 步骤{step_num}: {description} - 出错{Style.RESET_ALL}") def print_result(key, value): """打印结果信息""" print(f"{Fore.CYAN}📊 {key}: {value}{Style.RESET_ALL}") def get_session_folder(session_id): """获取会话专用的输出文件夹""" if not session_id: session_id = str(uuid.uuid4()) session_folder = os.path.join(config_paths["output_folder"], f"session_{session_id}") os.makedirs(session_folder, exist_ok=True) return session_folder, session_id # 生成prompts.yaml的函数 def generate_prompts_yaml(user_input=None): """ 动态生成prompts.yaml配置文件 """ if not user_input: user_input = "端午节海报，包含背景、活动亮点和图标" prompts_config = { "default_logo_text": "", "available_fonts": [ { "name": "Microsoft YaHei", "displayName": "微软雅黑", "tags": ["现代", "清晰"], "roles": ["title", "subtitle", "content"] }, { "name": "SimHei", "displayName": "黑体", "tags": ["通用", "标准"], "roles": ["title", "subtitle", "content"] } ], "NAMING_COLORS": { "primary": "#1976D2", "secondary": "#424242", "accent": "#FF5722" }, "STYLE_RULES": { "modern": { "primary_font": "Microsoft YaHei", "secondary_font": "SimHei" } }, "LOGO_RULES": { "default_position": "bottom", "fallback_text": "活动主办方" } } # 保存到临时文件 temp_prompts_path = os.path.join(config_paths["output_folder"], "temp_prompts.yaml") os.makedirs(os.path.dirname(temp_prompts_path), exist_ok=True) with open(temp_prompts_path, 'w', encoding='utf-8') as f: yaml.dump(prompts_config, f, allow_unicode=True, default_flow_style=False) print(f"{Fore.GREEN}✅ 临时prompts.yaml已生成: {temp_prompts_path}{Style.RESET_ALL}") return temp_prompts_path # 修复PSD合成接口 - 使用临时固定图片列表 def create_psd_from_images_wrapper(img_list, vue_layout_path, output_path): """ 临时接口：使用固定的图片文件创建PSD文件 使用outputs目录下的: lotus.jpg, nankai.png, stamp.jpg, background.png """ print(f"{Fore.CYAN}🎨 开始合成PSD文件(临时接口)...{Style.RESET_ALL}") try: # 使用固定的图片文件列表 fixed_image_files = ["background.png", "lotus.jpg", "nankai.png", "stamp.jpg"] image_paths = [] for img_file in fixed_image_files: img_path = os.path.join(config_paths["output_folder"], img_file) if os.path.exists(img_path): image_paths.append(img_path) print(f"{Fore.GREEN}✓ 找到图片: {img_file}{Style.RESET_ALL}") else: print(f"{Fore.YELLOW}⚠️ 图片不存在: {img_file} (路径: {img_path}){Style.RESET_ALL}") if not image_paths: print(f"{Fore.RED}❌ 没有找到任何指定的图片文件{Style.RESET_ALL}") return None print(f"{Fore.CYAN}📋 将合并以下图片 (共{len(image_paths)}张):") for i, path in enumerate(image_paths): print(f" {i + 1}. {os.path.basename(path)}") # 确保输出目录存在 os.makedirs(os.path.dirname(output_path), exist_ok=True) # 调用export_psd中的PSD创建函数 create_psd_impl( image_paths=image_paths, output_path=output_path, canvas_size=(1080, 1920), # 使用标准海报尺寸 mode='RGB' ) print(f"{Fore.GREEN}✅ PSD文件创建成功: {output_path}{Style.RESET_ALL}") # 验证PSD文件是否成功创建 if os.path.exists(output_path): file_size = os.path.getsize(output_path) / (1024 * 1024) # 转换为MB print(f"{Fore.CYAN}📁 PSD文件大小: {file_size:.2f} MB{Style.RESET_ALL}") return output_path except Exception as e: print(f"{Fore.RED}❌ PSD文件创建失败: {str(e)}{Style.RESET_ALL}") import traceback traceback.print_exc() return None # 增强Vue代码生成，确保包含文案内容 def generate_layout_prompt(user_input_analysis_result, parse_imglist, suggestions=None): """ 生成更完整的Vue布局提示，包含文案内容 """ width = user_input_analysis_result["width"] height = user_input_analysis_result["height"] theme = user_input_analysis_result.get("main_theme", "活动海报") # 构造图片信息字符串 images_info = "\n".join( [f"- {img['picture_name']} ({img['picture_description']})" for img in parse_imglist] ) # 构造文案信息 content_info = "" if suggestions: try: if 'layer6_title_content' in suggestions: title = suggestions['layer6_title_content'].get('content', theme) content_info += f"- 主标题: {title}\n" if 'layer7_subtitle_content' in suggestions: subtitle = suggestions['layer7_subtitle_content'].get('content', '精彩活动，敬请参与') content_info += f"- 副标题: {subtitle}\n" if 'layer5_logo_content' in suggestions: logo = suggestions['layer5_logo_content'].get('text', '主办方') content_info += f"- Logo文字: {logo}\n" except Exception as e: print(f"{Fore.YELLOW}⚠️ 文案信息解析错误: {e}{Style.RESET_ALL}") content_info = f"- 主标题: {theme}\n- 副标题: 精彩活动，敬请参与\n" # 调用DeepSeek生成动态排版Prompt system_prompt = "你是一个擅长前端开发的AI，专注于生成Vue.js代码。请根据提供的信息生成完整的Vue组件，包含所有必要的HTML结构和基础定位样式。" prompt = f""" 请生成一个Vue.js组件代码，用于{theme}海报，要求如下： 组件尺寸: {width}x{height}px 图片资源: {images_info} 文案内容: {content_info} 布局要求: 1. 背景图层: 使用第一张图片作为背景，占据整个组件区域 2. 主标题: 位于画布上方1/3处，居中显示 3. 副标题: 位于主标题下方，居中显示 4. 内容区域: 使用剩余图片，合理布局在中间区域 5. Logo区域: 位于底部，居中显示 技术要求: - 使用Vue 3 Composition API - 使用absolute定位进行精确布局 - 包含完整的template、script和style部分 - 确保所有文本内容都正确显示 - 图片使用相对路径引用 请生成完整可用的Vue组件代码，不要包含任何说明文字。 """ try: result, _ = call_deepseek(prompt=prompt, system_prompt=system_prompt, temperature=0.4) return result except Exception as e: print(f"{Fore.RED}❌ 布局提示生成失败: {e}{Style.RESET_ALL}") return generate_fallback_vue_code(theme, width, height) def generate_fallback_vue_code(theme, width=1080, height=1920): """ 生成备用的Vue代码 """ return f""" """ # 一键执行流程 def run_pipeline(user_input=None): """ 自动执行海报生成流程 """ try: print(f"{Fore.MAGENTA}{'=' * 50}") print(f"{Fore.MAGENTA}🎨 海报生成流程启动 🎨") print(f"{'=' * 50}{Style.RESET_ALL}") print_step(1, "加载配置文件") prompts_yaml_path = generate_prompts_yaml(user_input) load_config_from_file(prompts_yaml_path) print_step(1, "加载配置文件", "完成") print_step(2, "分析用户输入") user_input_analysis_result = llm_user_analysis(user_input) print_result("分析结果", user_input_analysis_result.get('main_theme', '未知')) print_step(2, "分析用户输入", "完成") print_step(3, "生成图片信息") system_prompt = user_input_analysis_result["analyzed_prompt"] parse_imglist = comfyui_img_info(user_input_analysis_result, system_prompt) print_result("生成图片数量", len(parse_imglist)) print_step(3, "生成图片信息", "完成") print_step(4, "生成文案建议") suggestions = get_poster_content_suggestions(user_input_analysis_result["analyzed_prompt"]) print(f"{Fore.CYAN}文案生成结果:") print(json.dumps(suggestions, indent=2, ensure_ascii=False)) # 保存文案到文件 suggestions_path = os.path.join(config_paths["output_folder"], "poster_content.json") with open(suggestions_path, "w", encoding="utf-8") as f: json.dump(suggestions, f, indent=2, ensure_ascii=False) print_step(4, "生成文案建议", "完成") print_step(5, "生成Vue排版") dynamic_prompt = generate_layout_prompt(user_input_analysis_result, parse_imglist, suggestions) vue_code = generate_vue_code(dynamic_prompt) vue_path = os.path.join(config_paths["output_folder"], "generated_code.vue") save_code(vue_code, file_path=vue_path) # 验证Vue文件是否成功生成 if os.path.exists(vue_path): print(f"{Fore.GREEN}✅ Vue文件已生成: {vue_path}{Style.RESET_ALL}") # 显示Vue代码的前几行用于验证 with open(vue_path, 'r', encoding='utf-8') as f: preview = f.read()[:500] print(f"{Fore.CYAN}Vue代码预览:\n{preview}...{Style.RESET_ALL}") else: print(f"{Fore.RED}❌ Vue文件生成失败{Style.RESET_ALL}") print_step(5, "生成Vue排版", "完成") print_step(6, "合成PSD文件") img_list = [(pic["picture_name"], pic["picture_type"]) for pic in parse_imglist] psd_path = os.path.join(config_paths["output_folder"], "final_poster.psd") result_path = create_psd_from_images_wrapper( img_list=img_list, vue_layout_path=vue_path, output_path=psd_path ) if result_path: print_step(6, "合成PSD文件", "完成") else: print_step(6, "合成PSD文件", "错误") print(f"\n{Fore.GREEN}{'=' * 50}") print(f"✅ 流程执行完成！") print(f"{'=' * 50}{Style.RESET_ALL}") return os.path.join(config_paths["output_folder"], "final_poster.png") except Exception as e: print_step("X", f"Pipeline执行", "错误") print(f"{Fore.RED}错误详情: {str(e)}{Style.RESET_ALL}") import traceback traceback.print_exc() return None # 本地运行函数 def run_local_pipeline(user_input=None): """ 本地运行整个管道流程，输出结果到控制台和文件系统。 """ print(f"{Fore.CYAN}🎬 Starting local pipeline with input: {Style.BRIGHT}{user_input}{Style.RESET_ALL}") output_path = run_pipeline(user_input) if output_path: print(f"{Fore.GREEN}🎊 Pipeline completed successfully!") print(f"{Fore.YELLOW}📁 Results saved to:") print(f" - Vue layout: {os.path.join(config_paths['output_folder'], 'generated_code.vue')}") print(f" - PSD file: {os.path.join(config_paths['output_folder'], 'final_poster.psd')}") print(f" - Content JSON: {os.path.join(config_paths['output_folder'], 'poster_content.json')}") print(f"{Fore.CYAN}💡 Check the outputs/ directory for generated files.{Style.RESET_ALL}") else: print(f"{Fore.RED}❌ Pipeline执行失败{Style.RESET_ALL}") # API路由 @app.get("/") def read_root(): return { "message": "AI海报生成系统API", "version": "1.0.0", "endpoints": { "generate_poster": "/api/generate-poster", # 主要接口 "download": "/api/download/{file_type}", "health": "/health", "status": "/api/status/{session_id}" } } @app.get("/health") def health_check(): """健康检查接口""" return {"status": "healthy", "timestamp": datetime.now().isoformat()} @app.post("/api/generate-poster", response_model=ApiResponse) async def generate_poster_api(request: PosterRequest): """ 一键生成完整海报（包含Vue代码和PSD文件）的主要API接口 """ try: session_folder, session_id = get_session_folder(request.session_id) print(f"{Fore.BLUE}🎨 开始生成海报...{Style.RESET_ALL}") print(f"{Fore.CYAN}用户输入: {request.user_input}{Style.RESET_ALL}") # === 步骤1: 生成配置文件 === print(f"{Fore.BLUE}📋 步骤1: 生成配置文件{Style.RESET_ALL}") temp_prompts_path = os.path.join(session_folder, "temp_prompts.yaml") prompts_config = { "default_logo_text": "", "available_fonts": [ { "name": "Microsoft YaHei", "displayName": "微软雅黑", "tags": ["现代", "清晰"], "roles": ["title", "subtitle", "content"] }, { "name": "SimHei", "displayName": "黑体", "tags": ["通用", "标准"], "roles": ["title", "subtitle", "content"] } ] } with open(temp_prompts_path, 'w', encoding='utf-8') as f: yaml.dump(prompts_config, f, allow_unicode=True, default_flow_style=False) load_config_from_file(temp_prompts_path) # === 步骤2: 分析用户输入 === print(f"{Fore.BLUE}📋 步骤2: 分析用户输入{Style.RESET_ALL}") user_input_analysis_result = llm_user_analysis(request.user_input) # === 步骤3: 生成图片信息 === print(f"{Fore.BLUE}📋 步骤3: 生成图片信息{Style.RESET_ALL}") system_prompt = user_input_analysis_result["analyzed_prompt"] parse_imglist = comfyui_img_info(user_input_analysis_result, system_prompt) # === 步骤4: 生成文案建议 === print(f"{Fore.BLUE}📋 步骤4: 生成文案建议{Style.RESET_ALL}") suggestions = get_poster_content_suggestions(user_input_analysis_result["analyzed_prompt"]) # 保存文案到会话文件夹 suggestions_path = os.path.join(session_folder, "poster_content.json") with open(suggestions_path, "w", encoding="utf-8") as f: json.dump(suggestions, f, indent=2, ensure_ascii=False) # === 步骤5: 生成Vue排版 === print(f"{Fore.BLUE}📋 步骤5: 生成Vue排版{Style.RESET_ALL}") dynamic_prompt = generate_layout_prompt(user_input_analysis_result, parse_imglist, suggestions) vue_code = generate_vue_code(dynamic_prompt) vue_path = os.path.join(session_folder, "generated_code.vue") save_code(vue_code, file_path=vue_path) # === 步骤6: 合成PSD文件 === print(f"{Fore.BLUE}📋 步骤6: 合成PSD文件{Style.RESET_ALL}") img_list = [(pic["picture_name"], pic["picture_type"]) for pic in parse_imglist] psd_path = os.path.join(session_folder, "final_poster.psd") # 修复PSD创建调用 try: # 使用固定的图片文件列表 fixed_image_files = ["background.png", "lotus.jpg", "nankai.png", "stamp.jpg"] image_paths = [] for img_file in fixed_image_files: img_path = os.path.join(config_paths["output_folder"], img_file) if os.path.exists(img_path): image_paths.append(img_path) if image_paths: # 确保输出目录存在 os.makedirs(os.path.dirname(psd_path), exist_ok=True) # 调用PSD创建函数 create_psd_impl( image_paths=image_paths, output_path=psd_path, canvas_size=(1080, 1920), mode='RGB' ) print(f"{Fore.GREEN}✅ PSD文件创建成功: {psd_path}{Style.RESET_ALL}") psd_created = True else: print(f"{Fore.YELLOW}⚠️ 没有找到图片文件，跳过PSD创建{Style.RESET_ALL}") psd_created = False except Exception as psd_error: print(f"{Fore.RED}❌ PSD文件创建失败: {str(psd_error)}{Style.RESET_ALL}") psd_created = False # 返回API响应 return ApiResponse( status="success", message="海报生成完成", data={ "vue_file": vue_path if os.path.exists(vue_path) else None, "psd_file": psd_path if psd_created else None, "content_file": suggestions_path, "analysis_result": user_input_analysis_result, "images_info": parse_imglist, "suggestions": suggestions }, session_id=session_id ) except Exception as e: print(f"{Fore.RED}❌ API错误: {str(e)}{Style.RESET_ALL}") import traceback traceback.print_exc() return ApiResponse( status="error", message=f"生成失败: {str(e)}", data=None, session_id=session_id if 'session_id' in locals() else None ) @app.get("/api/download/{file_type}") async def download_file(file_type: str, session_id: str = None): """ 下载生成的文件 """ try: if session_id: session_folder = os.path.join(config_paths["output_folder"], f"session_{session_id}") else: session_folder = config_paths["output_folder"] if file_type == "vue": file_path = os.path.join(session_folder, "generated_code.vue") media_type = "text/plain" elif file_type == "psd": file_path = os.path.join(session_folder, "final_poster.psd") media_type = "application/octet-stream" elif file_type == "json": file_path = os.path.join(session_folder, "poster_content.json") media_type = "application/json" else: return JSONResponse( status_code=400, content={"error": "不支持的文件类型"} ) if os.path.exists(file_path): return FileResponse( path=file_path, media_type=media_type, filename=os.path.basename(file_path) ) else: return JSONResponse( status_code=404, content={"error": "文件不存在"} ) except Exception as e: return JSONResponse( status_code=500, content={"error": f"下载失败: {str(e)}"} ) @app.get("/api/status/{session_id}") async def get_session_status(session_id: str): """ 获取会话状态 """ try: session_folder = os.path.join(config_paths["output_folder"], f"session_{session_id}") if not os.path.exists(session_folder): return JSONResponse( status_code=404, content={"error": "会话不存在"} ) files_status = { "vue_file": os.path.exists(os.path.join(session_folder, "generated_code.vue")), "psd_file": os.path.exists(os.path.join(session_folder, "final_poster.psd")), "content_file": os.path.exists(os.path.join(session_folder, "poster_content.json")) } return { "session_id": session_id, "files": files_status, "folder": session_folder } except Exception as e: return JSONResponse( status_code=500, content={"error": f"状态查询失败: {str(e)}"} ) if __name__ == "__main__": import uvicorn # 启动本地运行（可选） print(f"{Fore.BLUE}🔧 运行模式选择:{Style.RESET_ALL}") print(f"{Fore.YELLOW}1. 本地测试模式") print(f"2. API服务器模式{Style.RESET_ALL}") choice = input("请选择运行模式 (1/2): ").strip() if choice == "1": # 本地测试 test_input = input("请输入海报需求 (留空使用默认): ").strip() if not test_input: test_input = "端午节安康，" run_local_pipeline(test_input) else: # 启动API服务器 print(f"{Fore.GREEN}🚀 启动API服务器...{Style.RESET_ALL}") uvicorn.run(app, host="0.0.0.0", port=8000)