older FM support

This commit is contained in:
Karl Hudgell 2024-12-16 13:43:58 +00:00
parent c5a1867097
commit 39d6db68a3
3 changed files with 77 additions and 56 deletions

View File

@ -23,7 +23,7 @@ from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
logging.config.dictConfig(LOGGING_CONFIG) logging.config.dictConfig(LOGGING_CONFIG)
cut = 100 cut = 10
update = False update = False
use_gpu = False use_gpu = False
@ -41,13 +41,11 @@ rtf = RTF_Parser()
p = inflect.engine() 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.""" """Generate an image using the Comfy API."""
url = user_config["general"]["url"]
try: try:
# Initialize API and workflow # Initialize API and workflow
api = ComfyApiWrapper(url) api = ComfyApiWrapper(user_config["comfyui"]["comfyui_url"])
wf = ComfyWorkflowWrapper("./workflow_api.json") wf = ComfyWorkflowWrapper("./workflow_api.json")
# Set workflow parameters # 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("KSampler", "steps", steps)
wf.set_node_param("positive", "text", comfy_prompt) wf.set_node_param("positive", "text", comfy_prompt)
wf.set_node_param("Save Image", "filename_prefix", uid) 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 # Queue your workflow for completion
logging.debug(f"Generating image for UID: {uid}") logging.debug(f"Generating image for UID: {uid}")
results = api.queue_and_wait_images(wf, "Save Image") results = api.queue_and_wait_images(wf, "Save Image")
@ -112,7 +111,7 @@ def generate_prompts_for_players(players, app_config):
return prompts 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. 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 # Update or create configuration XML
if update: 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.") logging.debug("Configuration XML updated.")
else: else:
create_config_xml(output_folder) create_config_xml(output_folder, football_manager_version)
logging.debug("Configuration XML created.") logging.debug("Configuration XML created.")
except Exception as e: except Exception as e:
logging.error(f"Post-processing failed: {e}") logging.error(f"Post-processing failed: {e}")
@ -206,13 +205,11 @@ def main():
generate_image( generate_image(
uid, uid,
comfy_prompt, comfy_prompt,
user_config["general"]["model"],
args.num_inference_steps, args.num_inference_steps,
) )
try: try:
post_process_images( 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: except Exception as e:
logging.error(f"Post-processing failed: {e}") logging.error(f"Post-processing failed: {e}")

View File

@ -1,4 +1,7 @@
[general] [general]
url = football_manager_version = 2024
model = realisticVisionV60B1_v51HyperVAE.safetensors
output_dir = ./generated_images/ output_dir = ./generated_images/
[comfyui]
comfyui_url =
model = realisticVisionV60B1_v51HyperVAE.safetensors

View File

@ -1,63 +1,84 @@
import os import os
import xml.etree.ElementTree as ET 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 <list id="maps"> 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="preload", value="false")
ET.SubElement(root, "boolean", id="amap", value="false") ET.SubElement(root, "boolean", id="amap", value="false")
# Add the maps section
maps = ET.SubElement(root, "list", id="maps") 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): for filename in os.listdir(folder_path):
if filename.endswith(".png"): if filename.endswith(".png"):
# Extract the UID from the filename (remove the extension)
uid = os.path.splitext(filename)[0] uid = os.path.splitext(filename)[0]
add_uid_to_maps(maps, uid, football_manager_version)
# 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)
# Save the XML file # Save the XML file
output_path = os.path.join(folder_path, "config.xml") output_path = os.path.join(folder_path, "config.xml")
tree = ET.ElementTree(root)
tree.write(output_path, encoding="utf-8", xml_declaration=True) tree.write(output_path, encoding="utf-8", xml_declaration=True)
print(f"Config XML created at: {output_path}") print(f"Config XML created at: {output_path}")
def append_to_config_xml(file_path, uid_list): def append_to_config_xml(folder_path, uid_list, football_manager_version):
# Check if the file exists """
file_path = f"{file_path}/config.xml" Append new UIDs to an existing config.xml file.
if not os.path.exists(file_path): """
raise FileNotFoundError(f"The file {file_path} does not exist.") 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 # Get the maps list
tree = ET.parse(file_path) maps = get_or_create_maps_list(root)
root = tree.getroot()
# Find the <list id="maps"> element
maps = root.find(".//list[@id='maps']")
if maps is None:
raise ValueError("The XML file does not contain a <list id='maps'> element.")
# Add new UIDs # Add new UIDs
for uid in uid_list: for uid in uid_list:
# Check if the UID already exists (escaping 'from' in XPath) add_uid_to_maps(maps, uid, football_manager_version)
existing = None
for record in maps.findall("record"):
if record.attrib.get('from') == str(uid):
existing = record
break
if existing is None: # Save the updated XML file
# 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"})
# Write the updated XML back to the file
tree.write(file_path, encoding="utf-8", xml_declaration=True) tree.write(file_path, encoding="utf-8", xml_declaration=True)
print(f"Appended new UIDs to {file_path}") print(f"Appended new UIDs to {file_path}")