diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index dbeccee..f00dbd2 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -23,7 +23,7 @@ from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper logging.config.dictConfig(LOGGING_CONFIG) -cut = 100 +cut = 10 update = False use_gpu = False @@ -41,13 +41,11 @@ rtf = RTF_Parser() p = inflect.engine() -def generate_image(uid, comfy_prompt, model, steps): +def generate_image(uid, comfy_prompt, steps): """Generate an image using the Comfy API.""" - url = user_config["general"]["url"] - try: # Initialize API and workflow - api = ComfyApiWrapper(url) + api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"]) wf = ComfyWorkflowWrapper("./workflow_api.json") # Set workflow parameters @@ -55,8 +53,9 @@ def generate_image(uid, comfy_prompt, model, 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("Load Checkpoint", "ckpt_name", model) - + 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") @@ -112,7 +111,7 @@ def generate_prompts_for_players(players, app_config): return prompts -def post_process_images(output_folder, update, processed_players): +def post_process_images(output_folder, update, processed_players, football_manager_version): """ Handles post-processing tasks for generated images. @@ -132,10 +131,10 @@ def post_process_images(output_folder, update, processed_players): # Update or create configuration XML if update: - append_to_config_xml(output_folder, processed_players) + append_to_config_xml(output_folder, processed_players, football_manager_version) logging.debug("Configuration XML updated.") else: - create_config_xml(output_folder) + create_config_xml(output_folder, football_manager_version) logging.debug("Configuration XML created.") except Exception as e: logging.error(f"Post-processing failed: {e}") @@ -165,7 +164,7 @@ def main(): # Parse the RTF file try: - rtf_file = random.sample(rtf.parse_rtf(args.rtf_file),cut) + 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}") @@ -206,13 +205,11 @@ def main(): generate_image( uid, comfy_prompt, - user_config["general"]["model"], args.num_inference_steps, ) - try: post_process_images( - output_folder, update, [item[0] for item in players_to_process] + 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}") diff --git a/config.cfg.sample b/config.cfg.sample index 0bf7e4e..127c9d4 100644 --- a/config.cfg.sample +++ b/config.cfg.sample @@ -1,4 +1,7 @@ [general] -url = -model = realisticVisionV60B1_v51HyperVAE.safetensors -output_dir = ./generated_images/ \ No newline at end of file +football_manager_version = 2024 +output_dir = ./generated_images/ + +[comfyui] +comfyui_url = +model = realisticVisionV60B1_v51HyperVAE.safetensors \ No newline at end of file diff --git a/lib/generate_xml.py b/lib/generate_xml.py index 8b21135..d4676c6 100644 --- a/lib/generate_xml.py +++ b/lib/generate_xml.py @@ -1,63 +1,84 @@ import os import xml.etree.ElementTree as ET -def create_config_xml(folder_path): - # Define the root element - root = ET.Element("record") - # Add resource manager options +def parse_or_create_xml(file_path, create_if_missing=True): + """ + Parse an existing XML file or create a new one if it doesn't exist. + """ + if os.path.exists(file_path): + tree = ET.parse(file_path) + root = tree.getroot() + elif create_if_missing: + # Create a new XML structure if file doesn't exist + root = ET.Element("record") + tree = ET.ElementTree(root) + else: + raise FileNotFoundError(f"The file {file_path} does not exist.") + return tree, root + + +def get_or_create_maps_list(root): + """ + Get the element or create it if it doesn't exist. + """ + maps = root.find(".//list[@id='maps']") + if maps is None: + maps = ET.SubElement(root, "list", id="maps") + return maps + + +def add_uid_to_maps(maps, uid, football_manager_version): + """ + Add a UID to the maps list if it doesn't already exist. + """ + if football_manager_version == "2024": + prefix = "r-" + else: + prefix = "" + for record in maps.findall("record"): + if record.attrib.get("from") == str(uid): + return # UID already exists + ET.SubElement(maps, "record", attrib={"from": str(uid), "to": f"graphics/pictures/person/{prefix}{uid}/portrait"}) + + +def create_config_xml(folder_path, football_manager_version): + """ + Create a new config.xml file based on .png files in the folder. + """ + # Prepare the XML structure + root = ET.Element("record") ET.SubElement(root, "boolean", id="preload", value="false") ET.SubElement(root, "boolean", id="amap", value="false") - - # Add the maps section maps = ET.SubElement(root, "list", id="maps") - # Iterate through files in the folder + # Add records for each .png file for filename in os.listdir(folder_path): if filename.endswith(".png"): - # Extract the UID from the filename (remove the extension) uid = os.path.splitext(filename)[0] - - # Create a record for each file - ET.SubElement(maps, "record", attrib={"from": uid, "to": f"graphics/pictures/person/r-{uid}/portrait"}) - - # Create the XML tree - tree = ET.ElementTree(root) + add_uid_to_maps(maps, uid, football_manager_version) # Save the XML file output_path = os.path.join(folder_path, "config.xml") + tree = ET.ElementTree(root) tree.write(output_path, encoding="utf-8", xml_declaration=True) print(f"Config XML created at: {output_path}") -def append_to_config_xml(file_path, uid_list): - # Check if the file exists - file_path = f"{file_path}/config.xml" - if not os.path.exists(file_path): - raise FileNotFoundError(f"The file {file_path} does not exist.") +def append_to_config_xml(folder_path, uid_list, football_manager_version): + """ + Append new UIDs to an existing config.xml file. + """ + file_path = os.path.join(folder_path, "config.xml") + tree, root = parse_or_create_xml(file_path, create_if_missing=False) - # Parse the existing XML file - tree = ET.parse(file_path) - root = tree.getroot() - - # Find the element - maps = root.find(".//list[@id='maps']") - if maps is None: - raise ValueError("The XML file does not contain a element.") + # Get the maps list + maps = get_or_create_maps_list(root) # Add new UIDs for uid in uid_list: - # Check if the UID already exists (escaping 'from' in XPath) - existing = None - for record in maps.findall("record"): - if record.attrib.get('from') == str(uid): - existing = record - break - - if existing is None: - # Add the new record if the UID doesn't exist - ET.SubElement(maps, "record", attrib={"from": str(uid), "to": f"graphics/pictures/person/r-{uid}/portrait"}) + add_uid_to_maps(maps, uid, football_manager_version) - # Write the updated XML back to the file + # Save the updated XML file tree.write(file_path, encoding="utf-8", xml_declaration=True) - print(f"Appended new UIDs to {file_path}") \ No newline at end of file + print(f"Appended new UIDs to {file_path}")