Update app.py

app.py

@@ -1,5 +1,33 @@
 import gradio as gr
 import random
+import json
+from huggingface_hub import HfApi, hf_hub_download, upload_file
+
+REPO_ID = "throaway2854/promptGen"  # Replace with your actual Space name
+FILE_NAME = "saved_data.json"
+
+def load_data():
+    try:
+        file_path = hf_hub_download(repo_id=REPO_ID, filename=FILE_NAME)
+        with open(file_path, 'r') as f:
+            return json.load(f)
+    except:
+        return {
+            "scene_tags": [], "position_tags": [], "outfit_tags": [],
+            "camera_tags": [], "concept_tags": [], "lora_tags": [],
+            "characters": {}
+        }
+
+def save_data(data):
+    with open(FILE_NAME, 'w') as f:
+        json.dump(data, f)
+    api = HfApi()
+    api.upload_file(
+        path_or_fileobj=FILE_NAME,
+        path_in_repo=FILE_NAME,
+        repo_id=REPO_ID,
+        repo_type="space"
+    )
 
 def generate_prompt(scene_tags, num_people, position_tags, characters, outfit_tags, camera_tags, concept_tags, lora_tags, num_tags=3):
     all_tags = {
@@ -10,67 +38,85 @@ def generate_prompt(scene_tags, num_people, position_tags, characters, outfit_ta
         "concept": concept_tags.split(','),
     }
     
-    # Remove empty strings and strip whitespace
     for category in all_tags:
         all_tags[category] = [tag.strip() for tag in all_tags[category] if tag.strip()]
     
-    # Process characters
     character_list = [char.strip() for char in characters.split(';') if char.strip()]
     character_prompts = []
-    for i, character in enumerate(character_list):
+    for character in character_list:
         char_tags = character.split(',')
         char_name = char_tags[0].strip()
         char_traits = [tag.strip() for tag in char_tags[1:] if tag.strip()]
         char_prompt = f"{char_name}, " + ", ".join(random.sample(char_traits, min(num_tags, len(char_traits))))
         character_prompts.append(char_prompt)
     
-    # Generate main prompt
     selected_tags = []
     for category, tags in all_tags.items():
         if tags:
             selected_tags.extend(f"{tag}:{random.uniform(0.8, 1.2):.2f}" for tag in random.sample(tags, min(num_tags, len(tags))))
     
-    # Add number of people
     if num_people.strip():
         selected_tags.append(f"{num_people} people:1.1")
     
-    # Combine all parts of the prompt
     prompt_parts = character_prompts + selected_tags
     random.shuffle(prompt_parts)
     main_prompt = ", ".join(prompt_parts)
     
-    # Add LORA tags
     lora_list = [lora.strip() for lora in lora_tags.split(',') if lora.strip()]
     lora_prompt = " ".join(f"<lora:{lora}:1>" for lora in lora_list)
     
-    # Add fixed tags at the end
     fixed_tags = "source_anime, score_9, score_8_up, score_7_up, masterpiece, best quality, very aesthetic, absurdres, anime artwork, anime style, vibrant, studio anime, highly detailed"
     
     full_prompt = f"{main_prompt}, {fixed_tags} {lora_prompt}".strip()
     return full_prompt
 
+def update_data(data, key, value):
+    if isinstance(data[key], list):
+        data[key] = list(set(data[key] + [v.strip() for v in value.split(',') if v.strip()]))
+    elif isinstance(data[key], dict):
+        for char in [c.strip() for c in value.split(';') if c.strip()]:
+            char_name, *char_traits = [t.strip() for t in char.split(',')]
+            if char_name:
+                data[key][char_name] = list(set(data[key].get(char_name, []) + char_traits))
+    save_data(data)
+    return data
+
 def create_ui():
+    data = load_data()
+
     with gr.Blocks() as demo:
-        gr.Markdown("# Advanced Pony SDXL Prompt Generator")
+        gr.Markdown("# Advanced Pony SDXL Prompt Generator with Persistent Storage")
         
         with gr.Row():
             with gr.Column():
-                scene_input = gr.Textbox(label="Scene Tags (comma-separated)")
+                scene_input = gr.Textbox(label="Scene Tags (comma-separated)", value=", ".join(data["scene_tags"]))
                 num_people_input = gr.Textbox(label="Number of People")
-                position_input = gr.Textbox(label="Position Tags (comma-separated)")
-                characters_input = gr.Textbox(label="Characters (Format: Name, Trait1, Trait2; Name2, Trait1, Trait2)")
-                outfit_input = gr.Textbox(label="Outfit Tags (comma-separated)")
-                camera_input = gr.Textbox(label="Camera View/Angle Tags (comma-separated)")
-                concept_input = gr.Textbox(label="Concept Tags (comma-separated)")
-                lora_input = gr.Textbox(label="LORA Tags (comma-separated)")
+                position_input = gr.Textbox(label="Position Tags (comma-separated)", value=", ".join(data["position_tags"]))
+                characters_input = gr.Textbox(label="Characters (Format: Name, Trait1, Trait2; Name2, Trait1, Trait2)", 
+                                              value="; ".join([f"{name}, {', '.join(traits)}" for name, traits in data["characters"].items()]))
+                outfit_input = gr.Textbox(label="Outfit Tags (comma-separated)", value=", ".join(data["outfit_tags"]))
+                camera_input = gr.Textbox(label="Camera View/Angle Tags (comma-separated)", value=", ".join(data["camera_tags"]))
+                concept_input = gr.Textbox(label="Concept Tags (comma-separated)", value=", ".join(data["concept_tags"]))
+                lora_input = gr.Textbox(label="LORA Tags (comma-separated)", value=", ".join(data["lora_tags"]))
                 num_tags_slider = gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Number of tags per category")
                 generate_button = gr.Button("Generate Prompt")
             
             with gr.Column():
                 output = gr.Textbox(label="Generated Prompt", lines=5)
         
+        def update_and_generate(*args):
+            nonlocal data
+            data = update_data(data, "scene_tags", args[0])
+            data = update_data(data, "position_tags", args[2])
+            data = update_data(data, "characters", args[3])
+            data = update_data(data, "outfit_tags", args[4])
+            data = update_data(data, "camera_tags", args[5])
+            data = update_data(data, "concept_tags", args[6])
+            data = update_data(data, "lora_tags", args[7])
+            return generate_prompt(*args)
+
         generate_button.click(
-            generate_prompt,
+            update_and_generate,
             inputs=[scene_input, num_people_input, position_input, characters_input, outfit_input, camera_input, concept_input, lora_input, num_tags_slider],
             outputs=[output]
         )