diff --git a/comfy_fm_newgen.py b/comfy_fm_newgen.py index c5a3958..da6322b 100644 --- a/comfy_fm_newgen.py +++ b/comfy_fm_newgen.py @@ -25,6 +25,49 @@ from lib.text_chunker import chunk_prompt_for_clip # Profile functions are now handled entirely by GUI from lib.logging import LOGGING_CONFIG + +def save_prompt_mapping(uid, prompt): + """Save the prompt used for a specific image UID""" + try: + prompt_file = f"{output_folder}/prompts.json" + + # Load existing prompts + if os.path.exists(prompt_file): + with open(prompt_file, 'r') as f: + prompts_data = json.load(f) + else: + prompts_data = {} + + # Add or update the prompt for this UID + prompts_data[uid] = prompt + + # Save back to file + with open(prompt_file, 'w') as f: + json.dump(prompts_data, f, indent=2) + + logging.debug(f"Saved prompt for UID {uid}") + + except Exception as e: + logging.warning(f"Failed to save prompt for UID {uid}: {e}") + + +def get_prompt_for_image(uid): + """Get the prompt used for a specific image UID""" + try: + prompt_file = f"{output_folder}/prompts.json" + + if os.path.exists(prompt_file): + with open(prompt_file, 'r') as f: + prompts_data = json.load(f) + + return prompts_data.get(uid, "Prompt not found") + + return "No prompts file found" + + except Exception as e: + logging.warning(f"Failed to load prompt for UID {uid}: {e}") + return "Error loading prompt" + # from simple_term_menu import TerminalMenu from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler import torch @@ -354,6 +397,9 @@ def main(): uid = prompt.split(":")[0] comfy_prompt = prompt.split(":")[1] generate_image(uid, comfy_prompt) + + # Save the prompt for this image + save_prompt_mapping(uid, comfy_prompt) try: post_process_images( output_folder, diff --git a/gui.py b/gui.py index 26191c9..37568f4 100644 --- a/gui.py +++ b/gui.py @@ -18,7 +18,7 @@ from PIL import Image, ImageTk import logging # Import the main generation logic -from comfy_fm_newgen import generate_image, get_country_name, generate_prompts_for_players +from comfy_fm_newgen import generate_image, get_country_name, generate_prompts_for_players, get_prompt_for_image try: from lib.rtf_parser import RTF_Parser except ImportError: @@ -613,6 +613,14 @@ class FMFaceGeneratorGUI: self.total_generated_var = tk.StringVar(value="0") ttk.Label(info_frame, textvariable=self.total_generated_var).grid(row=1, column=1, sticky=tk.W, padx=5) + # Prompt display + prompt_frame = ttk.LabelFrame(preview_frame, text="Prompt Used", padding=10) + prompt_frame.pack(fill=tk.X, padx=10, pady=5) + + self.prompt_text = scrolledtext.ScrolledText(prompt_frame, height=4, wrap=tk.WORD) + self.prompt_text.pack(fill=tk.X, pady=5) + self.prompt_text.config(state=tk.DISABLED) # Make it read-only + # Navigation buttons nav_frame = ttk.Frame(preview_frame) nav_frame.pack(fill=tk.X, pady=10) @@ -974,6 +982,20 @@ class FMFaceGeneratorGUI: self.image_label.config(image=photo, text="") self.image_label.image = photo # Keep reference self.current_image_var.set(f"Player {uid}") + + # Load and display the prompt for this image + try: + prompt = get_prompt_for_image(uid) + self.prompt_text.config(state=tk.NORMAL) + self.prompt_text.delete(1.0, tk.END) + self.prompt_text.insert(tk.END, prompt) + self.prompt_text.config(state=tk.DISABLED) + except Exception as e: + self.prompt_text.config(state=tk.NORMAL) + self.prompt_text.delete(1.0, tk.END) + self.prompt_text.insert(tk.END, f"Error loading prompt: {str(e)}") + self.prompt_text.config(state=tk.DISABLED) + self.log_message(f"Preview updated: {uid}.png") except Exception as e: @@ -1036,6 +1058,8 @@ class FMFaceGeneratorGUI: if existing_images: self.log_message(f"Found {len(existing_images)} existing images") + # Show the first image if available + self.show_image(existing_images[0]) else: self.log_message("No existing images found")