Update app.py

app.py

@@ -134,18 +134,76 @@ def prompt_generator_app(data_manager):
         # For Character Selection
         with gr.Group():
             gr.Markdown("### Character Selection")
-            characters = data_manager.get_characters()
-            character_options = [char['name'] for char in characters]
-            character_select = gr.CheckboxGroup(choices=character_options, label="Select Characters")
-            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)")
-        
+
+            # Create placeholders for dynamic components
+            character_select = gr.State([])
+            selected_character_names = gr.State([])
+
+            characters_display = gr.HTML()
+            refresh_characters_button = gr.Button("Refresh Character List")
+
         # Number of people in the scene
         num_people = gr.Slider(minimum=0, maximum=10, step=1, value=1, label="Number of People in the Scene")
 
+        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)")
+
         generate_button = gr.Button("Generate Prompt")
         prompt_output = gr.Textbox(label="Generated Prompt", lines=5)
 
+        # Function to display characters with images and selection checkboxes
+        def display_characters():
+            characters = data_manager.get_characters()
+            if not characters:
+                return "", []
+
+            # Build HTML content for characters display
+            html_content = "<div style='display: flex; flex-wrap: wrap;'>"
+            for idx, char in enumerate(characters):
+                image_path = char.get('image_path', '')
+                image_url = ''
+                if image_path and os.path.exists(image_path):
+                    with open(image_path, "rb") as img_file:
+                        image_data = base64.b64encode(img_file.read()).decode('utf-8')
+                        image_url = f"data:image/png;base64,{image_data}"
+                else:
+                    image_url = ""
+
+                # Each character card
+                html_content += f"""
+                <div style='margin: 10px; text-align: center; width: 150px;'>
+                    <img src='{image_url}' style='width: 100%; height: auto;'/>
+                    <p>{char['name']}</p>
+                    <input type='checkbox' name='character_select' value='{char['name']}' onchange='updateSelectedCharacters()'/>
+                </div>
+                """
+            html_content += "</div>"
+
+            # JavaScript function to update selected characters
+            html_content += """
+            <script>
+            function updateSelectedCharacters() {
+                let selected = [];
+                let checkboxes = document.getElementsByName('character_select');
+                for (let i = 0; i < checkboxes.length; i++) {
+                    if (checkboxes[i].checked) {
+                        selected.push(checkboxes[i].value);
+                    }
+                }
+                gradioApp().querySelector('#selected_character_names').value = JSON.stringify(selected);
+            }
+            </script>
+            """
+
+            return html_content, []
+
+        # Display characters on load and when refresh is clicked
+        characters_display.update(display_characters)
+        refresh_characters_button.click(display_characters, outputs=characters_display)
+
+        # Hidden input to store selected character names
+        hidden_selected_characters = gr.Textbox(visible=False, elem_id='selected_character_names')
+
         def generate_prompt(*args):
             # args correspond to inputs in the order they are defined
             # Need to map args to variables
@@ -163,21 +221,23 @@ def prompt_generator_app(data_manager):
                     prompt_tags.extend(selected_tags)
             
             # Handle Characters
-            selected_characters = args[arg_idx]
+            selected_characters_json = args[arg_idx]
             random_chars = args[arg_idx + 1]
             num_random_chars = args[arg_idx + 2]
             num_people_in_scene = args[arg_idx + 3]
 
-            characters = data_manager.get_characters()
-
             arg_idx += 4
 
+            selected_character_names_list = json.loads(selected_characters_json) if selected_characters_json else []
+
+            characters = data_manager.get_characters()
+
             selected_chars = []
             if random_chars:
                 num = min(len(characters), int(num_random_chars))
                 selected_chars = random.sample(characters, num)
             else:
-                selected_chars = [char for char in characters if char['name'] in selected_characters]
+                selected_chars = [char for char in characters if char['name'] in selected_character_names_list]
 
             # Adjust the number of people in the scene
             # You can modify how additional characters are added if needed
@@ -201,7 +261,7 @@ def prompt_generator_app(data_manager):
         inputs_list = []
         for var_name in inputs:
             inputs_list.append(inputs[var_name])
-        inputs_list.extend([character_select, random_characters, num_characters, num_people])
+        inputs_list.extend([hidden_selected_characters, random_characters, num_characters, num_people])
 
         generate_button.click(generate_prompt, inputs=inputs_list, outputs=prompt_output)