Update app.py

app.py

@@ -5,6 +5,7 @@ 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 = [
@@ -107,7 +108,7 @@ class DataManager:
         image_path = os.path.join(self.images_folder, image_filename)
 
         # Save the image if provided
-        if image_path and os.path.exists(image_data):
+        if image_data and os.path.exists(image_data):
             try:
                 # Open and save the uploaded image to the images folder
                 image = Image.open(image_data)
@@ -134,7 +135,7 @@ class DataManager:
         for idx, char in enumerate(self.characters):
             if char['name'] == original_name:
                 # Handle image update
-                if updated_character['image']:
+                if updated_character['image'] and os.path.exists(updated_character['image']):
                     image_data = updated_character['image']  # File path
                     safe_name = "".join(c for c in updated_character['name'] if c.isalnum() or c in (' ', '_', '-')).rstrip()
                     image_filename = f"{safe_name}.png"
@@ -200,7 +201,13 @@ def prompt_generator_app(data_manager):
             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")
+                tag_num = gr.Slider(
+                    minimum=0,
+                    maximum=max_tags,
+                    step=1,
+                    value=default_value,
+                    label=f"Number of {category_name} Tags to Select"
+                )
             inputs_sliders[var_name + '_num'] = tag_num
             tag_displays[var_name] = (tag_display, tag_num)
 
@@ -218,17 +225,31 @@ def prompt_generator_app(data_manager):
                 return character_options
 
             character_options = get_character_options()
-            character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters", interactive=True)
+            character_select = gr.CheckboxGroup(
+                choices=character_options,
+                label="Select Characters",
+                interactive=True
+            )
 
             refresh_characters_button = gr.Button("Refresh Character List")
 
+            def refresh_character_selection():
+                new_options = get_character_options()
+                return gr.CheckboxGroup.update(choices=new_options)
+
             refresh_characters_button.click(
-                lambda: get_character_options(),
+                refresh_character_selection,
                 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)")
+        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)
@@ -342,8 +363,8 @@ def prompt_generator_app(data_manager):
                     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))
+                updates.append(tag_display.update(value=f"**Tags:** {tags_string}"))
+                updates.append(tag_num.update(maximum=max_tags, value=slider_value))
             return updates
 
         # Prepare the outputs list
@@ -355,10 +376,17 @@ def prompt_generator_app(data_manager):
 def character_creation_app(data_manager):
     with gr.Tab("Character Creation"):
         gr.Markdown("## Create and Manage Characters")
+
         # **Section 1: Create a New Character**
         with gr.Row():
-            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")
+            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")
         gender_input = gr.Radio(choices=["Boy", "Girl"], label="Gender")
         save_button = gr.Button("Save Character")
@@ -371,7 +399,11 @@ def character_creation_app(data_manager):
         # Dropdown to select a character for editing or deleting
         characters = data_manager.get_characters()
         character_names = [char['name'] for char in characters]
-        selected_character = gr.Dropdown(choices=character_names, label="Select a Character to Edit/Delete", interactive=True)
+        selected_character = gr.Dropdown(
+            choices=character_names,
+            label="Select a Character to Edit/Delete",
+            interactive=True
+        )
 
         # Buttons for Edit and Delete
         edit_button = gr.Button("Edit Selected Character")
@@ -379,17 +411,31 @@ def character_creation_app(data_manager):
         action_output = gr.Textbox(label="Action Status", interactive=False)
 
         # **Section 3: Edit Character**
-        edit_section = gr.Row(visible=False)
+        edit_section = gr.Column(visible=False)
         with edit_section:
