Update app.py

app.py

@@ -5,7 +5,7 @@ import random
 from PIL import Image
 import base64
 import io
-from pathlib import Path
+from functools import partial
 
 # Define categories at the top so they are accessible throughout the code
 categories = [
@@ -58,11 +58,11 @@ class DataManager:
 
     def save_characters(self):
         with open(self.characters_file, 'w') as f:
-            json.dump(self.characters, f)
+            json.dump(self.characters, f, indent=4)
 
     def save_persistent_tags(self):
         with open(self.persistent_tags_file, 'w') as f:
-            json.dump(self.persistent_tags, f)
+            json.dump(self.persistent_tags, f, indent=4)
 
     def load_category_tags(self):
         if os.path.exists(self.category_tags_file):
@@ -73,7 +73,7 @@ class DataManager:
 
     def save_category_tags(self):
         with open(self.category_tags_file, 'w') as f:
-            json.dump(self.category_tags, f)
+            json.dump(self.category_tags, f, indent=4)
 
     def get_category_tags(self, category_var_name):
         # Return the tags list for the given category variable name
@@ -110,7 +110,7 @@ class DataManager:
         # Save the image if provided
         if image_data:
             try:
-                # Copy the uploaded image to the images folder
+                # Open and save the uploaded image to the images folder
                 image = Image.open(image_data)
                 image.save(image_path)
                 character['image_path'] = image_path
@@ -145,6 +145,11 @@ class DataManager:
                         image = Image.open(image_path)
                         image.save(new_image_path)
                         updated_character['image_path'] = new_image_path
+                        # Remove old image if name has changed
+                        if original_name != updated_character['name']:
+                            old_image_path = char.get('image_path')
+                            if old_image_path and os.path.exists(old_image_path):
+                                os.remove(old_image_path)
                     except Exception as e:
                         print(f"Error updating image: {e}")
                         updated_character['image_path'] = char.get('image_path')
@@ -176,8 +181,9 @@ class DataManager:
                 return True
         return False
 
-def character_creation_app(data_manager):
+def character_creation_app(data_manager, refresh_characters_fn):
     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)")
@@ -198,6 +204,7 @@ def character_creation_app(data_manager):
             character = {'name': name, 'traits': traits, 'gender': gender, 'image': image_path}
 
             data_manager.add_character(character)
+            refresh_characters_fn()
             return f"Character '{name}' saved successfully."
 
         save_button.click(save_character, inputs=[name_input, traits_input, image_input, gender_input], outputs=output)
@@ -214,8 +221,8 @@ def character_creation_app(data_manager):
         def list_characters():
             characters = data_manager.get_characters()
             if not characters:
-                return "No characters created yet."
-
+                return [gr.Markdown("No characters created yet.")]
+            
             components = []
             for char in characters:
                 with gr.Accordion(label=char['name'], open=False):
@@ -232,14 +239,15 @@ def character_creation_app(data_manager):
                             with gr.Row():
                                 edit_btn = gr.Button("Edit", variant="primary")
                                 delete_btn = gr.Button("Delete", variant="secondary")
-                            edit_output = gr.Textbox(label="", interactive=False)
-                            delete_output = gr.Textbox(label="", interactive=False)
+                            edit_output = gr.Textbox(label="", interactive=False, visible=False)
+                            delete_output = gr.Textbox(label="", interactive=False, visible=False)
 
                             # Edit Interface (initially hidden)
-                            edit_interface = gr.Row(visible=False)
+                            edit_interface = gr.Column(visible=False)
                             with edit_interface:
-                                edit_name = gr.Textbox(label="Character Name", value=char['name'])
-                                edit_traits = gr.Textbox(label="Traits/Appearance Tags (comma separated)", value=', '.join(char['traits']))
+                                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_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")
@@ -261,6 +269,7 @@ def character_creation_app(data_manager):
                                     }
                                     success = data_manager.update_character(original_name, updated_char)
                                     if success:
+                                        refresh_characters_fn()
                                         return f"Character '{new_name}' updated successfully."
                                     else:
                                         return "Failed to update character."
@@ -272,11 +281,11 @@ def character_creation_app(data_manager):
                                 )
 
                             # Define edit functionality
