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 = [
@@ -21,7 +20,7 @@ categories = [
 ]
 
 class DataManager:
-    def __init__(self, base_dir='data'):
+    def __init__(self, base_dir='/data'):
         self.base_dir = base_dir
         # Ensure the base directory exists
         if not os.path.exists(self.base_dir):
@@ -58,11 +57,11 @@ class DataManager:
 
     def save_characters(self):
         with open(self.characters_file, 'w') as f:
-            json.dump(self.characters, f, indent=4)
+            json.dump(self.characters, f)
 
     def save_persistent_tags(self):
         with open(self.persistent_tags_file, 'w') as f:
-            json.dump(self.persistent_tags, f, indent=4)
+            json.dump(self.persistent_tags, f)
 
     def load_category_tags(self):
         if os.path.exists(self.category_tags_file):
@@ -73,7 +72,7 @@ class DataManager:
 
     def save_category_tags(self):
         with open(self.category_tags_file, 'w') as f:
-            json.dump(self.category_tags, f, indent=4)
+            json.dump(self.category_tags, f)
 
     def get_category_tags(self, category_var_name):
         # Return the tags list for the given category variable name
@@ -84,7 +83,7 @@ class DataManager:
         self.save_category_tags()
 
     def get_characters(self):
-        # Ensure images paths are up-to-date
+        # Load character images
         for char in self.characters:
             image_path = char.get('image_path')
             if image_path and os.path.exists(image_path):
@@ -102,21 +101,16 @@ class DataManager:
 
     def add_character(self, character):
         # Save image to disk and store the filename
-        image_data = character['image']  # This is the file path
+        image_data = character['image']  # This is base64 encoded string
         safe_name = "".join(c for c in 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)
 
-        # Save the image if provided
-        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)
-                image.save(image_path)
-                character['image_path'] = image_path
-            except Exception as e:
-                print(f"Error saving image: {e}")
-                character['image_path'] = None
+        # Decode the base64 image data and save it
+        if image_data:
+            image = Image.open(io.BytesIO(base64.b64decode(image_data.split(",")[1])))
+            image.save(image_path)
+            character['image_path'] = image_path
         else:
             character['image_path'] = None
 
@@ -130,56 +124,34 @@ 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'] 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"
-                    new_image_path = os.path.join(self.images_folder, image_filename)
-
-                    try:
-                        image = Image.open(image_data)
-                        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')
-                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")
+        # 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."
+            character = {'name': name, 'traits': traits, 'gender': gender, 'image': None}
+
+            # 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
+
+            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)
 
 def prompt_generator_app(data_manager):
     with gr.Tab("Prompt Generator"):
@@ -188,7 +160,7 @@ def prompt_generator_app(data_manager):
         # Add a refresh tags button
         refresh_tags_button = gr.Button("Refresh Tags")
 
-        inputs_sliders = {}
+        inputs = {}
         tag_displays = {}
         for category_name, var_name in categories:
             tags_list = data_manager.get_category_tags(var_name)
@@ -201,14 +173,8 @@ 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"
-                )
-            inputs_sliders[var_name + '_num'] = tag_num
+                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
@@ -225,31 +191,18 @@ 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():
+            def refresh_characters():
                 new_options = get_character_options()
-                return [gr.CheckboxGroup.update(choices=new_options)]
+                return gr.update(choices=new_options)
 
-            refresh_characters_button.click(
-                refresh_character_selection,
-                outputs=[character_select]
-            )
+            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)"
-        )
+        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)
@@ -343,18 +296,15 @@ def prompt_generator_app(data_manager):
         # Prepare the list of inputs for the generate_prompt function
         inputs_list = []
         for category_name, var_name in categories:
-            inputs_list.append(inputs_sliders[f"{var_name}_num"])
+            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
-        )
+        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)
@@ -366,192 +316,15 @@ 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]
-                tag_display.update(value=f"**Tags:** {tags_string}")
-                tag_num.update(maximum=max_tags, value=slider_value)
-            return None  # No outputs to return
+                updates.append(gr.update(value=f"**Tags:** {tags_string}"))
+                updates.append(gr.update(maximum=max_tags, value=slider_value))
+            return updates
 
-        # Connect the refresh_tags function to the refresh_tags_button
-        refresh_tags_button.click(
-            refresh_tags,
-            inputs=None,
-            outputs=None
-        )
+        # Prepare the outputs list
+        outputs = [component for pair in tag_displays.values() for component in pair]
 
