Update app.py

app.py

@@ -6,12 +6,27 @@ from PIL import Image
 import base64
 import io
 
+# Define categories at the top so they are accessible throughout the code
+categories = [
+    ('Scene', 'scene_tags'),
+    ('Position', 'position_tags'),
+    ('Outfit', 'outfit_tags'),
+    ('Camera View/Angle', 'camera_tags'),
+    ('Concept', 'concept_tags'),
+    ('Facial Expression', 'facial_expression_tags'),
+    ('Pose', 'pose_tags'),
+    ('Additional', 'additional_tags'),
+    ('LORA', 'lora_tags')
+]
+
 class DataManager:
-    def __init__(self, characters_file='characters.json', 
-                 persistent_tags_file='persistent_tags.json', 
+    def __init__(self, characters_file='characters.json',
+                 persistent_tags_file='persistent_tags.json',
+                 category_tags_file='category_tags.json',
                  images_folder='character_images'):
         self.characters_file = characters_file
         self.persistent_tags_file = persistent_tags_file
+        self.category_tags_file = category_tags_file
         self.images_folder = images_folder
 
         # Make sure image folder exists
@@ -19,7 +34,7 @@ class DataManager:
             os.makedirs(images_folder)
 
         self.load_data()
-    
+
     def load_data(self):
         # Load characters
         if os.path.exists(self.characters_file):
@@ -27,7 +42,7 @@ class DataManager:
                 self.characters = json.load(f)
         else:
             self.characters = []
-        
+
         # Load persistent tags
         if os.path.exists(self.persistent_tags_file):
             with open(self.persistent_tags_file, 'r') as f:
@@ -35,6 +50,9 @@ class DataManager:
         else:
             self.persistent_tags = []
 
+        # Load category tags
+        self.load_category_tags()
+
     def save_characters(self):
         with open(self.characters_file, 'w') as f:
             json.dump(self.characters, f)
@@ -42,7 +60,43 @@ class DataManager:
     def save_persistent_tags(self):
         with open(self.persistent_tags_file, 'w') as f:
             json.dump(self.persistent_tags, f)
-    
+
+    def load_category_tags(self):
+        if os.path.exists(self.category_tags_file):
+            with open(self.category_tags_file, 'r') as f:
+                self.category_tags = json.load(f)
+        else:
+            self.category_tags = {}
+
+    def save_category_tags(self):
+        with open(self.category_tags_file, 'w') as f:
+            json.dump(self.category_tags, f)
+
+    def get_category_tags(self, category_var_name):
+        # Return the tags list for the given category variable name
+        return self.category_tags.get(category_var_name, [])
+
+    def set_category_tags(self, category_var_name, tags_list):
+        self.category_tags[category_var_name] = tags_list
+        self.save_category_tags()
+
+    def get_characters(self):
+        # Load character images
+        for char in self.characters:
+            image_path = char.get('image_path')
+            if image_path and os.path.exists(image_path):
+                char['image'] = image_path
+            else:
+                char['image'] = None
+        return self.characters
+
+    def set_persistent_tags(self, tags_list):
+        self.persistent_tags = tags_list
+        self.save_persistent_tags()
+
+    def get_persistent_tags(self):
+        return self.persistent_tags
+
     def add_character(self, character):
         # Save image to disk and store the filename
         image_data = character['image']  # This is base64 encoded string
@@ -66,24 +120,6 @@ class DataManager:
 
         self.characters.append(character)
         self.save_characters()
-    
-    def get_characters(self):
-        # Load character images
-        for char in self.characters:
-            image_path = char.get('image_path')
-            if image_path and os.path.exists(image_path):
-                char['image'] = image_path
-            else:
-                char['image'] = None
-        return self.characters
-    
-    def set_persistent_tags(self, tags_string):
-        if isinstance(tags_string, str):
-            self.persistent_tags = [t.strip() for t in tags_string.split(',') if t.strip()]
-            self.save_persistent_tags()
-    
-    def get_persistent_tags(self):
-        return self.persistent_tags
 
 def character_creation_app(data_manager):
     with gr.Tab("Character Creation"):
@@ -110,30 +146,27 @@ def character_creation_app(data_manager):
                 character['image'] = None
 
             data_manager.add_character(character)
-            # Clear inputs after saving
             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"):
-        # Input fields for each category
-        categories = [
-            ('Scene', 'scene_tags'),
-            ('Position', 'position_tags'),
-            ('Outfit', 'outfit_tags'),
-            ('Camera View/Angle', 'camera_tags'),
-            ('Concept', 'concept_tags'),
-            ('Additional', 'additional_tags'),
-            ('LORA', 'lora_tags')
-        ]
+        gr.Markdown("## Prompt Generator")
 
         inputs = {}
         for category_name, var_name in categories:
-            with gr.Row():
-                tag_input = gr.Textbox(label=f"{category_name} Tags (comma separated)")
-                tag_num = gr.Slider(minimum=0, maximum=10, step=1, value=1, label=f"Number of {category_name} Tags to Select")
-            inputs[f"{var_name}_input"] = tag_input
+            tags_list = data_manager.get_category_tags(var_name)
+            tags_string = ', '.join(tags_list)
+            max_tags = len(tags_list)
+            if max_tags == 0:
+                default_value = 0
+            else:
+                default_value = min(1, max_tags)
+            with gr.Box():
+                gr.Markdown(f"### {category_name}")
+                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[f"{var_name}_num"] = tag_num
 
         # For Character Selection
@@ -166,16 +199,15 @@ def prompt_generator_app(data_manager):
         generate_button = gr.Button("Generate Prompt")
         prompt_output = gr.Textbox(label="Generated Prompt", lines=5)
 
-        def generate_prompt(*args):
+        def generate_prompt(*args, data_manager=data_manager, categories=categories):
             arg_idx = 0
 
             prompt_tags = []
             for category_name, var_name in categories:
-                tags_input = args[arg_idx]
-                tags_num = args[arg_idx + 1]
-                arg_idx += 2
+                tags_list = data_manager.get_category_tags(var_name)
+                tags_num = args[arg_idx]
+                arg_idx += 1
 
-                tags_list = [tag.strip() for tag in tags_input.split(',')] if tags_input else []
                 if tags_list and tags_num > 0:
                     selected_tags = random.sample(tags_list, min(len(tags_list), int(tags_num)))
                     prompt_tags.extend(selected_tags)
@@ -255,31 +287,56 @@ def prompt_generator_app(data_manager):
 
         # Prepare the list of inputs for the generate_prompt function
         inputs_list = []
-        for var_name in inputs:
-            inputs_list.append(inputs[var_name])
+        for category_name, var_name in categories:
+            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)
 
