Update app.py

app.py

@@ -1,14 +1,16 @@
 import gradio as gr
 import random
 import json
+import os
 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"
+DATA_FILE = "saved_data.json"
+IMAGES_DIR = "character_images"
 
 def load_data():
     try:
-        file_path = hf_hub_download(repo_id=REPO_ID, filename=FILE_NAME)
+        file_path = hf_hub_download(repo_id=REPO_ID, filename=DATA_FILE)
         with open(file_path, 'r') as f:
             return json.load(f)
     except:
@@ -19,17 +21,31 @@ def load_data():
         }
 
 def save_data(data):
-    with open(FILE_NAME, 'w') as f:
+    with open(DATA_FILE, 'w') as f:
         json.dump(data, f)
     api = HfApi()
     api.upload_file(
-        path_or_fileobj=FILE_NAME,
-        path_in_repo=FILE_NAME,
+        path_or_fileobj=DATA_FILE,
+        path_in_repo=DATA_FILE,
         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):
+def save_character_image(name, image):
+    if not os.path.exists(IMAGES_DIR):
+        os.makedirs(IMAGES_DIR)
+    image_path = os.path.join(IMAGES_DIR, f"{name}.png")
+    image.save(image_path)
+    api = HfApi()
+    api.upload_file(
+        path_or_fileobj=image_path,
+        path_in_repo=f"{IMAGES_DIR}/{name}.png",
+        repo_id=REPO_ID,
+        repo_type="space"
+    )
+    return image_path
+
+def generate_prompt(scene_tags, num_people, position_tags, selected_characters, outfit_tags, camera_tags, concept_tags, lora_tags, tag_counts, data):
     all_tags = {
         "scene": scene_tags.split(','),
         "position": position_tags.split(','),
@@ -41,19 +57,17 @@ def generate_prompt(scene_tags, num_people, position_tags, characters, outfit_ta
     for category in all_tags:
         all_tags[category] = [tag.strip() for tag in all_tags[category] if tag.strip()]
     
-    character_list = [char.strip() for char in characters.split(';') if char.strip()]
     character_prompts = []
-    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)
+    for char_name in selected_characters:
+        if char_name in data["characters"]:
+            char_traits = data["characters"][char_name]["traits"]
+            char_prompt = f"{char_name}, " + ", ".join(random.sample(char_traits, min(tag_counts["character"], len(char_traits))))
+            character_prompts.append(char_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))))
+            selected_tags.extend(f"{tag}:{random.uniform(0.8, 1.2):.2f}" for tag in random.sample(tags, min(tag_counts[category], len(tags))))
     
     if num_people.strip():
         selected_tags.append(f"{num_people} people:1.1")
