added logging

.gitignore

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

comfy_fm_newgen.py

@@ -7,30 +7,38 @@ 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"]
 
+    try:
         # Initialize API and workflow
         api = ComfyApiWrapper(url)
         wf = ComfyWorkflowWrapper("./workflow_api.json")
@@ -42,11 +50,15 @@ def generate_image(uid, comfy_prompt, model, steps):
         wf.set_node_param("Save Image", "filename_prefix", uid)
         wf.set_node_param("Load Checkpoint", "ckpt_name", model)
 
-    # queue your workflow for completion
+        # Queue your workflow for completion
+        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:
                 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,7 +75,8 @@ 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:
+            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])
@@ -81,9 +95,11 @@ def generate_prompts_for_players(players, app_config):
                 facial_characteristics=facial_characteristics or "no facial hair",
                 hair=f"{hair_length} {hair_colour} {hair}",
             )
-        print(f"Generated prompt: {prompt}")
+            logging.debug(f"Generated prompt: {prompt}")
             prompt = f"{player[0]}:{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
+    try:
         hair_length = list(set(item[5] 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
+    try:
         with open("config.json", "r") as f:
             app_config = json.load(f)
-    print(f"{len(rtf_file)} images will be generated")
+        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,
-            comfy_prompt,
-            user_config["general"]["model"],
-            args.num_inference_steps
+            uid, comfy_prompt, user_config["general"]["model"], args.num_inference_steps
         )
+
+    try:
         remove_bg_from_files_in_dir(output_folder)
         resize_images(output_folder)
         create_config_xml(output_folder)
-    print("\nImage generation complete for players in RTF file.")
+        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"],
+    },
+}