import argparse import os import random import json import configparser import pycountry import inflect import logging import sys import logging import logging.config from tqdm import tqdm from lib.rtf_parser import RTF_Parser from lib.remove_bg import remove_bg_from_files_in_dir from lib.generate_xml import create_config_xml, append_to_config_xml from lib.resize_images import resize_images from lib.xml_reader import extract_from_values from lib.logging import LOGGING_CONFIG # from simple_term_menu import TerminalMenu from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper logging.config.dictConfig(LOGGING_CONFIG) cut = 10 update = False use_gpu = False # Load user configurations user_config = configparser.ConfigParser() try: user_config.read("config.cfg") output_folder = user_config["general"]["output_dir"] logging.debug("Configuration loaded successfully.") except KeyError as e: logging.error(f"Missing configuration key: {e}") sys.exit(1) rtf = RTF_Parser() p = inflect.engine() def generate_image(uid, comfy_prompt, steps): """Generate an image using the Comfy API.""" try: # Initialize API and workflow api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"]) wf = ComfyWorkflowWrapper("./workflow_api.json") # Set workflow parameters wf.set_node_param("KSampler", "seed", random.getrandbits(32)) wf.set_node_param("KSampler", "steps", steps) wf.set_node_param("positive", "text", comfy_prompt) wf.set_node_param("Save Image", "filename_prefix", uid) wf.set_node_param( "Load Checkpoint", "ckpt_name", user_config["comfyui"]["model"] ) # Queue your workflow for completion logging.debug(f"Generating image for UID: {uid}") results = api.queue_and_wait_images(wf, "Save Image") for filename, image_data in results.items(): with open(f"./generated_images/{uid}.png", "wb+") as f: f.write(image_data) logging.debug(f"Image generated successfully for UID: {uid}") except Exception as e: logging.error(f"Failed to generate image for UID: {uid}. Error: {e}") def get_country_name(app_config, country_code): # First check if it's a custom mapping if country_code in app_config["facial_characteristics"]: return app_config["facial_characteristics"][country_code] # Use pycountry for standard country codes country = pycountry.countries.get(alpha_3=country_code) if country: return country.name return "Unknown Country" def generate_prompts_for_players(players, app_config): """Generate images for a specific player and configuration.""" prompts = [] for player in players: try: logging.debug(f"Generating prompt for {player[0]} - {player[8]}") os.makedirs(output_folder, exist_ok=True) country = get_country_name(app_config, player[1]) facial_characteristics = random.choice(app_config["facial_characteristics"]) hair_length = app_config["hair_length"][player[5]] hair_colour = app_config["hair_color"][player[6]] skin_tone = app_config["skin_tone_map"][player[7]] player_age = p.number_to_words(player[3]) hair = random.choice(app_config["hair"]) # Format the prompt prompt = app_config["prompt"].format( skin_tone=skin_tone, age=player_age, country=country, facial_characteristics=facial_characteristics or "no facial hair", hair=f"{hair_length} {hair_colour}", ) logging.debug(f"Generated prompt: {prompt}") prompt = f"{player[0]}:{prompt}" prompts.append(prompt) except KeyError as e: logging.warning(f"Key error while generating prompt for player: {e}") return prompts def post_process_images(output_folder, update, processed_players, football_manager_version): """ Handles post-processing tasks for generated images. Args: output_folder (str): Path to the folder where images are stored. update (bool): Flag to determine if XML config should be updated. processed_players (list): List of processed player IDs. """ try: # Resize images to desired dimensions resize_images(output_folder) logging.debug("Images resized successfully.") # Remove background from images using GPU if available remove_bg_from_files_in_dir(output_folder, use_gpu=use_gpu) logging.debug("Background removed from images.") # Update or create configuration XML if update: append_to_config_xml(output_folder, processed_players, football_manager_version) logging.debug("Configuration XML updated.") else: create_config_xml(output_folder, football_manager_version) logging.debug("Configuration XML created.") except Exception as e: logging.error(f"Post-processing failed: {e}") raise # Re-raise the exception to ensure the script stops if post-processing fails. def main(): """Main function for generating images.""" parser = argparse.ArgumentParser(description="Generate images for country groups") parser.add_argument( "--rtf_file", type=str, default=None, help="Path to the RTF file to be processed", ) parser.add_argument( "--num_inference_steps", type=int, default=6, help="Number of inference steps. Defaults to 6", ) args = parser.parse_args() if not args.rtf_file: logging.error("Please pass in a RTF file as --rtf_file") sys.exit(1) # Parse the RTF file try: rtf_file = random.sample(rtf.parse_rtf(args.rtf_file), cut) logging.info(f"Parsed RTF file successfully. Found {len(rtf_file)} players.") except FileNotFoundError: logging.error(f"RTF file not found: {args.rtf_file}") sys.exit(1) # Load configurations try: with open("config.json", "r") as f: app_config = json.load(f) logging.debug("Application configuration loaded successfully.") except FileNotFoundError: logging.error("config.json file not found.") sys.exit(1) # Check for processed try: if update: values_from_config = extract_from_values( f"{user_config['general']['output_dir']}config.xml" ) # Extract the IDs from list_a ids_in_b = [item for item in values_from_config] # Filter list_a to remove inner lists whose first item matches an ID in list_b players_to_process = [item for item in rtf_file if item[0] not in ids_in_b] else: players_to_process = rtf_file except FileNotFoundError: logging.error("config.json file not found.") sys.exit(1) if len(players_to_process) > 0: print(f"Processing {len(players_to_process)} players") logging.info(f"Processing {len(players_to_process)} players") prompts = generate_prompts_for_players(players_to_process, app_config) for prompt in tqdm(prompts, desc="Generating Images"): uid = prompt.split(":")[0] comfy_prompt = prompt.split(":")[1] generate_image( uid, comfy_prompt, args.num_inference_steps, ) try: post_process_images( output_folder, update, [item[0] for item in players_to_process],user_config["general"]["football_manager_version"] ) except Exception as e: logging.error(f"Post-processing failed: {e}") else: print(f"{len(rtf_file)} players already processed") logging.info(f"{len(rtf_file)} players already processed") logging.info("Image generation complete for players in RTF file.") if __name__ == "__main__": main()