ai_service/scripts/run_pipeline.py

import os
from dotenv import load_dotenv
import yaml
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, HTTPException
from fastapi.responses import FileResponse
from ComfyUI.flux_con import comfyui_img_info  # 导入已实现的图像生成函数

# 配置路径
config_paths = {
    "font": "../configs/font.yaml",
    "output_folder": "../outputs/",
}

# 加载环境变量和配置
load_dotenv()
app = FastAPI()

with open(config_paths["font"], "r", encoding="utf-8") as f:
    fonts_config = yaml.safe_load(f)


# 假设的 DeepSeek 提示词分析接口
def llm_user_analysis(user_input):
    # 占位符实现，模拟 DeepSeek 分析
    if not user_input:
        user_input = "端午节海报，包含背景和活动亮点"
    return {
        "analyzed_prompt": user_input,
        "keywords": ["端午节", "背景", "活动亮点"],
        "width": 1080,
        "height": 1920,
        "batch_size": 2
    }


# 假设的 PSD 合成接口（占位符）
def create_psd_from_images(img_list, vue_layout_path, output_path):
    print(f"Generating PSD from {img_list} using layout {vue_layout_path} to {output_path}")
    return None


# 动态生成 prompts.yaml 基于用户输入
def generate_prompts_yaml(user_input):
    prompts_data = {
        "user_prompt": user_input,
        "generated_at": "02:50 PM HKT on Monday, June 09, 2025",
        "description": f"Generated prompt based on user input: {user_input}"
    }
    prompts_yaml_path = os.path.join(config_paths["output_folder"], "prompts.yaml")
    with open(prompts_yaml_path, "w", encoding="utf-8") as f:
        yaml.dump(prompts_data, f)
    return prompts_yaml_path


# 生成动态Vue排版Prompt，支持多张图片
def generate_layout_prompt(user_input_analysis_result, parse_imglist):
    width = user_input_analysis_result["width"]
    height = user_input_analysis_result["height"]

    # 构造图片信息字符串
    images_info = "\n".join(
        [f"- {img['picture_name']} ({img['picture_description']})" for img in parse_imglist]
    )

    # 调用DeepSeek生成动态排版Prompt
    system_prompt = "你是一个擅长前端开发的AI，专注于生成Vue.js代码。给定组件尺寸和多张图片信息，自动分析并建议合理的排版位置，使用absolute定位，仅关注位置，不包含样式描述。"
    prompt = (
        f"根据组件尺寸 {width}x{height}px 和以下图片信息，生成一个Vue组件代码的Prompt，包含以下部分并指定动态排版位置：\n"
        f"图片信息：\n{images_info}\n"
        f"要求：\n"
        f"1. 背景图层：使用第一张图片，占据整个组件区域。"
        f"2. 主体图层：包含标题和副标题，位于合适位置，居中。"
        f"3. 活动亮点：使用剩余图片，位于合适位置，居中，使用网格或分层布局展示所有图片。"
        f"4. 页脚：包含主办单位信息和logo图片，位于合适位置，居中。"
        f"组件尺寸为{width}x{height}px，布局使用absolute定位，仅关注排版位置。"
        f"返回一个格式化的Prompt字符串，供生成Vue代码使用。"
    )

    result, _ = call_deepseek(prompt=prompt, system_prompt=system_prompt, temperature=0.4)
    return result


# 一键执行流程
def run_pipeline(user_input=None):
    """
    自动执行海报生成流程：
    1. 使用 llm_user_analysis 分析用户提示词得到 user_input_analysis_result。
    2. 使用 comfyui_img_info 生成 parse_imglist 列表。
    3. 根据 user_input_analysis_result 生成文案 suggestions。
    4. 结合 parse_imglist 和动态生成的 prompts.yaml 生成 Vue 排版文件。
    5. 将 parse_imglist 组装成 img_list，调用 create_psd_from_images 合成 PSD。
    :param user_input: 用户提供的提示词（可选，默认使用占位符）
    :return: 合成海报的路径
    """
    # 步骤 1: 加载配置
    load_config_from_file(generate_prompts_yaml(user_input))  # 动态生成 prompts.yaml

    # 步骤 2: 分析用户提示词
    user_input_analysis_result = llm_user_analysis(user_input)
    print(f"Analyzed result: {user_input_analysis_result}")

    # 步骤 3: 生成图片信息
    system_prompt = user_input_analysis_result["analyzed_prompt"]
    parse_imglist = comfyui_img_info(user_input_analysis_result, system_prompt)
    print(f"Generated image list: {parse_imglist}")

    # 步骤 4: 生成文案
    suggestions = get_poster_content_suggestions(user_input_analysis_result)
    print(f"Generated suggestions: {suggestions}")

    # 步骤 5: 生成 Vue 排版（动态Prompt，支持多张图片）
    dynamic_prompt = generate_layout_prompt(user_input_analysis_result, parse_imglist)
    print(f"Generated dynamic layout prompt: {dynamic_prompt}")
    vue_code = generate_vue_code(dynamic_prompt)
    save_code(vue_code, file_path=os.path.join(config_paths["output_folder"], "generated_code.vue"))
    print("Generated Vue layout")

    # 步骤 6: 合成 PSD
    img_list = [(pic["picture_name"], pic["picture_type"]) for pic in parse_imglist]
    create_psd_from_images(
        img_list=img_list,
        vue_layout_path=os.path.join(config_paths["output_folder"], "generated_code.vue"),
        output_path=os.path.join(config_paths["output_folder"], "final_poster.psd")
    )
    print("Generated PSD file")

    return os.path.join(config_paths["output_folder"], "final_poster.png")  # 假设 PNG 作为最终输出


# 本地运行函数
def run_local_pipeline(user_input=None):
    """
    本地运行整个管道流程，输出结果到控制台和文件系统。
    :param user_input: 用户提供的提示词（可选，默认使用占位符）
    :return: None，打印运行结果
    """
    print(f"Starting local pipeline with input: {user_input}")
    output_path = run_pipeline(user_input)
    print(f"Pipeline completed. 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"- Final poster (placeholder): {output_path}")
    print("Check the outputs/ directory for generated files.")


# API 端点
@app.get("/generate_poster")
async def generate_poster(user_input: str = None):
    output_path = run_pipeline(user_input)
    return FileResponse(output_path, media_type="image/png")


# 可选 API 端点
@app.get("/generate_poster_psd")
async def generate_poster_psd(user_input: str = None):
    try:
        output_path = run_pipeline(user_input)
        return {
            "image_url": output_path,
            "font_info": fonts_config,
            "psd_path": os.path.join(config_paths["output_folder"], "final_poster.psd")
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"PSD generation failed: {str(e)}")


if __name__ == "__main__":
    import uvicorn

    # 启动本地运行（可选）
    run_local_pipeline("端午节海报，包含背景、活动亮点和图标")
    # 或者启动API服务器
    # uvicorn.run(app, host="0.0.0.0", port=8000)