@@ -73,54 +87,108 @@ def generate_prompt(scene_tags, num_people, position_tags, characters, outfit_ta
 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_character(name, traits, image, data):
+    if name:
+        data["characters"][name] = {
+            "traits": [trait.strip() for trait in traits.split(',') if trait.strip()],
+            "image": save_character_image(name, image) if image else None
+        }
+        save_data(data)
+    return data, gr.update(choices=list(data["characters"].keys()))
+
 def create_ui():
     data = load_data()
 
     with gr.Blocks() as demo:
-        gr.Markdown("# Advanced Pony SDXL Prompt Generator with Persistent Storage")
+        gr.Markdown("# Advanced Pony SDXL Prompt Generator with Character Creation")
         
-        with gr.Row():
-            with gr.Column():
-                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)", 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.Tabs():
+            with gr.TabItem("Prompt Generator"):
+                with gr.Row():
+                    with gr.Column():
+                        scene_input = gr.Textbox(label="Scene Tags (comma-separated)", value=", ".join(data["scene_tags"]))
+                        scene_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of scene tags")
+                        
+                        num_people_input = gr.Textbox(label="Number of People")
+                        
+                        position_input = gr.Textbox(label="Position Tags (comma-separated)", value=", ".join(data["position_tags"]))
+                        position_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of position tags")
+                        
+                        character_select = gr.CheckboxGroup(label="Select Characters", choices=list(data["characters"].keys()))
+                        character_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of character traits")
+                        
+                        outfit_input = gr.Textbox(label="Outfit Tags (comma-separated)", value=", ".join(data["outfit_tags"]))
+                        outfit_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of outfit tags")
+                        
+                        camera_input = gr.Textbox(label="Camera View/Angle Tags (comma-separated)", value=", ".join(data["camera_tags"]))
+                        camera_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of camera tags")
+                        
+                        concept_input = gr.Textbox(label="Concept Tags (comma-separated)", value=", ".join(data["concept_tags"]))
+                        concept_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of concept tags")
+                        
+                        lora_input = gr.Textbox(label="LORA Tags (comma-separated)", value=", ".join(data["lora_tags"]))
+                        lora_count = gr.Slider(minimum=1, maximum=10, step=1, value=3, label="Number of LORA tags")
+                        
+                        generate_button = gr.Button("Generate Prompt")
+                    
+                    with gr.Column():
+                        output = gr.Textbox(label="Generated Prompt", lines=5)
+                        
+                        with gr.Row():
+                            for char_name, char_data in data["characters"].items():
+                                if char_data["image"]:
+                                    gr.Image(value=char_data["image"], label=char_name, height=100, width=100)
             
-            with gr.Column():
-                output = gr.Textbox(label="Generated Prompt", lines=5)
-        
+            with gr.TabItem("Character Creation"):
+                with gr.Row():
+                    with gr.Column():
+                        char_name_input = gr.Textbox(label="Character Name")
+                        char_traits_input = gr.Textbox(label="Character Traits (comma-separated)")
+                        char_image_input = gr.Image(label="Character Image", type="pil")
+                        create_char_button = gr.Button("Create/Update Character")
+                    
+                    with gr.Column():
+                        char_gallery = gr.Gallery(label="Existing Characters", show_label=False, elem_id="char_gallery").style(grid=2, height="auto")
+
         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)
+            scene_tags, num_people, position_tags, selected_characters, outfit_tags, camera_tags, concept_tags, lora_tags, *tag_counts = args
+            data = update_data(data, "scene_tags", scene_tags)
+            data = update_data(data, "position_tags", position_tags)
+            data = update_data(data, "outfit_tags", outfit_tags)
+            data = update_data(data, "camera_tags", camera_tags)
+            data = update_data(data, "concept_tags", concept_tags)
+            data = update_data(data, "lora_tags", lora_tags)
+            
+            tag_count_dict = {
+                "scene": tag_counts[0], "position": tag_counts[1], "character": tag_counts[2],
+                "outfit": tag_counts[3], "camera": tag_counts[4], "concept": tag_counts[5], "lora": tag_counts[6]
+            }
+            
+            return generate_prompt(scene_tags, num_people, position_tags, selected_characters, outfit_tags, camera_tags, concept_tags, lora_tags, tag_count_dict, data)
 
         generate_button.click(
             update_and_generate,
-            inputs=[scene_input, num_people_input, position_input, characters_input, outfit_input, camera_input, concept_input, lora_input, num_tags_slider],
+            inputs=[scene_input, num_people_input, position_input, character_select, outfit_input, camera_input, concept_input, lora_input,
+                    scene_count, position_count, character_count, outfit_count, camera_count, concept_count, lora_count],
             outputs=[output]
         )
-    
+
+        def update_char_gallery():
+            return gr.Gallery(value=[char_data["image"] for char_data in data["characters"].values() if char_data["image"]])
+
+        create_char_button.click(
+            create_character,
+            inputs=[char_name_input, char_traits_input, char_image_input],
+            outputs=[gr.State(data), character_select]
+        ).then(
+            update_char_gallery,
+            outputs=[char_gallery]
+        )
+
     return demo
 
 if __name__ == "__main__":