-            edit_name = gr.Textbox(label="New Character Name", placeholder="Enter new character name")
-            edit_traits = gr.Textbox(label="New Traits/Appearance Tags (comma separated)", placeholder="e.g., red hair, blue eyes")
+            with gr.Row():
+                edit_name = gr.Textbox(
+                    label="New Character Name",
+                    placeholder="Enter new character name"
+                )
+                edit_traits = gr.Textbox(
+                    label="New Traits/Appearance Tags (comma separated)",
+                    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")
+            edit_gender = gr.Radio(
+                choices=["Boy", "Girl"],
+                label="Gender"
+            )
             update_button = gr.Button("Update Character")
             update_output = gr.Textbox(label="Update Status", interactive=False)
 
         # **Section 4: Delete Confirmation**
-        delete_confirmation = gr.Textbox(label="Are you sure you want to delete the selected character? Type 'YES' to confirm.", interactive=True, visible=False)
+        delete_confirmation = gr.Textbox(
+            label="Are you sure you want to delete this character? Type 'YES' to confirm.",
+            interactive=True,
+            visible=False
+        )
         confirm_delete_button = gr.Button("Confirm Delete", visible=False)
         cancel_delete_button = gr.Button("Cancel Delete", visible=False)
 
@@ -416,7 +462,15 @@ def character_creation_app(data_manager):
         def show_edit(selected_name):
             if not selected_name:
                 return "Please select a character to edit.", gr.update(visible=False)
-            return "", gr.update(visible=True)
+            # Fetch character details
+            character = next((char for char in data_manager.get_characters() if char['name'] == selected_name), None)
+            if character:
+                edit_name_value = character['name']
+                edit_traits_value = ', '.join(character['traits'])
+                edit_gender_value = character['gender']
+                return "", gr.update(visible=True)
+            else:
+                return f"Character '{selected_name}' not found.", gr.update(visible=False)
 
         edit_button.click(
             show_edit,
@@ -437,7 +491,8 @@ def character_creation_app(data_manager):
                 'name': new_name,
                 'traits': new_traits,
                 'gender': new_gender,
-                'image': new_image_path if new_image_path else [char for char in data_manager.get_characters() if char['name'] == original_name][0]['image']
+                'image': new_image_path if new_image_path else \
+                         next((char['image'] for char in data_manager.get_characters() if char['name'] == original_name), None)
             }
             success = data_manager.update_character(original_name, updated_char)
             if success:
@@ -465,7 +520,7 @@ def character_creation_app(data_manager):
 
         # **Function to Perform Deletion**
         def perform_deletion(selected_name, confirmation):
-            if confirmation != "YES":
+            if confirmation.strip().upper() != "YES":
                 return "Deletion cancelled.", gr.update(visible=False), gr.update(visible=False)
             success = data_manager.delete_character(selected_name)
             if success:
@@ -497,7 +552,10 @@ def tags_app(data_manager):
             gr.Markdown(f"### {category_name} Tags")
             tags_list = data_manager.get_category_tags(var_name)
             tags_string = ', '.join(tags_list)
-            tag_input = gr.Textbox(label=f"{category_name} Tags (comma separated)", value=tags_string)
+            tag_input = gr.Textbox(
+                label=f"{category_name} Tags (comma separated)",
+                value=tags_string
+            )
             save_button = gr.Button(f"Save {category_name} Tags")
             status_output = gr.Textbox(label="", interactive=False)
 
@@ -514,7 +572,10 @@ def tags_app(data_manager):
         # Persistent Tags
         gr.Markdown(f"### Persistent Tags")
         persistent_tags_string = ', '.join(data_manager.get_persistent_tags())
-        persistent_tags_input = gr.Textbox(label="Persistent Tags (comma separated)", value=persistent_tags_string)
+        persistent_tags_input = gr.Textbox(
+            label="Persistent Tags (comma separated)",
+            value=persistent_tags_string
+        )
         save_persistent_tags_button = gr.Button("Save Persistent Tags")
         persistent_status_output = gr.Textbox(label="", interactive=False)
 
@@ -523,7 +584,11 @@ def tags_app(data_manager):
             data_manager.set_persistent_tags(tags_list)
             return "Persistent tags saved successfully."
 
-        save_persistent_tags_button.click(save_persistent_tags_fn, 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')