import os import time from datetime import datetime import argparse import sys import traceback 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_single_content import random TEXT_POSBILITY = 0.3 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": 5, "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")) ai_agent.close() # 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)): tweet_content_list = [] for j in range(config_file["variants"]): time.sleep(random.random()) ai_agent = AI_Agent(config_file["api_url"], config_file["model"], config_file["api_key"]) tweet_content, gen_result = generate_single_content( ai_agent, content_system_prompt, tweet_topic_record.topics_list[i], config_file["prompts_dir"], config_file["resource_dir"], output_dir, run_id, i+1, j+1, config_file["content_temperature"] ) tweet_content_list.append(tweet_content.get_json_file()) ai_agent.close() if not tweet_content: print(f"生成第{i+1}篇文章的第{j+1}个变体失败，跳过") continue object_name = tweet_topic_record.topics_list[i]["object"] try: object_name = object_name.split(".")[0] except: pass try: object_name = object_name.split("景点信息-")[1] except: pass # 处理对象名称中可能包含的"+"等特殊字符 # if "+" in object_name: # # 取第一个景点名称作为主要对象 # object_name = object_name.split("+")[0].strip() # print(f"对象名称包含多个景点，使用第一个景点：{object_name}") # 检查图片路径是否存在 img_dir_path = f"/root/autodl-tmp/sanming_img/modify/{object_name}" if not os.path.exists(img_dir_path): print(f"图片目录不存在：{img_dir_path}，跳过该对象") continue info_directory = [ f"/root/autodl-tmp/sanming_img/相机/{object_name}/description.txt" ] # 检查描述文件是否存在 if not os.path.exists(info_directory[0]): print(f"描述文件不存在：{info_directory[0]}，使用生成的内容替代") info_directory = [] poster_num = config_file["variants"] input_dir = img_dir_path # 使用前面检查过的目录路径 target_size = (900, 1200) result_path = [] content_gen = contentGen.ContentGenerator() response = content_gen.run(info_directory, poster_num, tweet_content_list) print(response) poster_config_summary = posterGen.PosterConfig(response) for j_index in range(config_file["variants"]): poster_config = poster_config_summary.get_config_by_index(j_index) img_dir = os.path.join(output_dir, f"{i+1}_{j_index+1}") try: # 创建输出目录 collage_output_dir = os.path.join(img_dir, "collage_img") os.makedirs(collage_output_dir, exist_ok=True) poster_output_dir = os.path.join(img_dir, "poster") os.makedirs(poster_output_dir, exist_ok=True) # 处理图片目录 img_list = simple_collage.process_directory( input_dir, target_size=target_size, output_count=1, output_dir=collage_output_dir ) print(img_list) if not img_list or len(img_list) == 0: print(f"未能生成拼贴图片，跳过海报生成") continue # 生成海报 poster_gen = posterGen.PosterGenerator() if random.random() < TEXT_POSBILITY: text_data = { "title": f"{poster_config['main_title']}", "subtitle": "", "additional_texts": [ {"text": f"{poster_config['texts'][0]}", "position": "bottom", "size_factor": 0.5}, {"text": f"{poster_config['texts'][1]}", "position": "bottom", "size_factor": 0.5} ] } else: text_data = { "title": f"{poster_config['main_title']}", "subtitle": "", "additional_texts": [ {"text": f"{poster_config['texts'][0]}", "position": "bottom", "size_factor": 0.5}, # {"text": f"{poster_config['texts'][1]}", "position": "bottom", "size_factor": 0.5} ] } print(text_data) img_path = img_list[0]['path'] print(f"使用图片路径: {img_path}") output_path = os.path.join(poster_output_dir, f"poster.jpg") result_path.append(poster_gen.create_poster(img_path, text_data, output_path)) except Exception as e: print(f"海报生成过程中出错: {e}") traceback.print_exc() # 打印完整的堆栈跟踪信息 continue if __name__ == "__main__": main()