TravelContentCreator/main.py

import os
import time
from datetime import datetime
import argparse
import sys

from core.ai_agent import AI_Agent
from core.topic_parser import TopicParser
import core.contentGen as contentGen
import core.posterGen as posterGen
import core.simple_collage as simple_collage
from utils.resource_loader import ResourceLoader
from utils.tweet_generator import prepare_topic_generation, generate_topics, generate_content

def main():
    config_file = {
        "date": "4月17日",
        "num": 5,
        "model": "qwenQWQ",
        "api_url": "vllm",
        "api_key": "EMPTY",
        "topic_system_prompt": "/root/autodl-tmp/TravelContentCreator/SelectPrompt/systemPrompt.txt",
        "topic_user_prompt": "/root/autodl-tmp/TravelContentCreator/SelectPrompt/userPrompt.txt",
        "content_system_prompt": "/root/autodl-tmp/TravelContentCreator/genPrompts/systemPrompt.txt",
        "resource_dir": [{
                            "type": "Object",
                            "num": 4,
                            "file_path": ["/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-尚书第.txt",
                                          "/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-明清园.txt",
                                          "/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-泰宁古城.txt",
                                          "/root/autodl-tmp/TravelContentCreator/resource/Object/景点信息-甘露寺.txt"
                            ]},
                            {
                                "type": "Product",
                                "num": 0,
                                "file_path": []
                            }
                         ],
        "prompts_dir": "/root/autodl-tmp/TravelContentCreator/genPrompts",
        "output_dir": "/root/autodl-tmp/TravelContentCreator/result",
        "variants": 2,
        "topic_temperature": 0.2,
        "content_temperature": 0.3
    }

    if True:
        # 1. 首先生成选题
        ai_agent, system_prompt, user_prompt, output_dir = prepare_topic_generation(
            config_file["date"], config_file["num"], config_file["topic_system_prompt"], config_file["topic_user_prompt"],
            config_file["api_url"], config_file["model"], config_file["api_key"], config_file["prompts_dir"], config_file["resource_dir"], config_file["output_dir"]
        )
        
        run_id, tweet_topic_record = generate_topics(
            ai_agent, system_prompt, user_prompt, config_file["output_dir"],
            config_file["topic_temperature"], 0.5, 1.5
        )
        
        output_dir = os.path.join(config_file["output_dir"], run_id)
        os.makedirs(output_dir, exist_ok=True)
        tweet_topic_record.save_topics(os.path.join(output_dir, "tweet_topic.json"))
        tweet_topic_record.save_prompt(os.path.join(output_dir, "tweet_prompt.txt"))
        # raise Exception("选题生成失败，退出程序")
        if not run_id or not tweet_topic_record:
            print("选题生成失败，退出程序")
            return
        
        # 2. 然后生成内容
        print("\n开始根据选题生成内容...")
        
        # 加载内容生成的系统提示词
        content_system_prompt = ResourceLoader.load_system_prompt(config_file["content_system_prompt"])

            
        if not content_system_prompt:
            print("内容生成系统提示词为空，使用选题生成的系统提示词")
            content_system_prompt = system_prompt
        
        # 直接使用同一个AI Agent实例
        for i in range(len(tweet_topic_record.topics_list)):
            result = generate_content(
                ai_agent, content_system_prompt, tweet_topic_record.topics_list[i], output_dir, run_id, config_file["prompts_dir"], config_file["resource_dir"],
                config_file["variants"], config_file["content_temperature"]
            )
            object_name = tweet_topic_record.topics_list[i].object
            info_directory = [
                f"/root/autodl-tmp/sanming_img/相机/{object_name}/description.txt"
            ]
            poster_num = 1
            tweet_content = f"""
            {result[0].get_result_json()}
            """
            input_dir = f"/root/autodl-tmp/sanming_img/modify/{object_name}"
            output_dir = "/root/autodl-tmp/poster_generate_result"
            target_size = (900, 1200)
            result_path = []
            
            
            content_gen = contentGen.ContentGenerator()
            response = content_gen.run(info_directory, poster_num, tweet_content)
            print(response)
            img_list = simple_collage.process_directory(input_dir, target_size=target_size, output_count=poster_num, output_dir=output_dir)
            print(img_list)
            poster_gen = posterGen.PosterGenerator()
            poster_config = posterGen.PosterConfig(response)
            for index, item in enumerate(poster_config.get_config()):
                text_data = {
                    "title": f"{item['main_title']}",
                    "subtitle": "",
                    "additional_texts": [
                        {"text": f"{item['texts'][0]}", "position": "bottom", "size_factor": 0.5},
                        {"text": f"{item['texts'][1]}", "position": "bottom", "size_factor": 0.5}
                    ]
                }
                img_path = img_list[index]['path']
                print(img_path)
                result_path.append(poster_gen.create_poster(img_path, text_data, f"{index}.jpg"))
        

if __name__ == "__main__":
    main()