-                            def toggle_edit_interface(visible, *args):
-                                return gr.update(visible=visible)
+                            def toggle_edit_visibility(event, current_visibility=False):
+                                return not current_visibility
 
                             edit_btn.click(
-                                fn=lambda: gr.update(visible=True),
+                                toggle_edit_visibility,
                                 inputs=None,
                                 outputs=edit_interface
                             )
@@ -288,29 +297,28 @@ def character_creation_app(data_manager):
                             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
+
                             delete_btn.click(
                                 lambda name: confirm_delete(name),
                                 inputs=[char['name']],
-                                outputs=delete_output
+                                outputs=[delete_output, delete_output]
                             )
 
                             delete_confirm = gr.Button("Confirm Delete", visible=False)
                             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), message
-                                return gr.update(visible=False), message
-
                             delete_btn.click(
-                                show_delete_buttons,
-                                inputs=[delete_output],
-                                outputs=[delete_confirm, delete_output]
+                                lambda: gr.update(visible=True),
+                                inputs=None,
+                                outputs=[delete_confirm]
                             )
 
                             # Handle delete confirmation
@@ -325,188 +333,195 @@ def character_creation_app(data_manager):
                                 outputs=delete_output
                             )
 
-                    components.append("")  # Placeholder for spacing
-
+                    # Add spacing or separator if needed
+                    components.append("")
             return components
 
-        characters_container.load(fn=list_characters, outputs=characters_container)
+        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=list_characters, outputs=characters_container)
 
-def prompt_generator_app(data_manager):
-    with gr.Tab("Prompt Generator"):
-        gr.Markdown("## Prompt Generator")
+        refresh_button.click(fn=refresh_characters, outputs=characters_container)
 
-        # Add a refresh tags button
-        refresh_tags_button = gr.Button("Refresh Tags")
+    def prompt_generator_app(data_manager):
+        with gr.Tab("Prompt Generator"):
+            gr.Markdown("## Prompt Generator")
 
-        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
+            # Add a refresh tags button
+            refresh_tags_button = gr.Button("Refresh Tags")
 
-            character_options = get_character_options()
-            character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters", interactive=True)
+            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")
 
-            refresh_characters_button = gr.Button("Refresh Character List")
+                # 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
 
-            def refresh_characters():
-                new_options = get_character_options()
-                return gr.CheckboxGroup.update(choices=new_options)
+                character_options = get_character_options()
+                character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters", interactive=True)
 
-            refresh_characters_button.click(refresh_characters, outputs=character_select)
+                refresh_characters_button = gr.Button("Refresh Character List")
 
-        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)")
+                def refresh_characters_fn():
+                    new_options = get_character_options()
+                    return new_options
 
-        generate_button = gr.Button("Generate Prompt")
-        prompt_output = gr.Textbox(label="Generated Prompt", lines=5)
+                refresh_characters_button.click(
+                    lambda: get_character_options(),
+                    outputs=character_select
+                )
 
-        def generate_prompt(*args):
-            arg_idx = 0
+            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)")
 
-            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
+            generate_button = gr.Button("Generate Prompt")
+            prompt_output = gr.Textbox(label="Generated Prompt", lines=5)
 
-                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)
+            def generate_prompt(*args):
+                arg_idx = 0
 
-            # Handle Characters
-            selected_character_options = args[arg_idx]
-            random_chars = args[arg_idx + 1]
-            num_random_chars = args[arg_idx + 2]
+                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
 
-            arg_idx += 3
+                    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)
 
-            characters = data_manager.get_characters()
+                # Handle Characters
+                selected_character_options = args[arg_idx]
+                random_chars = args[arg_idx + 1]
+                num_random_chars = args[arg_idx + 2]
 
-            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])
+                arg_idx += 3
 
-        generate_button.click(generate_prompt, inputs=inputs_list, outputs=prompt_output)
+                characters = data_manager.get_characters()
 
-        # 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
+                if random_chars:
+                    num = min(len(characters), int(num_random_chars))
+                    selected_chars = random.sample(characters, num)
                 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)
+                    # 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 tags_app(data_manager):
     with gr.Tab("Tags"):
@@ -549,9 +564,11 @@ def main():
     data_manager = DataManager(base_dir='/data')
     with gr.Blocks() as demo:
         with gr.Tabs():
-            prompt_generator_app(data_manager)
-            character_creation_app(data_manager)
-            tags_app(data_manager)
+            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)
 
     demo.launch()