Update app.py

app.py

@@ -7,7 +7,9 @@ import base64
 import io
 
 class DataManager:
-    def __init__(self, characters_file='characters.json', persistent_tags_file='persistent_tags.json', images_folder='character_images'):
+    def __init__(self, characters_file='characters.json', 
+                 persistent_tags_file='persistent_tags.json', 
+                 images_folder='character_images'):
         self.characters_file = characters_file
         self.persistent_tags_file = persistent_tags_file
         self.images_folder = images_folder
@@ -138,62 +140,28 @@ def prompt_generator_app(data_manager):
         with gr.Group():
             gr.Markdown("### Character Selection")
 
-            # Initialize selected_characters_state
-            selected_characters_state = gr.State([])
-
-            # Add the Refresh Character List button
-            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
 
-            # Create a placeholder for the character_select component
-            character_select = gr.Gallery(label="Select Characters", columns=3, height='auto')
+            character_options = get_character_options()
+            character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters", interactive=True)
 
-            # Load and display characters
-            def load_characters():
-                characters = data_manager.get_characters()
-                if not characters:
-                    return [], []
-                else:
-                    # Build character images and labels
-                    character_images = []
-                    for idx, char in enumerate(characters):
-                        image_path = char.get('image_path', '')
-                        if image_path and os.path.exists(image_path):
-                            image_url = image_path
-                        else:
-                            image_url = None  # Or a placeholder image
-                        # Adding the character's name and gender as a caption
-                        caption = f"{char['name']} ({char['gender']})"
-                        character_images.append((image_url, caption))
-                    return character_images, characters
-
-            # Function to update selected characters
-            def update_selected_characters(selected_indices):
-                characters = data_manager.get_characters()
-                selected_indices = selected_indices or []
-                selected_chars = [characters[i] for i in selected_indices]
-                return selected_chars
+            refresh_characters_button = gr.Button("Refresh Character List")
 
-            # Function to refresh characters
             def refresh_characters():
-                character_images, characters = load_characters()
-                # Update character_select with new images
-                return gr.update(value=character_images)
-
-            # Load initial characters
-            character_images, characters = load_characters()
-            if character_images:
-                character_select.value = character_images
-            else:
-                gr.Markdown("No characters available. Please add characters in the Character Creation tab.")
-
-            # Update selected characters when selection changes
-            character_select.select(fn=update_selected_characters, inputs=None, outputs=selected_characters_state)
+                new_options = get_character_options()
+                return gr.update(choices=new_options)
 
-            # Update character gallery when refresh button is clicked
-            refresh_characters_button.click(fn=refresh_characters, inputs=None, 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=0, 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)
@@ -213,7 +181,7 @@ def prompt_generator_app(data_manager):
                     prompt_tags.extend(selected_tags)
 
             # Handle Characters
-            selected_chars = args[arg_idx]
+            selected_character_options = args[arg_idx]
             random_chars = args[arg_idx + 1]
             num_random_chars = args[arg_idx + 2]
 
@@ -225,8 +193,14 @@ def prompt_generator_app(data_manager):
                 num = min(len(characters), int(num_random_chars))
                 selected_chars = random.sample(characters, num)
             else:
-                # selected_chars is already set from the selected characters in the gallery
-                selected_chars = selected_chars or []
+                # 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')
@@ -283,8 +257,8 @@ def prompt_generator_app(data_manager):
         inputs_list = []
         for var_name in inputs:
             inputs_list.append(inputs[var_name])
-        # Add selected characters state to inputs
-        inputs_list.extend([selected_characters_state, random_characters, num_characters])
+        # 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)