-def settings_app(data_manager):
-    with gr.Tab("Settings"):
-        persistent_tags_input = gr.Textbox(label="Persistent Tags (comma separated)")
+def tags_app(data_manager):
+    with gr.Tab("Tags"):
+        gr.Markdown("## Edit Tags for Each Category")
+
+        for category_name, var_name in categories:
+            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)
+            save_button = gr.Button(f"Save {category_name} Tags")
+            status_output = gr.Textbox(label="", interactive=False)
+
+            # Capture var_name and category_name in the closure
+            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()]
+                    data_manager.set_category_tags(var_name, tags_list)
+                    return f"{category_name} tags saved successfully."
+                return fn
+
+            save_fn = make_save_category_tags_fn(var_name, category_name)
+            save_button.click(save_fn, inputs=tag_input, outputs=status_output)
+
+        # 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)
         save_persistent_tags_button = gr.Button("Save Persistent Tags")
-        status_output = gr.Textbox(label="Status", interactive=False)
+        persistent_status_output = gr.Textbox(label="", interactive=False)
 
         def save_persistent_tags(tags_string):
-            data_manager.set_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, inputs=[persistent_tags_input], outputs=status_output)
+        save_persistent_tags_button.click(save_persistent_tags, inputs=persistent_tags_input, outputs=persistent_status_output)
 
 def main():
     data_manager = DataManager()
     with gr.Blocks() as demo:
         prompt_generator_app(data_manager)
         character_creation_app(data_manager)
-        settings_app(data_manager)
+        tags_app(data_manager)
 
     demo.launch()