added logging

.gitignore

@@ -1,3 +1,4 @@
 venv/*
 config.cfg
 generated_images/*
+script.log

comfy_fm_newgen.py

@@ -7,46 +7,58 @@ import pycountry
 import inflect
 import logging
 import sys
+import logging
+import logging.config
 
 from tqdm import tqdm
 from lib.rtf_parser import RTF_Parser
 from lib.remove_bg import remove_bg_from_files_in_dir
 from lib.generate_xml import create_config_xml
 from lib.resize_images import resize_images
+from lib.logging import LOGGING_CONFIG
 from comfy_api_simplified import ComfyApiWrapper, ComfyWorkflowWrapper
 
-# logging.basicConfig(stream=sys.stdout, level=logging.INFO)
+logging.config.dictConfig(LOGGING_CONFIG)
 
 # Load user configurations
 user_config = configparser.ConfigParser()
-user_config.read("config.cfg")
+try:
+    user_config.read("config.cfg")
+    output_folder = user_config["general"]["output_dir"]
+    logging.debug("Configuration loaded successfully.")
+except KeyError as e:
+    logging.error(f"Missing configuration key: {e}")
+    sys.exit(1)
 
 rtf = RTF_Parser()
 p = inflect.engine()
 
-output_folder = user_config["general"]["output_dir"]
-
 
 def generate_image(uid, comfy_prompt, model, steps):
     """Generate an image using the Comfy API."""
     url = user_config["general"]["url"]
 
-    # Initialize API and workflow
+    try:
-    api = ComfyApiWrapper(url)
+        # Initialize API and workflow
-    wf = ComfyWorkflowWrapper("./workflow_api.json")
+        api = ComfyApiWrapper(url)
+        wf = ComfyWorkflowWrapper("./workflow_api.json")
 
-    # Set workflow parameters
+        # Set workflow parameters
-    wf.set_node_param("KSampler", "seed", random.getrandbits(32))
+        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("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", model)
 
-    # queue your workflow for completion
+        # Queue your workflow for completion
-    results = api.queue_and_wait_images(wf, "Save Image")
+        logging.debug(f"Generating image for UID: {uid}")
-    for filename, image_data in results.items():
+        results = api.queue_and_wait_images(wf, "Save Image")
-        with open(f"./generated_images/{uid}.png", "wb+") as f:
+        for filename, image_data in results.items():
-            f.write(image_data)
+            with open(f"./generated_images/{uid}.png", "wb+") as f:
+                f.write(image_data)
+        logging.debug(f"Image generated successfully for UID: {uid}")
+    except Exception as e:
+        logging.error(f"Failed to generate image for UID: {uid}. Error: {e}")
 
 
 def get_country_name_from_code(code):
@@ -55,6 +67,7 @@ def get_country_name_from_code(code):
         country = pycountry.countries.get(alpha_3=code.upper())
         return country.name if country else "Unknown country code"
     except KeyError:
+        logging.warning(f"Invalid country code provided: {code}")
         return "Invalid country code"
 
 
@@ -62,28 +75,31 @@ def generate_prompts_for_players(players, app_config):
     """Generate images for a specific player and configuration."""
     prompts = []
     for player in players:
-        print(f"\nGenerating image for {player[0]} - {player[8]}")
+        try:
-        os.makedirs(output_folder, exist_ok=True)
+            logging.debug(f"Generating prompt for {player[0]} - {player[8]}")
+            os.makedirs(output_folder, exist_ok=True)
 
-        country = get_country_name_from_code(player[1])
+            country = get_country_name_from_code(player[1])
-        facial_characteristics = random.choice(app_config["facial_characteristics"])
+            facial_characteristics = random.choice(app_config["facial_characteristics"])
-        hair_length = app_config["hair_length"][player[5]]
+            hair_length = app_config["hair_length"][player[5]]
-        hair_colour = app_config["hair_color"][player[6]]
+            hair_colour = app_config["hair_color"][player[6]]
-        skin_tone = app_config["skin_tone_map"][player[7]]
+            skin_tone = app_config["skin_tone_map"][player[7]]
-        player_age = p.number_to_words(player[3])
+            player_age = p.number_to_words(player[3])
-        hair = random.choice(app_config["hair"])
+            hair = random.choice(app_config["hair"])
 
-        # Format the prompt
+            # Format the prompt
-        prompt = app_config["prompt"].format(
+            prompt = app_config["prompt"].format(
-            skin_tone=skin_tone,
+                skin_tone=skin_tone,
-            age=player_age,
+                age=player_age,
-            country=country,
+                country=country,
-            facial_characteristics=facial_characteristics or "no facial hair",
+                facial_characteristics=facial_characteristics or "no facial hair",
-            hair=f"{hair_length} {hair_colour} {hair}",
+                hair=f"{hair_length} {hair_colour} {hair}",
-        )
+            )
-        print(f"Generated prompt: {prompt}")
+            logging.debug(f"Generated prompt: {prompt}")
-        prompt = f"{player[0]}:{prompt}"
+            prompt = f"{player[0]}:{prompt}"
-        prompts.append(prompt)
+            prompts.append(prompt)
+        except KeyError as e:
+            logging.warning(f"Key error while generating prompt for player: {e}")
     return prompts
 
 
@@ -105,34 +121,52 @@ def main():
     args = parser.parse_args()
 
     if not args.rtf_file:
-        raise Exception("Please pass in a RTF file as --rtf_file")
+        logging.error("Please pass in a RTF file as --rtf_file")
+        sys.exit(1)
 
     # Parse the RTF file
-    rtf_file = rtf.parse_rtf(args.rtf_file)[:5]
+    try:
+        rtf_file = rtf.parse_rtf(args.rtf_file)[:10]
+        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)
 
     # Extract unique values
-    hair_length = list(set(item[5] for item in rtf_file))
+    try:
-    hair_colour = list(set(item[6] for item in rtf_file))
+        hair_length = list(set(item[5] for item in rtf_file))
-    skin_tone = list(set(item[7] for item in rtf_file))
+        hair_colour = list(set(item[6] for item in rtf_file))
+        skin_tone = list(set(item[7] for item in rtf_file))
+    except IndexError as e:
+        logging.error(f"Error extracting player attributes: {e}")
+        sys.exit(1)
 
     # Load configurations
-    with open("config.json", "r") as f:
+    try:
-        app_config = json.load(f)
+        with open("config.json", "r") as f:
-    print(f"{len(rtf_file)} images will be generated")
+            app_config = json.load(f)
+        logging.debug("Application configuration loaded successfully.")
+    except FileNotFoundError:
+        logging.error("config.json file not found.")
+        sys.exit(1)
+
     prompts = generate_prompts_for_players(rtf_file, app_config)
     for prompt in tqdm(prompts, desc="Generating Images"):
         uid = prompt.split(":")[0]
         comfy_prompt = prompt.split(":")[1]
         generate_image(
-            uid,
+            uid, comfy_prompt, user_config["general"]["model"], args.num_inference_steps
-            comfy_prompt,
-            user_config["general"]["model"],
-            args.num_inference_steps
         )
-    remove_bg_from_files_in_dir(output_folder)
+
-    resize_images(output_folder)
+    try:
-    create_config_xml(output_folder)
+        remove_bg_from_files_in_dir(output_folder)
-    print("\nImage generation complete for players in RTF file.")
+        resize_images(output_folder)
+        create_config_xml(output_folder)
+        logging.debug("Post-processing complete.")
+    except Exception as e:
+        logging.error(f"Post-processing failed: {e}")
+
+    logging.info("Image generation complete for players in RTF file.")
 
 
 if __name__ == "__main__":

lib/logging.py

@@ -0,0 +1,31 @@
+LOGGING_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "file": {
+            "format": "%(asctime)s - %(levelname)s - %(message)s"
+        },
+        "console": {
+            "format": "%(levelname)s - %(message)s"
+        },
+    },
+    "handlers": {
+        "file_handler": {
+            "class": "logging.FileHandler",
+            "level": "INFO",
+            "formatter": "file",
+            "filename": "script.log",
+            "mode": "a",  # Append to the log file
+        },
+        "console_handler": {
+            "class": "logging.StreamHandler",
+            "level": "WARNING",
+            "formatter": "console",
+            "stream": "ext://sys.stdout",
+        },
+    },
+    "root": {
+        "level": "DEBUG",  # Root logger captures all logs
+        "handlers": ["file_handler", "console_handler"],
+    },
+}