-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"
-            )
-        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")
-        save_output = gr.Textbox(label="Status", interactive=False)
-
-        # **Section 2: Existing Characters**
-        gr.Markdown("---")
-        gr.Markdown("## Existing Characters")
-
-        # Dropdown to select a character for editing or deleting
-        def get_character_names():
-            characters = data_manager.get_characters()
-            return [char['name'] for char in characters]
-
-        selected_character = gr.Dropdown(
-            choices=get_character_names(),
-            label="Select a Character to Edit/Delete",
-            interactive=True
-        )
-
-        # Buttons for Edit and Delete
-        edit_button = gr.Button("Edit Selected Character")
-        delete_button = gr.Button("Delete Selected Character")
-        action_output = gr.Textbox(label="Action Status", interactive=False)
-
-        # **Section 3: Edit Character**
-        edit_section = gr.Row(visible=False)
-        with edit_section:
-            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"
-            )
-            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 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)
-
-        # **Function to Save New Character**
-        def save_character(name, traits, image_path, gender):
-            if not name.strip() or not traits.strip() or not gender:
-                return "Please enter all fields."
-            # 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."
-
-            character = {'name': name, 'traits': traits, 'gender': gender, 'image': image_path}
-            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=save_output
-        )
-
-        # **Function to Show Edit Section**
-        def show_edit(selected_name):
-            if not selected_name:
-                return "Please select a character to edit.", gr.update(visible=False)
-            # Fetch character details
-            character = next((char for char in data_manager.get_characters() if char['name'] == selected_name), None)
-            if character:
-                # Pre-fill the edit fields
-                edit_name.value = character['name']
-                edit_traits.value = ', '.join(character['traits'])
-                edit_gender.value = character['gender']
-                edit_image.value = None  # Reset image
-                return "", gr.update(visible=True)
-            else:
-                return f"Character '{selected_name}' not found.", gr.update(visible=False)
-
-        edit_button.click(
-            show_edit,
-            inputs=selected_character,
-            outputs=[action_output, edit_section]
-        )
-
-        # **Function to Update Character**
-        def update_character(original_name, new_name, new_traits, new_image_path, new_gender):
-            if not new_name.strip() or not new_traits.strip() or not new_gender:
-                return "Please enter all fields.", gr.update(visible=False)
-            # If the name has changed, check for duplicates
-            if new_name != original_name:
-                existing_names = [char['name'] for char 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.", gr.update(visible=True)
-            updated_char = {
-                'name': new_name,
-                'traits': new_traits,
-                'gender': new_gender,
-                'image': new_image_path if new_image_path else \
-                         next((char for char in data_manager.get_characters() if char['name'] == original_name), {}).get('image')
-            }
-            success = data_manager.update_character(original_name, updated_char)
-            if success:
-                # Update the dropdown choices
-                return f"Character '{new_name}' updated successfully.", gr.update(visible=False)
-            else:
-                return "Failed to update character.", gr.update(visible=True)
-
-        update_button.click(
-            update_character,
-            inputs=[selected_character, edit_name, edit_traits, edit_image, edit_gender],
-            outputs=[update_output, edit_section]
-        )
-
-        # **Function to Show Delete Confirmation**
-        def show_delete_confirmation(selected_name):
-            if not selected_name:
-                return "Please select a character to delete.", gr.update(visible=False), gr.update(visible=False)
-            return f"Are you sure you want to delete '{selected_name}'? Type 'YES' to confirm.", gr.update(visible=True), gr.update(visible=True)
-
-        delete_button.click(
-            show_delete_confirmation,
-            inputs=selected_character,
-            outputs=[action_output, delete_confirmation, confirm_delete_button]
-        )
-
-        # **Function to Perform Deletion**
-        def perform_deletion(selected_name, confirmation):
-            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:
-                return f"Character '{selected_name}' deleted successfully.", gr.update(visible=False), gr.update(visible=False)
-            else:
-                return f"Failed to delete character '{selected_name}'.", gr.update(visible=False), gr.update(visible=False)
-
-        confirm_delete_button.click(
-            perform_deletion,
-            inputs=[selected_character, delete_confirmation],
-            outputs=[action_output, delete_confirmation, confirm_delete_button]
-        )
-
-        # **Function to Cancel Deletion**
-        def cancel_deletion():
-            return "Deletion cancelled.", gr.update(visible=False), gr.update(visible=False)
-
-        cancel_delete_button.click(
-            cancel_deletion,
-            inputs=None,
-            outputs=[action_output, delete_confirmation, confirm_delete_button]
-        )
+        # 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"):
@@ -561,13 +334,11 @@ 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)
 
+            # 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()]
@@ -581,31 +352,24 @@ 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)
 
-        def save_persistent_tags_fn(tags_string):
+        def save_persistent_tags(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_fn,
-            inputs=persistent_tags_input,
-            outputs=persistent_status_output
-        )
+        save_persistent_tags_button.click(save_persistent_tags, inputs=persistent_tags_input, outputs=persistent_status_output)
 
 def main():
-    data_manager = DataManager(base_dir='data')
+    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)
+        prompt_generator_app(data_manager)
+        character_creation_app(data_manager)
+        tags_app(data_manager)
+
     demo.launch()
 
 if __name__ == "__main__":