Update app.py

app.py

@@ -5,7 +5,6 @@ import random
 from PIL import Image
 import base64
 import io
-from functools import partial
 
 # Define categories at the top so they are accessible throughout the code
 categories = [
@@ -84,7 +83,7 @@ class DataManager:
         self.save_category_tags()
 
     def get_characters(self):
-        # Load character images
+        # Ensure images paths are up-to-date
         for char in self.characters:
             image_path = char.get('image_path')
             if image_path and os.path.exists(image_path):
@@ -181,33 +180,209 @@ class DataManager:
                 return True
         return False
 
-def character_creation_app(data_manager, refresh_characters_fn):
+def prompt_generator_app(data_manager):
+    with gr.Tab("Prompt Generator"):
+        gr.Markdown("## Prompt Generator")
+
+        # Add a refresh tags button
+        refresh_tags_button = gr.Button("Refresh Tags")
+
+        inputs = {}
+        tag_displays = {}
+        for category_name, var_name in categories:
+            tags_list = data_manager.get_category_tags(var_name)
+            tags_string = ', '.join(tags_list)
+            max_tags = len(tags_list)
+            if max_tags == 0:
+                default_value = 0
+            else:
+                default_value = min(1, max_tags)
+            with gr.Group():
+                gr.Markdown(f"### {category_name}")
+                tag_display = gr.Markdown(f"**Tags:** {tags_string}")
+                tag_num = gr.Slider(minimum=0, maximum=max_tags, step=1, value=default_value, label=f"Number of {category_name} Tags to Select")
+            inputs[f"{var_name}_num"] = tag_num
+            tag_displays[var_name] = (tag_display, tag_num)
+
+        # For Character Selection
+        with gr.Group():
+            gr.Markdown("### Character Selection")
+
+            # Get the list of characters
+            def get_character_options():
+                characters = data_manager.get_characters()
+                character_options = []
+                for char in characters:
+                    option_label = f"{char['name']} ({char['gender']})"
+                    character_options.append(option_label)
+                return character_options
+
+            character_options = get_character_options()
+            character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters", interactive=True)
+
+            refresh_characters_button = gr.Button("Refresh Character List")
+
+            def refresh_characters():
+                new_options = get_character_options()
+                return gr.CheckboxGroup.update(choices=new_options)
+
+            refresh_characters_button.click(refresh_characters, outputs=character_select)
+
+        random_characters = gr.Checkbox(label="Select Random Characters")
+        num_characters = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Characters (if random)")
+
+        generate_button = gr.Button("Generate Prompt")
+        prompt_output = gr.Textbox(label="Generated Prompt", lines=5)
+
+        def generate_prompt(*args):
+            arg_idx = 0
+
+            prompt_tags = []
+            for category_name, var_name in categories:
+                tags_list = data_manager.get_category_tags(var_name)
+                tags_num = args[arg_idx]
+                arg_idx += 1
+
+                if tags_list and tags_num > 0:
+                    selected_tags = random.sample(tags_list, min(len(tags_list), int(tags_num)))
+                    prompt_tags.extend(selected_tags)
+
+            # Handle Characters
+            selected_character_options = args[arg_idx]
+            random_chars = args[arg_idx + 1]
+            num_random_chars = args[arg_idx + 2]
+
+            arg_idx += 3
+
+            characters = data_manager.get_characters()
+
+            if random_chars:
+                num = min(len(characters), int(num_random_chars))
+                selected_chars = random.sample(characters, num)
+            else:
+                # Extract selected character names from options
+                selected_chars = []
+                for option in selected_character_options:
+                    name = option.split(' (')[0]
+                    for char in characters:
+                        if char['name'] == name:
+                            selected_chars.append(char)
+                            break
+
+            # Determine the number of boys and girls
+            num_girls = sum(1 for char in selected_chars if char.get('gender') == 'Girl')
+            num_boys = sum(1 for char in selected_chars if char.get('gender') == 'Boy')
+
+            # Build the initial character count tags
+            character_count_tags = []
+            if num_girls > 0:
+                character_count_tags.append(f"{num_girls}girl" if num_girls == 1 else f"{num_girls}girls")
+            if num_boys > 0:
+                character_count_tags.append(f"{num_boys}boy" if num_boys == 1 else f"{num_boys}boys")
+
+            prompt_parts = []
+
+            if character_count_tags:
+                prompt_parts.append(', '.join(character_count_tags))
+
+            # Build character descriptions
+            character_descriptions = []
+            for idx, char in enumerate(selected_chars):
+                # Get traits for the character
+                traits = ', '.join(char['traits'])
+                # Create a description for each character
+                # For SDXL models, use the format "[char1 description] AND [char2 description]"
+                # Each character's description is enclosed in parentheses
+                character_description = f"({traits})"
+                character_descriptions.append(character_description)
+
+            # Join character descriptions appropriately for SDXL models
+            if character_descriptions:
+                character_descriptions_str = ' AND '.join(character_descriptions)
+                prompt_parts.append(character_descriptions_str)
+
+            # Append selected prompt tags from categories
+            if prompt_tags:
+                prompt_tags_str = ', '.join(prompt_tags)
+                prompt_parts.append(prompt_tags_str)
+
+            # Load persistent tags
+            persistent_tags = data_manager.get_persistent_tags()
+            if persistent_tags:
+                persistent_tags_str = ', '.join(persistent_tags)
+                prompt_parts.append(persistent_tags_str)
+
+            # Add ending tags
+            ending_tags = "source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres, anime artwork, anime style, vibrant, studio anime, highly detailed"
+            prompt_parts.append(ending_tags)
+
+            prompt_string = ', '.join(prompt_parts)
+
+            return prompt_string
+
+        # Prepare the list of inputs for the generate_prompt function
+        inputs_list = []
+        for category_name, var_name in categories:
+            inputs_list.append(inputs[f"{var_name}_num"])
+        # Add character_select directly to inputs
+        inputs_list.extend([character_select, random_characters, num_characters])
+
+        generate_button.click(generate_prompt, inputs=inputs_list, outputs=prompt_output)
+
+        # Function to refresh tags display and sliders
+        def refresh_tags():
+            updates = []
+            for category_name, var_name in categories:
+                # Reload tags from data_manager
+                tags_list = data_manager.get_category_tags(var_name)
+                tags_string = ', '.join(tags_list)
+                max_tags = len(tags_list)
+                if max_tags == 0:
+                    slider_value = 0
+                else:
+                    slider_value = min(1, max_tags)
+                # Update the tag display and slider
+                tag_display, tag_num = tag_displays[var_name]
+                updates.append(gr.Markdown.update(value=f"**Tags:** {tags_string}"))
+                updates.append(gr.Slider.update(maximum=max_tags, value=slider_value))
+            return updates
+
+        # Prepare the outputs list
+        outputs = [component for pair in tag_displays.values() for component in pair]
+
+        # Connect the refresh_tags function to the refresh_tags_button
+        refresh_tags_button.click(refresh_tags, outputs=outputs)
+
+def character_creation_app(data_manager):
     with gr.Tab("Character Creation"):
         gr.Markdown("## Create a New Character")
         with gr.Row():
-            name_input = gr.Textbox(label="Character Name")
-            traits_input = gr.Textbox(label="Traits/Appearance Tags (comma separated)")
+            name_input = gr.Textbox(label="Character Name", placeholder="Enter unique character name")
+            traits_input = gr.Textbox(label="Traits/Appearance Tags (comma separated)", placeholder="e.g., blue hair, green eyes, tall")
         image_input = gr.Image(label="Upload Character Image", type="filepath")
-        # New gender selection input
         gender_input = gr.Radio(choices=["Boy", "Girl"], label="Gender")
         save_button = gr.Button("Save Character")
         output = gr.Textbox(label="Status", interactive=False)
 
-        def save_character(name, traits, image_path, gender):
-            if not name or not traits or not gender:
-                return "Please enter all fields."
+        def save_character(name, traits, image_path, gender, characters_container):
+            if not name.strip() or not traits.strip() or not gender:
+                return "Please enter all fields.", None
             # Check for duplicate names
             existing_names = [char['name'] for char in data_manager.get_characters()]
             if name in existing_names:
-                return f"Character with name '{name}' already exists. Please choose a different name."
+                return f"Character with name '{name}' already exists. Please choose a different name.", None
 
             character = {'name': name, 'traits': traits, 'gender': gender, 'image': image_path}
 
             data_manager.add_character(character)
-            refresh_characters_fn()
-            return f"Character '{name}' saved successfully."
+            # After adding, re-render the character list
+            return f"Character '{name}' saved successfully.", list_characters(data_manager)
 
-        save_button.click(save_character, inputs=[name_input, traits_input, image_input, gender_input], outputs=output)
+        save_button.click(
+            save_character, 
+            inputs=[name_input, traits_input, image_input, gender_input, gr.State()],
+            outputs=[output, "characters_renderer"]
+        )
 
         # Divider
         gr.Markdown("---")
@@ -215,20 +390,19 @@ def character_creation_app(data_manager, refresh_characters_fn):
         # Display Existing Characters
         gr.Markdown("## Existing Characters")
 
-        # Container to hold the list of characters
-        characters_container = gr.Column()
+        # Renderer for characters
+        characters_renderer = gr.Column(label="Characters")
 
-        def list_characters():
+        def list_characters(data_manager):
             characters = data_manager.get_characters()
             if not characters:
-                return [gr.Markdown("No characters created yet.")]
+                return gr.Markdown("No characters created yet.")
             
             components = []
             for char in characters:
                 with gr.Accordion(label=char['name'], open=False):
                     with gr.Row():
                         if char['image']:
-                            # Display the image using filepath
                             img_display = gr.Image(value=char['image'], label="Character Image", type="filepath", interactive=False)
                         else:
                             img_display = gr.Markdown("**No Image Provided**")
@@ -246,67 +420,64 @@ def character_creation_app(data_manager, refresh_characters_fn):
                             edit_interface = gr.Column(visible=False)
                             with edit_interface:
                                 with gr.Row():
-                                    edit_name = gr.Textbox(label="Character Name", value=char['name'])
-                                    edit_traits = gr.Textbox(label="Traits/Appearance Tags (comma separated)", value=', '.join(char['traits']))
+                                    edit_name = gr.Textbox(label="Character Name", value=char['name'], placeholder="Enter new name")
+                                    edit_traits = gr.Textbox(label="Traits/Appearance Tags (comma separated)", value=', '.join(char['traits']), placeholder="e.g., red hair, blue eyes")
                                 edit_image = gr.Image(label="Upload New Character Image", type="filepath")
                                 edit_gender = gr.Radio(choices=["Boy", "Girl"], label="Gender", value=char['gender'])
                                 save_edit_btn = gr.Button("Save Changes")
                                 save_edit_output = gr.Textbox(label="Edit Status", interactive=False)
 
                                 def save_edit(original_name, new_name, new_traits, new_image_path, new_gender):
-                                    if not new_name or not new_traits or not new_gender:
-                                        return "Please enter all fields."
+                                    if not new_name.strip() or not new_traits.strip() or not new_gender:
+                                        return "Please enter all fields.", None
                                     # If the name has changed, check for duplicates
                                     if new_name != original_name:
                                         existing_names = [c['name'] for c in data_manager.get_characters()]
                                         if new_name in existing_names:
-                                            return f"Character with name '{new_name}' already exists. Please choose a different name."
+                                            return f"Character with name '{new_name}' already exists. Please choose a different name.", None
                                     updated_char = {
                                         'name': new_name,
                                         'traits': new_traits,
                                         'gender': new_gender,
-                                        'image': new_image_path
+                                        'image': new_image_path if new_image_path else data_manager.characters[[c['name'] for c in data_manager.get_characters()].index(original_name)]['image']
                                     }
                                     success = data_manager.update_character(original_name, updated_char)
                                     if success:
-                                        refresh_characters_fn()
-                                        return f"Character '{new_name}' updated successfully."
+                                        return f"Character '{new_name}' updated successfully.", list_characters(data_manager)
                                     else:
-                                        return "Failed to update character."
+                                        return "Failed to update character.", None
 
                                 save_edit_btn.click(
                                     save_edit, 
-                                    inputs=[gr.State(char['name']), edit_name, edit_traits, edit_image, edit_gender], 
-                                    outputs=save_edit_output
+                                    inputs=[gr.State(char['name']), edit_name, edit_traits, edit_image, edit_gender],
+                                    outputs=[save_edit_output, "characters_renderer"]
                                 )
 
                             # Define edit functionality
-                            def toggle_edit_visibility(event, current_visibility=False):
-                                return not current_visibility
+                            def toggle_edit_visibility():
+                                return gr.update(visible=True)
 
                             edit_btn.click(
-                                toggle_edit_visibility,
+                                lambda: True,
                                 inputs=None,
                                 outputs=edit_interface
                             )
 
                             # Define delete functionality
                             def confirm_delete(char_name):
-                                return f"Are you sure you want to delete '{char_name}'?"
+                                return f"Are you sure you want to delete '{char_name}'?", True
 
-                            def perform_delete(char_name):
-                                success = data_manager.delete_character(char_name)
-                                if success:
-                                    refresh_characters_fn()
-                                    return f"Character '{char_name}' deleted successfully."
-                                else:
-                                    return f"Failed to delete character '{char_name}'."
-
-                            def show_confirmation(message):
-                                return message, True
+                            def perform_delete(char_name, confirm):
+                                if confirm:
+                                    success = data_manager.delete_character(char_name)
+                                    if success:
+                                        return f"Character '{char_name}' deleted successfully.", list_characters(data_manager)
+                                    else:
+                                        return f"Failed to delete character '{char_name}'.", None
+                                return "Deletion cancelled.", None
 
                             delete_btn.click(
-                                lambda name: confirm_delete(name),
+                                confirm_delete,
                                 inputs=[char['name']],
                                 outputs=[delete_output, delete_output]
                             )
@@ -315,213 +486,36 @@ def character_creation_app(data_manager, refresh_characters_fn):
                             delete_cancel = gr.Button("Cancel Delete", visible=False)
 
                             # Show confirmation buttons when delete is clicked
+                            def show_delete_buttons(message):
+                                if "Are you sure" in message:
+                                    return gr.update(visible=True), gr.update(visible=True)
+                                return gr.update(visible=False), gr.update(visible=False)
+
                             delete_btn.click(
-                                lambda: gr.update(visible=True),
-                                inputs=None,
-                                outputs=[delete_confirm]
+                                show_delete_buttons,
+                                inputs=[delete_output],
+                                outputs=[delete_confirm, delete_cancel]
                             )
 
                             # Handle delete confirmation
                             delete_confirm.click(
                                 perform_delete,
-                                inputs=[char['name']],
-                                outputs=delete_output
+                                inputs=[char['name'], gr.Checkbox(label="", value=True)],
+                                outputs=[delete_output, "characters_renderer"]
                             )
 
+                            # Handle delete cancellation
                             delete_cancel.click(
-                                lambda: "",
-                                outputs=delete_output
+                                lambda: ("Deletion cancelled.", None),
+                                inputs=None,
+                                outputs=[delete_output, "characters_renderer"]
                             )
+                return components
 
-                    # Add spacing or separator if needed
-                    components.append("")
-            return components
-
-        def refresh_characters():
-            return list_characters()
-
-        characters_container.render(list_characters())
-
-        # Refresh Button to reload characters
-        refresh_button = gr.Button("Refresh Characters")
-
-        refresh_button.click(fn=refresh_characters, outputs=characters_container)
-
-    def prompt_generator_app(data_manager):
-        with gr.Tab("Prompt Generator"):
-            gr.Markdown("## Prompt Generator")
-
-            # Add a refresh tags button
-            refresh_tags_button = gr.Button("Refresh Tags")
-
-            inputs = {}
-            tag_displays = {}
-            for category_name, var_name in categories:
-                tags_list = data_manager.get_category_tags(var_name)
-                tags_string = ', '.join(tags_list)
-                max_tags = len(tags_list)
-                if max_tags == 0:
-                    default_value = 0
-                else:
-                    default_value = min(1, max_tags)
-                with gr.Group():
-                    gr.Markdown(f"### {category_name}")
-                    tag_display = gr.Markdown(f"**Tags:** {tags_string}")
-                    tag_num = gr.Slider(minimum=0, maximum=max_tags, step=1, value=default_value, label=f"Number of {category_name} Tags to Select")
-                inputs[f"{var_name}_num"] = tag_num
-                tag_displays[var_name] = (tag_display, tag_num)
-
-            # For Character Selection
-            with gr.Group():
-                gr.Markdown("### Character Selection")
-
-                # Get the list of characters
-                def get_character_options():
-                    characters = data_manager.get_characters()
-                    character_options = []
-                    for char in characters:
-                        option_label = f"{char['name']} ({char['gender']})"
-                        character_options.append(option_label)
-                    return character_options
-
-                character_options = get_character_options()
-                character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters", interactive=True)
-
-                refresh_characters_button = gr.Button("Refresh Character List")
-
-                def refresh_characters_fn():
-                    new_options = get_character_options()
-                    return new_options
-
-                refresh_characters_button.click(
-                    lambda: get_character_options(),
-                    outputs=character_select
-                )
-
-            random_characters = gr.Checkbox(label="Select Random Characters")
-            num_characters = gr.Slider(minimum=1, maximum=10, step=1, value=1, label="Number of Characters (if random)")
-
-            generate_button = gr.Button("Generate Prompt")
-            prompt_output = gr.Textbox(label="Generated Prompt", lines=5)
+        def list_characters_render():
+            return list_characters(data_manager)
 
-            def generate_prompt(*args):
-                arg_idx = 0
-
-                prompt_tags = []
-                for category_name, var_name in categories:
-                    tags_list = data_manager.get_category_tags(var_name)
-                    tags_num = args[arg_idx]
-                    arg_idx += 1
-
-                    if tags_list and tags_num > 0:
-                        selected_tags = random.sample(tags_list, min(len(tags_list), int(tags_num)))
-                        prompt_tags.extend(selected_tags)
-
-                # Handle Characters
-                selected_character_options = args[arg_idx]
-                random_chars = args[arg_idx + 1]
-                num_random_chars = args[arg_idx + 2]
-
-                arg_idx += 3
-
-                characters = data_manager.get_characters()
-
-                if random_chars:
-                    num = min(len(characters), int(num_random_chars))
-                    selected_chars = random.sample(characters, num)
-                else:
-                    # Extract selected character names from options
-                    selected_chars = []
-                    for option in selected_character_options:
-                        name = option.split(' (')[0]
-                        for char in characters:
-                            if char['name'] == name:
-                                selected_chars.append(char)
-                                break
-
-                # Determine the number of boys and girls
-                num_girls = sum(1 for char in selected_chars if char.get('gender') == 'Girl')
-                num_boys = sum(1 for char in selected_chars if char.get('gender') == 'Boy')
-
-                # Build the initial character count tags
-                character_count_tags = []
-                if num_girls > 0:
-                    character_count_tags.append(f"{num_girls}girl" if num_girls == 1 else f"{num_girls}girls")
-                if num_boys > 0:
-                    character_count_tags.append(f"{num_boys}boy" if num_boys == 1 else f"{num_boys}boys")
-
-                prompt_parts = []
-
-                if character_count_tags:
-                    prompt_parts.append(', '.join(character_count_tags))
-
-                # Build character descriptions
-                character_descriptions = []
-                for idx, char in enumerate(selected_chars):
-                    # Get traits for the character
-                    traits = ', '.join(char['traits'])
-                    # Create a description for each character
-                    # For SDXL models, use the format "[char1 description] AND [char2 description]"
-                    # Each character's description is enclosed in parentheses
-                    character_description = f"({traits})"
-                    character_descriptions.append(character_description)
-
-                # Join character descriptions appropriately for SDXL models
-                if character_descriptions:
-                    character_descriptions_str = ' AND '.join(character_descriptions)
-                    prompt_parts.append(character_descriptions_str)
-
-                # Append selected prompt tags from categories
-                if prompt_tags:
-                    prompt_tags_str = ', '.join(prompt_tags)
-                    prompt_parts.append(prompt_tags_str)
-
-                # Load persistent tags
-                persistent_tags = data_manager.get_persistent_tags()
-                if persistent_tags:
-                    persistent_tags_str = ', '.join(persistent_tags)
-                    prompt_parts.append(persistent_tags_str)
-
-                # Add ending tags
-                ending_tags = "source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres, anime artwork, anime style, vibrant, studio anime, highly detailed"
-                prompt_parts.append(ending_tags)
-
-                prompt_string = ', '.join(prompt_parts)
-
-                return prompt_string
-
-            # Prepare the list of inputs for the generate_prompt function
-            inputs_list = []
-            for category_name, var_name in categories:
-                inputs_list.append(inputs[f"{var_name}_num"])
-            # Add character_select directly to inputs
-            inputs_list.extend([character_select, random_characters, num_characters])
-
-            generate_button.click(generate_prompt, inputs=inputs_list, outputs=prompt_output)
-
-            # Function to refresh tags display and sliders
-            def refresh_tags():
-                updates = []
-                for category_name, var_name in categories:
-                    # Reload tags from data_manager
-                    tags_list = data_manager.get_category_tags(var_name)
-                    tags_string = ', '.join(tags_list)
-                    max_tags = len(tags_list)
-                    if max_tags == 0:
-                        slider_value = 0
-                    else:
-                        slider_value = min(1, max_tags)
-                    # Update the tag display and slider
-                    tag_display, tag_num = tag_displays[var_name]
-                    updates.append(gr.Markdown.update(value=f"**Tags:** {tags_string}"))
-                    updates.append(gr.Slider.update(maximum=max_tags, value=slider_value))
-                return updates
-
-            # Prepare the outputs list
-            outputs = [component for pair in tag_displays.values() for component in pair]
-
-            # Connect the refresh_tags function to the refresh_tags_button
-            refresh_tags_button.click(refresh_tags, outputs=outputs)
+        characters_renderer.render(list_characters(data_manager))
 
 def tags_app(data_manager):
     with gr.Tab("Tags"):
@@ -535,7 +529,6 @@ def tags_app(data_manager):
             save_button = gr.Button(f"Save {category_name} Tags")
             status_output = gr.Textbox(label="", interactive=False)
 
-            # Function to save tags
             def make_save_category_tags_fn(var_name, category_name):
                 def fn(tags_string):
                     tags_list = [t.strip() for t in tags_string.split(',') if t.strip()]
@@ -553,22 +546,20 @@ def tags_app(data_manager):
         save_persistent_tags_button = gr.Button("Save Persistent Tags")
         persistent_status_output = gr.Textbox(label="", interactive=False)
 
-        def save_persistent_tags(tags_string):
+        def save_persistent_tags_fn(tags_string):
             tags_list = [t.strip() for t in tags_string.split(',') if t.strip()]
             data_manager.set_persistent_tags(tags_list)
             return "Persistent tags saved successfully."
 
-        save_persistent_tags_button.click(save_persistent_tags, inputs=persistent_tags_input, outputs=persistent_status_output)
+        save_persistent_tags_button.click(save_persistent_tags_fn, inputs=persistent_tags_input, outputs=persistent_status_output)
 
 def main():
     data_manager = DataManager(base_dir='/data')
     with gr.Blocks() as demo:
         with gr.Tabs():
-            with gr.Row():
-                with gr.Column():
-                    prompt_generator_app(data_manager)
-                    character_creation_app(data_manager, refresh_characters_fn=lambda: None)
-                tags_app(data_manager)
+            prompt_generator_app(data_manager)
+            character_creation_app(data_manager)
+            tags_app(data_manager)
 
     demo.launch()