latest rework with profiles

This commit is contained in:
Karl Hudgell 2024-12-21 14:08:10 +00:00
parent 2511ef1594
commit aa574f9c01
7 changed files with 186 additions and 65 deletions

4
.gitignore vendored
View File

@ -1,8 +1,8 @@
venv/*
config.cfg
generated_*/*
script.log
**/*.pyc
*.rtf
build/
dist/
dist/
user_config.cfg

14
.vscode/launch.json vendored
View File

@ -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"
// ]
}
]
}

View File

@ -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"
},

View File

@ -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}")

80
lib/general.py Normal file
View File

@ -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)

View File

@ -1,7 +1,8 @@
[general]
football_manager_version = 2024
output_dir = ./generated_images/
[comfyui]
comfyui_url =
model = realisticVisionV60B1_v51HyperVAE.safetensors
model = realisticVisionV60B1_v51HyperVAE.safetensors
[profile:NewGens]
football_manager_version = 2024
rtf_file = ./NewGen.rtf
output_dir = ./NewGens/

View File

@ -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,