import gradio as gr import random import json import os DATA_DIR = "/data" IMAGES_DIR = "/images" DATA_FILE = os.path.join(DATA_DIR, "saved_data.json") def load_data(): if os.path.exists(DATA_FILE): with open(DATA_FILE, 'r') as f: return json.load(f) return { "scene_tags": [], "position_tags": [], "outfit_tags": [], "camera_tags": [], "concept_tags": [], "lora_tags": [], "characters": {} } def save_data(data): os.makedirs(DATA_DIR, exist_ok=True) with open(DATA_FILE, 'w') as f: json.dump(data, f) def save_character_image(name, image): os.makedirs(IMAGES_DIR, exist_ok=True) image_path = os.path.join(IMAGES_DIR, f"{name}.png") image.save(image_path) 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(','), "outfit": outfit_tags.split(','), "camera": camera_tags.split(','), "concept": concept_tags.split(','), } all_tags = {k: [tag.strip() for tag in v if tag.strip()] for k, v in all_tags.items()} character_prompts = [ f"{char_name}, " + ", ".join(random.sample(data["characters"][char_name]["traits"], min(tag_counts["character"], len(data["characters"][char_name]["traits"])))) for char_name in selected_characters if char_name in data["characters"] ] selected_tags = [ f"{tag}:{random.uniform(0.8, 1.2):.2f}" for category, tags in all_tags.items() 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") prompt_parts = character_prompts + selected_tags random.shuffle(prompt_parts) main_prompt = ", ".join(prompt_parts) lora_list = [lora.strip() for lora in lora_tags.split(',') if lora.strip()] lora_prompt = " ".join(f"" for lora in lora_list) 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" return f"{main_prompt}, {fixed_tags} {lora_prompt}".strip() def update_data(data, key, value): data[key] = list(set(data[key] + [v.strip() for v in value.split(',') if v.strip()])) 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 Character Creation") 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) char_images = [char_data["image"] for char_data in data["characters"].values() if char_data["image"]] gr.Gallery(value=char_images, label="Character Images", show_label=True, elem_id="char_gallery", columns=2, rows=2, height="auto") 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=True, elem_id="char_gallery", columns=2, rows=2, height="auto") def update_and_generate(*args): nonlocal data 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, 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(): char_images = [char_data["image"] for char_data in data["characters"].values() if char_data["image"]] return gr.Gallery(value=char_images) 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__": demo = create_ui() demo.launch()