Update app.py

app.py

@@ -5,6 +5,7 @@ import random
 from PIL import Image
 import base64
 import io
+from pathlib import Path
 
 # Define categories at the top so they are accessible throughout the code
 categories = [
@@ -124,35 +125,208 @@ class DataManager:
         self.characters.append(character)
         self.save_characters()
 
+    def update_character(self, original_name, updated_character):
+        # Find the character by original_name
+        for idx, char in enumerate(self.characters):
+            if char['name'] == original_name:
+                # Handle image update
+                if updated_character['image']:
+                    image_data = updated_character['image']  # base64 string
+                    safe_name = "".join(c for c in updated_character['name'] if c.isalnum() or c in (' ', '_', '-')).rstrip()
+                    image_filename = f"{safe_name}.png"
+                    image_path = os.path.join(self.images_folder, image_filename)
+
+                    # Decode and save the new image
+                    image = Image.open(io.BytesIO(base64.b64decode(image_data.split(",")[1])))
+                    image.save(image_path)
+                    updated_character['image_path'] = image_path
+                else:
+                    # If no new image provided, retain the old image_path
+                    updated_character['image_path'] = char.get('image_path')
+
+                updated_character.pop('image', None)
+
+                # Process traits
+                if isinstance(updated_character['traits'], str):
+                    updated_character['traits'] = updated_character['traits'].split(',')
+                    updated_character['traits'] = [t.strip() for t in updated_character['traits']]
+
+                self.characters[idx] = updated_character
+                self.save_characters()
+                return True
+        return False
+
+    def delete_character(self, name):
+        # Find and remove the character
+        for idx, char in enumerate(self.characters):
+            if char['name'] == name:
+                image_path = char.get('image_path')
+                if image_path and os.path.exists(image_path):
+                    os.remove(image_path)
+                del self.characters[idx]
+                self.save_characters()
+                return True
+        return False
+
 def character_creation_app(data_manager):
     with gr.Tab("Character Creation"):
         with gr.Row():
             name_input = gr.Textbox(label="Character Name")
             traits_input = gr.Textbox(label="Traits/Appearance Tags (comma separated)")
-        image_input = gr.Image(label="Upload Character Image", type="filepath")
+        image_input = gr.Image(label="Upload Character Image", type="base64")
         # 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):
+        def save_character(name, traits, image_data, gender):
             if not name or not traits or not gender:
                 return "Please enter all fields."
-            character = {'name': name, 'traits': traits, 'gender': gender, 'image': 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."
 
-            # Read and encode image if provided
-            if image_path:
-                with open(image_path, "rb") as img_file:
-                    image_data = base64.b64encode(img_file.read()).decode('utf-8')
-                    character['image'] = f"data:image/png;base64,{image_data}"
-            else:
-                character['image'] = None
+            character = {'name': name, 'traits': traits, 'gender': gender, 'image': image_data}
 
             data_manager.add_character(character)
             return f"Character '{name}' saved successfully."
 
         save_button.click(save_character, inputs=[name_input, traits_input, image_input, gender_input], outputs=output)
 
+        # Divider
+        gr.Markdown("---")
+
+        # Display Existing Characters
+        gr.Markdown("## Existing Characters")
+
+        # Container to hold the list of characters
+        characters_container = gr.Column()
+
+        def list_characters():
+            characters = data_manager.get_characters()
+            if not characters:
+                return "No characters created yet."
+
+            components = []
+            for char in characters:
+                with gr.AccordionItem(label=char['name']):
+                    with gr.Row():
+                        if char['image']:
+                            image = Image.open(char['image'])
+                            image_thumbnail = image.copy()
+                            image_thumbnail.thumbnail((150, 150))
+                            buffered = io.BytesIO()
+                            image_thumbnail.save(buffered, format="PNG")
+                            img_str = base64.b64encode(buffered.getvalue()).decode()
+                            img_data = f"data:image/png;base64,{img_str}"
+                            gr.Image(value=img_data, label="Character Image", width=150).render()
+                        else:
+                            gr.Markdown("**No Image Provided**").render()
+                        with gr.Column():
+                            gr.Markdown(f"**Name:** {char['name']}").render()
+                            gr.Markdown(f"**Gender:** {char['gender']}").render()
+                            gr.Markdown(f"**Traits:** {', '.join(char['traits'])}").render()
+                            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)
+
+                            # Define edit functionality
+                            def initiate_edit(char_name):
+                                return gr.update(visible=True), char_name
+
+                            edit_interface = gr.Accordion(label=f"Edit {char['name']}", open=False, 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_image = gr.Image(label="Upload New Character Image", type="base64")
+                                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, new_gender):
+                                    if not new_name or not new_traits or not new_gender:
+                                        return "Please enter all fields."
+                                    # 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."
+                                    updated_char = {
+                                        'name': new_name,
+                                        'traits': new_traits,
+                                        'gender': new_gender,
+                                        'image': new_image
+                                    }
+                                    success = data_manager.update_character(original_name, updated_char)
+                                    if success:
+                                        return f"Character '{new_name}' updated successfully."
+                                    else:
+                                        return "Failed to update character."
+
+                                save_edit_btn.click(
+                                    save_edit, 
+                                    inputs=[gr.State(char['name']), edit_name, edit_traits, edit_image, edit_gender], 
+                                    outputs=save_edit_output
+                                )
+
+                            # Define delete functionality
+                            def confirm_delete(char_name):
+                                return f"Are you sure you want to delete '{char_name}'?"
+
+                            def perform_delete(char_name):
+                                success = data_manager.delete_character(char_name)
+                                if success:
+                                    return f"Character '{char_name}' deleted successfully."
+                                else:
+                                    return f"Failed to delete character '{char_name}'."
+
+                            delete_btn.click(
+                                lambda name: confirm_delete(name),
+                                inputs=[char['name']],
+                                outputs=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]
+                            )
+
+                            # Handle delete confirmation
+                            delete_confirm.click(
+                                perform_delete,
+                                inputs=[char['name']],
+                                outputs=delete_output
+                            )
+
+                            delete_cancel.click(
+                                lambda: "",
+                                outputs=delete_output
+                            )
+
+                    components.append("")  # Placeholder for spacing
+
+            return gr.Row(components)
+
+        characters_container.load(fn=list_characters, outputs=characters_container)
+
+        # 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")
@@ -366,9 +540,10 @@ def tags_app(data_manager):
 def main():
     data_manager = DataManager(base_dir='/data')
     with gr.Blocks() as demo:
-        prompt_generator_app(data_manager)
-        character_creation_app(data_manager)
-        tags_app(data_manager)
+        with gr.Tabs():
+            prompt_generator_app(data_manager)
+            character_creation_app(data_manager)
+            tags_app(data_manager)
 
     demo.launch()