import os import time from datetime import datetime import argparse import sys import traceback import json from core.ai_agent import AI_Agent # from core.topic_parser import TopicParser # No longer needed directly in main? 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 generate_single_content, run_topic_generation_pipeline # Import the new pipeline function from utils.prompt_manager import PromptManager # Import PromptManager import random TEXT_POSBILITY = 0.3 # Consider moving this to config if it varies def load_config(config_path="poster_gen_config.json"): """Loads configuration from a JSON file.""" if not os.path.exists(config_path): print(f"Error: Configuration file '{config_path}' not found.") print("Please copy 'example_config.json' to 'poster_gen_config.json' and customize it.") sys.exit(1) try: with open(config_path, 'r', encoding='utf-8') as f: config = json.load(f) # Basic validation (can be expanded) required_keys = ["api_url", "model", "api_key", "resource_dir", "prompts_dir", "output_dir", "num", "variants", "topic_system_prompt", "topic_user_prompt", "content_system_prompt", "image_base_dir"] if not all(key in config for key in required_keys): missing_keys = [key for key in required_keys if key not in config] print(f"Error: Config file '{config_path}' is missing required keys: {missing_keys}") sys.exit(1) # Resolve relative paths based on config location or a defined base path if necessary # For simplicity, assuming paths in config are relative to project root or absolute return config except json.JSONDecodeError: print(f"Error: Could not decode JSON from '{config_path}'. Check the file format.") sys.exit(1) except Exception as e: print(f"Error loading configuration from '{config_path}': {e}") sys.exit(1) # Removed generate_topics_step function definition from here # Its logic is now in utils.tweet_generator.run_topic_generation_pipeline def generate_content_and_posters_step(config, run_id, tweet_topic_record): """Generates content and posters based on generated topics.""" if not run_id or not tweet_topic_record: print("Missing run_id or topics data. Skipping content and poster generation.") return print("\nStep 2: Generating Content and Posters...") base_output_dir = config["output_dir"] output_dir = os.path.join(base_output_dir, run_id) # Directory for this specific run # --- Pre-load resources and initialize shared objects --- # Initialize PromptManager once for this phase try: prompt_manager = PromptManager(config) except Exception as e: print(f"Error initializing PromptManager: {e}") traceback.print_exc() return # Load content generation system prompt once content_system_prompt = ResourceLoader.load_system_prompt(config["content_system_prompt"]) if not content_system_prompt: print("Warning: Content generation system prompt is empty. Using default logic if available or might fail.") # Initialize AI Agent once for the entire content generation phase ai_agent = None try: print(f"Initializing AI Agent ({config['model']})...") ai_agent = AI_Agent(config["api_url"], config["model"], config["api_key"]) except Exception as e: print(f"Error initializing AI Agent: {e}. Cannot proceed with content generation.") traceback.print_exc() return # Cannot continue without AI agent # Check image base directory image_base_dir = config.get("image_base_dir", None) if not image_base_dir or not os.path.isdir(image_base_dir): print(f"Error: 'image_base_dir' ({image_base_dir}) not specified or not a valid directory in config. Cannot locate images.") if ai_agent: ai_agent.close() # Close agent if initialized return camera_image_subdir = config.get("camera_image_subdir", "相机") # Default '相机' modify_image_subdir = config.get("modify_image_subdir", "modify") # Default 'modify' # Initialize ContentGenerator and PosterGenerator once if they are stateless # Assuming they are stateless for now content_gen = contentGen.ContentGenerator() poster_gen_instance = posterGen.PosterGenerator() # --- Process each topic --- for i, topic in enumerate(tweet_topic_record.topics_list): topic_index = i + 1 print(f"\nProcessing Topic {topic_index}/{len(tweet_topic_record.topics_list)}: {topic.get('title', 'N/A')}") tweet_content_list = [] # --- Content Generation Loop (using the single AI Agent) --- for j in range(config["variants"]): variant_index = j + 1 print(f" Generating Variant {variant_index}/{config['variants']}...") # Get prompts for this specific topic item using PromptManager content_system_prompt, content_user_prompt = prompt_manager.get_content_prompts(topic) if not content_system_prompt or not content_user_prompt: print(f" Skipping Variant {variant_index} due to missing content prompts.") continue # Skip this variant if prompts failed time.sleep(random.random() * 0.5) # Slightly reduced delay try: # Use the pre-initialized AI Agent and pass the specific prompts tweet_content, gen_result = generate_single_content( ai_agent, content_system_prompt, content_user_prompt, topic, output_dir, run_id, topic_index, variant_index, config.get("content_temperature", 0.3) ) if tweet_content: # Assuming get_json_file() returns a dictionary or similar structure tweet_content_data = tweet_content.get_json_file() if tweet_content_data: tweet_content_list.append(tweet_content_data) else: print(f" Warning: generate_single_content for Topic {topic_index}, Variant {variant_index} returned empty data.") else: print(f" Failed to generate content for Topic {topic_index}, Variant {variant_index}. Skipping.") except Exception as e: print(f" Error during content generation for Topic {topic_index}, Variant {variant_index}: {e}") # Decide if traceback is needed here, might be too verbose for loop errors # traceback.print_exc() # Do NOT close the agent here if not tweet_content_list: print(f" No valid content generated for Topic {topic_index}. Skipping poster generation.") continue # --- Poster Generation Setup --- object_name = topic.get("object", "") if not object_name: print(" Warning: Topic object name is missing. Cannot determine image path.") continue # Clean object name (consider making this a utility function) try: # More robust cleaning might be needed depending on actual object name formats object_name_cleaned = object_name.split(".")[0].replace("景点信息-", "").strip() if not object_name_cleaned: print(f" Warning: Object name '{object_name}' resulted in empty string after cleaning.") continue object_name = object_name_cleaned except Exception as e: print(f" Warning: Could not fully clean object name '{object_name}': {e}") # Continue with potentially unclean name? Or skip? # Let's continue for now, path checks below might catch issues. # Construct and check image paths using config base dir input_img_dir_path = os.path.join(image_base_dir, modify_image_subdir, object_name) camera_img_dir_path = os.path.join(image_base_dir, camera_image_subdir, object_name) description_file_path = os.path.join(camera_img_dir_path, "description.txt") if not os.path.exists(input_img_dir_path) or not os.path.isdir(input_img_dir_path): print(f" Image directory not found or not a directory: '{input_img_dir_path}'. Skipping poster generation for this topic.") continue info_directory = [] if os.path.exists(description_file_path): info_directory = [description_file_path] print(f" Using description file: {description_file_path}") else: print(f" Description file not found: '{description_file_path}'. Using generated content for poster text.") # --- Generate Text Configurations for Posters --- try: # Pass the list of content data directly poster_text_configs_raw = content_gen.run(info_directory, config["variants"], tweet_content_list) # print(f" Raw poster text configs: {poster_text_configs_raw}") # For debugging if not poster_text_configs_raw: print(" Warning: ContentGenerator returned empty configuration data.") continue # Skip if no text configs generated poster_config_summary = posterGen.PosterConfig(poster_text_configs_raw) except Exception as e: print(f" Error running ContentGenerator or parsing poster configs: {e}") traceback.print_exc() continue # Skip poster generation for this topic # --- Poster Generation Loop --- poster_num = config["variants"] # Same as content variants target_size = tuple(config.get("poster_target_size", [900, 1200])) for j_index in range(poster_num): variant_index = j_index + 1 print(f" Generating Poster {variant_index}/{poster_num}...") try: poster_config = poster_config_summary.get_config_by_index(j_index) if not poster_config: print(f" Warning: Could not get poster config for index {j_index}. Skipping.") continue variant_output_dir = os.path.join(output_dir, f"{topic_index}_{variant_index}") collage_output_dir = os.path.join(variant_output_dir, "collage_img") poster_output_dir = os.path.join(variant_output_dir, "poster") os.makedirs(collage_output_dir, exist_ok=True) os.makedirs(poster_output_dir, exist_ok=True) # --- Image Collage --- img_list = simple_collage.process_directory( input_img_dir_path, target_size=target_size, output_count=1, # Assuming 1 collage image per poster variant output_dir=collage_output_dir ) if not img_list or len(img_list) == 0 or not img_list[0].get('path'): print(f" Failed to generate collage image for Variant {variant_index}. Skipping poster.") continue collage_img_path = img_list[0]['path'] print(f" Using collage image: {collage_img_path}") # --- Create Poster (using the single poster_gen_instance) --- text_data = { "title": poster_config.get('main_title', 'Default Title'), "subtitle": "", "additional_texts": [] } texts = poster_config.get('texts', []) if texts: text_data["additional_texts"].append({"text": texts[0], "position": "bottom", "size_factor": 0.5}) if len(texts) > 1 and random.random() < TEXT_POSBILITY: text_data["additional_texts"].append({"text": texts[1], "position": "bottom", "size_factor": 0.5}) final_poster_path = os.path.join(poster_output_dir, "poster.jpg") result_path = poster_gen_instance.create_poster(collage_img_path, text_data, final_poster_path) if result_path: print(f" Successfully generated poster: {result_path}") else: print(f" Poster generation function did not return a valid path.") except Exception as e: print(f" Error during poster generation for Variant {variant_index}: {e}") traceback.print_exc() continue # Continue to next variant # --- Cleanup --- # Close the AI Agent after processing all topics if ai_agent: print("\nClosing AI Agent...") ai_agent.close() def main(): # No argparse for now, directly load default config config = load_config() # Load from poster_gen_config.json # Execute steps sequentially # Step 1: Generate Topics (using the function from utils.tweet_generator) run_id, tweet_topic_record = run_topic_generation_pipeline(config) # Step 2: Generate Content and Posters (if Step 1 was successful) if run_id and tweet_topic_record: generate_content_and_posters_step(config, run_id, tweet_topic_record) else: print("Exiting due to issues in topic generation.") if __name__ == "__main__": main()