Add prompt logging option

Author: ionite34

src/inference_core_nodes/prompt_expansion/prompt_expansion.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import random
+import logging
 from functools import cache
 from pathlib import Path
 
@@ -16,6 +17,8 @@
 
 CONFIGS_DIR = Path(__file__).parent.joinpath("configs")
 
+logger = logging.getLogger(__name__)
+
 fooocus_magic_split = [", extremely", ", intricate,"]
 
 disallowed_chars_table = str.maketrans("", "", "[]【】()（）|:：")
@@ -110,6 +113,7 @@ def INPUT_TYPES(s):
                 "model_name": (folder_paths.get_filename_list("prompt_expansion"),),
                 "text": ("STRING", {"multiline": True}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
+                "log_prompt": ("BOOLEAN", {"default": False})
             },
         }
 
@@ -127,7 +131,7 @@ def INPUT_TYPES(s):
 
     @staticmethod
     @torch.no_grad()
-    def expand_prompt(model_name: str, text: str, seed: int, log_prompt: str):
+    def expand_prompt(model_name: str, text: str, seed: int, log_prompt: bool):
         expansion = load_expansion_runner(model_name)
 
         prompt = remove_empty_str([safe_str(text)], default="")[0]
@@ -142,6 +146,10 @@ def expand_prompt(model_name: str, text: str, seed: int, log_prompt: str):
 
         expansion_text = expansion(prompt, seed)
         expanded_prompt = join_prompts(prompt, expansion_text)
+        
+        if log_prompt:
+            logger.info(f"Prompt: {prompt}")
+            logger.info(f"Expanded Prompt: {expanded_prompt}")
 
         return expanded_prompt, seed