diff --git a/.gitignore b/.gitignore index 1fed40b..87f5263 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,8 @@ venv/* -config.cfg generated_*/* script.log **/*.pyc *.rtf build/ -dist/ \ No newline at end of file +dist/ +user_config.cfg diff --git a/.vscode/launch.json b/.vscode/launch.json index 7146044..be5e826 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -11,12 +11,14 @@ "program": "${file}", "console": "integratedTerminal", "justMyCode": false, - "args": [ - "--num_inference_steps", - "10", - "--rtf_file", - "./NewGen.rtf" - ] + // "args": [ + // "--num_inference_steps", + // "6", + // "--rtf_file", + // "./AllInDB.rtf", + // "--player_uuid", + // "2000252303" + // ] } ] } \ No newline at end of file diff --git a/app_config.json b/app_config.json index 4379b6d..f223d2e 100644 --- a/app_config.json +++ b/app_config.json @@ -35,11 +35,6 @@ null ], "hair": [ - "short", - "long", - "bald", - "buzz cut", - "medium-length", "curly", "wavy", "spiky", @@ -83,8 +78,8 @@ }, "hair_length": { "0": "Bald", - "1": "Short", - "2": "Medium", + "1": "Really Short", + "2": "Medium Length", "3": "Long", "4": "Bald" }, diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index ff1b6ec..5e9d3e0 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -16,6 +16,12 @@ from lib.remove_bg import remove_bg_from_file_list 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.general import ( + choose_profile, + create_or_update, + process_player_or_file, + get_player_input, +) from lib.logging import LOGGING_CONFIG # from simple_term_menu import TerminalMenu @@ -26,12 +32,15 @@ logging.config.dictConfig(LOGGING_CONFIG) cut = None update = False use_gpu = False +process_player = False # Load user configurations user_config = configparser.ConfigParser() try: - user_config.read("user_config.cfg") - output_folder = user_config["general"]["output_dir"] + user_config.read("./user_config.cfg") + selected_profile = choose_profile("./user_config.cfg") + selected_profile = f"profile:{selected_profile}" + output_folder = user_config[selected_profile]["output_dir"] logging.debug("Configuration loaded successfully.") except KeyError as e: logging.error(f"Missing configuration key: {e}") @@ -41,7 +50,7 @@ rtf = RTF_Parser() p = inflect.engine() -def generate_image(uid, comfy_prompt, steps): +def generate_image(uid, comfy_prompt): """Generate an image using the Comfy API.""" try: # Initialize API and workflow @@ -50,7 +59,7 @@ def generate_image(uid, comfy_prompt, steps): # Set workflow parameters wf.set_node_param("KSampler", "seed", random.getrandbits(32)) - wf.set_node_param("KSampler", "steps", steps) + # 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( @@ -60,7 +69,9 @@ def generate_image(uid, comfy_prompt, steps): 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: + with open( + f"{user_config[selected_profile]['output_dir']}{uid}.png", "wb+" + ) as f: f.write(image_data) logging.debug(f"Image generated successfully for UID: {uid}") except Exception as e: @@ -93,7 +104,10 @@ def generate_prompts_for_players(players, app_config): 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"]) + if int(player[5]) > 1: + hair_extra = random.choice(app_config["hair"]) + else: + hair_extra = "" # Format the prompt prompt = app_config["prompt"].format( @@ -101,7 +115,7 @@ def generate_prompts_for_players(players, app_config): age=player_age, country=country, facial_characteristics=facial_characteristics or "no facial hair", - hair=f"{hair_length} {hair_colour}", + hair=f"{hair_length} {hair_colour} {hair_extra}", ) logging.debug(f"Generated prompt: {prompt}") prompt = f"{player[0]}:{prompt}" @@ -111,7 +125,9 @@ def generate_prompts_for_players(players, app_config): return prompts -def post_process_images(output_folder, update, processed_players, football_manager_version): +def post_process_images( + output_folder, update, processed_players, football_manager_version +): """ Handles post-processing tasks for generated images. @@ -126,15 +142,19 @@ def post_process_images(output_folder, update, processed_players, football_manag logging.debug("Images resized successfully.") # Remove background from images using GPU if available - remove_bg_from_file_list(output_folder,processed_players, use_gpu=use_gpu) + remove_bg_from_file_list(output_folder, processed_players, 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) + append_to_config_xml( + output_folder, processed_players, football_manager_version + ) logging.debug("Configuration XML updated.") else: - create_config_xml(output_folder,processed_players, football_manager_version) + create_config_xml( + output_folder, processed_players, football_manager_version + ) logging.debug("Configuration XML created.") except Exception as e: logging.error(f"Post-processing failed: {e}") @@ -143,32 +163,30 @@ def post_process_images(output_folder, update, processed_players, football_manag 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() + # 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( + # "--player_uuid", + # type=int, + # default=None, + # help="Player UUID to generate", + # ) + # 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) + # if not args.rtf_file: + # logging.error("Please pass in a RTF file as --rtf_file") + # sys.exit(1) # Load configurations try: @@ -179,17 +197,43 @@ def main(): logging.error("app_config.json file not found.") sys.exit(1) + # Parse the RTF file + try: + # rtf_file = random.sample(rtf.parse_rtf(args.rtf_file), cut) + rtf_location = user_config[selected_profile]["rtf_file"] + rtf_file = rtf.parse_rtf(rtf_location) + logging.info(f"Parsed RTF file successfully. Found {len(rtf_file)} players.") + except FileNotFoundError: + logging.error(f"RTF file not found: {rtf_location}") + sys.exit(1) + + update = create_or_update() + process_player = process_player_or_file() + if process_player: + player_uuid = get_player_input() # Check for processed try: if update: values_from_config = extract_from_values( - f"{user_config['general']['output_dir']}config.xml" + f"{user_config[selected_profile]['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] + if process_player: + players_to_process = [ + inner_list + for inner_list in players_to_process + if int(inner_list[0]) == player_uuid + ] + elif process_player: + players_to_process = [ + inner_list + for inner_list in rtf_file + if int(inner_list[0]) == player_uuid + ] else: players_to_process = rtf_file except FileNotFoundError: @@ -202,14 +246,13 @@ def main(): 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, - ) + generate_image(uid, comfy_prompt) try: post_process_images( - output_folder, update, [item[0] for item in players_to_process],user_config["general"]["football_manager_version"] + output_folder, + update, + [item[0] for item in players_to_process], + user_config[selected_profile]["football_manager_version"], ) except Exception as e: logging.error(f"Post-processing failed: {e}") diff --git a/lib/general.py b/lib/general.py new file mode 100644 index 0000000..00a0e23 --- /dev/null +++ b/lib/general.py @@ -0,0 +1,80 @@ +import configparser + +def list_profiles(cfg_file): + config = configparser.ConfigParser() + config.read(cfg_file) + + profiles = [section.split(':', 1)[1] for section in config.sections() if section.startswith('profile:')] + return profiles + + +def choose_profile(cfg_file): + profiles = list_profiles(cfg_file) + if not profiles: + print("No profiles found.") + return None + + print("Available Profiles:") + for i, profile in enumerate(profiles, 1): + print(f"{i}. {profile}") + + while True: + try: + choice = int(input("Enter the number of the profile you want to use: ")) + if 1 <= choice <= len(profiles): + return profiles[choice - 1] + else: + print("Invalid number. Please try again.") + except ValueError: + print("Invalid input. Please enter a number.") + +def create_or_update(): + values = ["Create", "Update"] + print("Create or Update:") + for i, profile in enumerate(values, 1): + print(f"{i}. {profile}") + + while True: + try: + choice = int(input("Enter the number of the choice: ")) + if 1 <= choice <= len(values): + if choice == 1: + return False + else: + return True + else: + print("Invalid number. Please try again.") + except ValueError: + print("Invalid input. Please enter a number.") + + +def process_player_or_file(): + values = ["Whole File", "Specific Player"] + print("Process whole rtf file or a specific player?:") + for i, profile in enumerate(values, 1): + print(f"{i}. {profile}") + + while True: + try: + choice = int(input("Enter the number of the choice: ")) + if 1 <= choice <= len(values): + if choice == 1: + return False + else: + return True + else: + print("Invalid number. Please try again.") + except ValueError: + print("Invalid input. Please enter a number.") + + +def get_player_input(): + # Prompt the user to enter a UUID + player_uuid = input("Please enter the player UUID: ").strip() + + # Validate that the UUID is not empty + while not player_uuid: + print("Player UUID cannot be empty. Please try again.") + player_uuid = input("Please enter the player UUID: ").strip() + + return int(player_uuid) \ No newline at end of file diff --git a/user_config.cfg.sample b/user_config.cfg.sample index 127c9d4..e16a7c5 100644 --- a/user_config.cfg.sample +++ b/user_config.cfg.sample @@ -1,7 +1,8 @@ -[general] -football_manager_version = 2024 -output_dir = ./generated_images/ - [comfyui] comfyui_url = -model = realisticVisionV60B1_v51HyperVAE.safetensors \ No newline at end of file +model = realisticVisionV60B1_v51HyperVAE.safetensors + +[profile:NewGens] +football_manager_version = 2024 +rtf_file = ./NewGen.rtf +output_dir = ./NewGens/ \ No newline at end of file diff --git a/workflow_api.json b/workflow_api.json index 58c3ca9..e96d874 100644 --- a/workflow_api.json +++ b/workflow_api.json @@ -2,8 +2,8 @@ "3": { "inputs": { "seed": 531046367, - "steps": 10, - "cfg": 1.0, + "steps": 6, + "cfg": 1.5, "sampler_name": "dpmpp_2m_sde", "scheduler": "karras", "denoise": 1,