latest rework with profiles

.gitignore

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

.vscode/launch.json

@@ -11,12 +11,14 @@
             "program": "${file}",
             "console": "integratedTerminal",
             "justMyCode": false,
-            "args": [
+            // "args": [
-                "--num_inference_steps",
+            //     "--num_inference_steps",
-                "10",
+            //     "6",
-                "--rtf_file",
+            //     "--rtf_file",
-                "./NewGen.rtf"
+            //     "./AllInDB.rtf",
-            ]
+            //     "--player_uuid",
+            //     "2000252303"
+            // ]
         }
     ]
 }

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",
+    "1": "Really Short",
-    "2": "Medium",
+    "2": "Medium Length",
     "3": "Long",
     "4": "Bald"
   },

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")
+    user_config.read("./user_config.cfg")
-    output_folder = user_config["general"]["output_dir"]
+    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.
 
@@ -131,10 +147,14 @@ def post_process_images(output_folder, update, processed_players, football_manag
 
         # 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 = argparse.ArgumentParser(description="Generate images for country groups")
-    parser.add_argument(
+    # parser.add_argument(
-        "--rtf_file",
+    #     "--rtf_file",
-        type=str,
+    #     type=str,
-        default=None,
+    #     default=None,
-        help="Path to the RTF file to be processed",
+    #     help="Path to the RTF file to be processed",
-    )
+    # )
-    parser.add_argument(
+    # parser.add_argument(
-        "--num_inference_steps",
+    #     "--player_uuid",
-        type=int,
+    #     type=int,
-        default=6,
+    #     default=None,
-        help="Number of inference steps. Defaults to 6",
+    #     help="Player UUID to generate",
-    )
+    # )
-    args = parser.parse_args()
+    # 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:
+    # if not args.rtf_file:
-        logging.error("Please pass in a RTF file as --rtf_file")
+    #     logging.error("Please pass in a RTF file as --rtf_file")
-        sys.exit(1)
+    #     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:
@@ -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(
+            generate_image(uid, comfy_prompt)
-                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"]
+                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}")

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)

user_config.cfg.sample

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

workflow_api.json

@@ -2,8 +2,8 @@
   "3": {
     "inputs": {
       "seed": 531046367,
-      "steps": 10,
+      "steps": 6,
-      "cfg": 1.0,
+      "cfg": 1.5,
       "sampler_name": "dpmpp_2m_sde",
       "scheduler": "karras",
       "denoise": 1,