import gradio as gr import json import logging import torch from PIL import Image import spaces from diffusers import DiffusionPipeline import copy import random import time from mod import (models, clear_cache, get_repo_safetensors, change_base_model, description_ui, num_loras, compose_lora_json, is_valid_lora, fuse_loras, get_trigger_word) from flux import (search_civitai_lora, select_civitai_lora, search_civitai_lora_json, download_my_lora, get_all_lora_tupled_list, apply_lora_prompt, update_loras) from tagger.tagger import predict_tags_wd, compose_prompt_to_copy from tagger.fl2cog import predict_tags_fl2_cog from tagger.fl2flux import predict_tags_fl2_flux # Load LoRAs from JSON file with open('loras.json', 'r') as f: loras = json.load(f) # Initialize the base model base_model = models[0] pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16) MAX_SEED = 2**32-1 class calculateDuration: def __init__(self, activity_name=""): self.activity_name = activity_name def __enter__(self): self.start_time = time.time() return self def __exit__(self, exc_type, exc_value, traceback): self.end_time = time.time() self.elapsed_time = self.end_time - self.start_time if self.activity_name: print(f"Elapsed time for {self.activity_name}: {self.elapsed_time:.6f} seconds") else: print(f"Elapsed time: {self.elapsed_time:.6f} seconds") def update_selection(evt: gr.SelectData, width, height): selected_lora = loras[evt.index] new_placeholder = f"Type a prompt for {selected_lora['title']}" lora_repo = selected_lora["repo"] updated_text = f"### Selected: [{lora_repo}](https://huggingface.co/{lora_repo}) ✨" if "aspect" in selected_lora: if selected_lora["aspect"] == "portrait": width = 768 height = 1024 elif selected_lora["aspect"] == "landscape": width = 1024 height = 768 return ( gr.update(placeholder=new_placeholder), updated_text, evt.index, width, height, ) @spaces.GPU(duration=70) def generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress): pipe.to("cuda") generator = torch.Generator(device="cuda").manual_seed(seed) with calculateDuration("Generating image"): # Generate image image = pipe( prompt=f"{prompt} {trigger_word}", num_inference_steps=steps, guidance_scale=cfg_scale, width=width, height=height, generator=generator, joint_attention_kwargs={"scale": lora_scale}, ).images[0] return image def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, lora_json, progress=gr.Progress(track_tqdm=True)): if selected_index is None and not is_valid_lora(lora_json): gr.Info("LoRA isn't selected.") # raise gr.Error("You must select a LoRA before proceeding.") if is_valid_lora(lora_json): with calculateDuration("Loading LoRA weights"): fuse_loras(pipe, lora_json) trigger_word = get_trigger_word(lora_json) elif selected_index is not None: selected_lora = loras[selected_index] lora_path = selected_lora["repo"] trigger_word = selected_lora["trigger_word"] # Load LoRA weights with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"): if "weights" in selected_lora: pipe.load_lora_weights(lora_path, weight_name=selected_lora["weights"]) else: pipe.load_lora_weights(lora_path) # Set random seed for reproducibility with calculateDuration("Randomizing seed"): if randomize_seed: seed = random.randint(0, MAX_SEED) image = generate_image(prompt, trigger_word, steps, seed, cfg_scale, width, height, lora_scale, progress) pipe.to("cpu") if selected_index is not None or is_valid_lora(lora_json): pipe.unload_lora_weights() clear_cache() return image, seed run_lora.zerogpu = True css = ''' #gen_btn{height: 100%} #title{text-align: center} #title h1{font-size: 3em; display:inline-flex; align-items:center} #title img{width: 100px; margin-right: 0.5em} #gallery .grid-wrap{height: 10vh} ''' with gr.Blocks(theme=gr.themes.Soft(), fill_width=True, css=css) as app: with gr.Tab("FLUX LoRA the Explorer"): title = gr.HTML( """

LoRAFLUX LoRA the Explorer Mod

""", elem_id="title", ) selected_index = gr.State(None) with gr.Row(): with gr.Column(scale=3): with gr.Group(): with gr.Accordion("Generate Prompt from Image", open=False): tagger_image = gr.Image(label="Input image", type="pil", sources=["upload", "clipboard"], height=256) with gr.Accordion(label="Advanced options", open=False): tagger_general_threshold = gr.Slider(label="Threshold", minimum=0.0, maximum=1.0, value=0.3, step=0.01, interactive=True) tagger_character_threshold = gr.Slider(label="Character threshold", minimum=0.0, maximum=1.0, value=0.8, step=0.01, interactive=True) neg_prompt = gr.Text(label="Negative Prompt", lines=1, max_lines=8, placeholder="", visible=False) v2_character = gr.Textbox(label="Character", placeholder="hatsune miku", scale=2, visible=False) v2_series = gr.Textbox(label="Series", placeholder="vocaloid", scale=2, visible=False) v2_copy = gr.Button(value="Copy to clipboard", size="sm", interactive=False, visible=False) tagger_algorithms = gr.CheckboxGroup(["Use WD Tagger", "Use CogFlorence-2.1-Large", "Use Florence-2-Flux"], label="Algorithms", value=["Use WD Tagger"]) tagger_generate_from_image = gr.Button(value="Generate Prompt from Image") prompt = gr.Textbox(label="Prompt", lines=1, placeholder="Type a prompt after selecting a LoRA") with gr.Column(scale=1, elem_id="gen_column"): generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn") with gr.Row(): with gr.Column(scale=3): selected_info = gr.Markdown("") gallery = gr.Gallery( [(item["image"], item["title"]) for item in loras], label="LoRA Gallery", allow_preview=False, columns=3, elem_id="gallery" ) with gr.Column(scale=4): result = gr.Image(label="Generated Image") with gr.Row(): with gr.Accordion("Advanced Settings", open=False): with gr.Column(): with gr.Row(): model_name = gr.Dropdown(label="Base Model", choices=models, value=models[0], allow_custom_value=True) with gr.Row(): cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, step=0.5, value=3.5) steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=28) with gr.Row(): width = gr.Slider(label="Width", minimum=256, maximum=1536, step=64, value=1024) height = gr.Slider(label="Height", minimum=256, maximum=1536, step=64, value=1024) with gr.Row(): randomize_seed = gr.Checkbox(True, label="Randomize seed") seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, randomize=True) lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, step=0.01, value=0.95) with gr.Column(): lora_repo_json = gr.JSON(value=[{}] * num_loras, visible=False) lora_repo = [None] * num_loras lora_weights = [None] * num_loras lora_trigger = [None] * num_loras lora_wt = [None] * num_loras lora_info = [None] * num_loras lora_copy = [None] * num_loras lora_md = [None] * num_loras lora_num = [None] * num_loras for i in range(num_loras): with gr.Group(): with gr.Row(): lora_repo[i] = gr.Dropdown(label=f"LoRA {int(i+1)} Repo", choices=get_all_lora_tupled_list(), info="Input LoRA Repo ID", value="", allow_custom_value=True) lora_weights[i] = gr.Dropdown(label=f"LoRA {int(i+1)} Filename", choices=[], info="Optional", value="", allow_custom_value=True) lora_trigger[i] = gr.Textbox(label=f"LoRA {int(i+1)} Trigger Prompt", value="") lora_wt[i] = gr.Slider(label=f"LoRA {int(i+1)} Scale", minimum=-2, maximum=2, step=0.01, value=1.00) with gr.Row(): lora_info[i] = gr.Textbox(label="", info="Example of prompt:", value="", show_copy_button=True, interactive=False, visible=False) lora_copy[i] = gr.Button(value="Copy example to prompt", visible=False) lora_md[i] = gr.Markdown(value="", visible=False) lora_num[i] = gr.Number(i, visible=False) with gr.Accordion("From URL", open=True, visible=True): with gr.Row(): lora_search_civitai_query = gr.Textbox(label="Query", placeholder="flux", lines=1) lora_search_civitai_basemodel = gr.CheckboxGroup(label="Search LoRA for", choices=["Flux.1 D", "Flux.1 S"], value=["Flux.1 D", "Flux.1 S"]) lora_search_civitai_submit = gr.Button("Search on Civitai") lora_search_civitai_result = gr.Dropdown(label="Search Results", choices=[("", "")], value="", allow_custom_value=True, visible=False) lora_search_civitai_json = gr.JSON(value={}, visible=False) lora_search_civitai_desc = gr.Markdown(value="", visible=False) lora_download_url = gr.Textbox(label="URL", placeholder="http://...my_lora_url.safetensors", lines=1) lora_download = gr.Button("Get and set LoRA") gallery.select( update_selection, inputs=[width, height], outputs=[prompt, selected_info, selected_index, width, height] ) gr.on( triggers=[generate_button.click, prompt.submit], fn=run_lora, inputs=[prompt, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, lora_repo_json], # outputs=[result, seed] ) model_name.change(change_base_model, [model_name], None) gr.on( triggers=[lora_search_civitai_submit.click, lora_search_civitai_query.submit], fn=search_civitai_lora, inputs=[lora_search_civitai_query, lora_search_civitai_basemodel], outputs=[lora_search_civitai_result, lora_search_civitai_desc, lora_search_civitai_submit, lora_search_civitai_query], scroll_to_output=True, queue=True, show_api=False, ) lora_search_civitai_json.change(search_civitai_lora_json, [lora_search_civitai_query, lora_search_civitai_basemodel], [lora_search_civitai_json], queue=True, show_api=True) # fn for api lora_search_civitai_result.change(select_civitai_lora, [lora_search_civitai_result], [lora_download_url, lora_search_civitai_desc], scroll_to_output=True, queue=False, show_api=False) gr.on( triggers=[lora_download.click, lora_download_url.submit], fn=download_my_lora, inputs=[lora_download_url, lora_repo[0]], outputs=[lora_repo[0]], scroll_to_output=True, queue=True, show_api=False, ) for i, l in enumerate(lora_repo): gr.on( triggers=[lora_repo[i].change, lora_wt[i].change], fn=update_loras, inputs=[prompt, lora_repo[i], lora_wt[i]], outputs=[prompt, lora_repo[i], lora_wt[i], lora_info[i], lora_copy[i], lora_md[i]], queue=False, trigger_mode="once", show_api=False, ).success(get_repo_safetensors, [lora_repo[i]], [lora_weights[i]], queue=False, show_api=False ).success(apply_lora_prompt, [lora_info[i]], [lora_trigger[i]], queue=False, show_api=False ).success(compose_lora_json, [lora_repo_json, lora_num[i], lora_repo[i], lora_wt[i], lora_weights[i], lora_trigger[i]], [lora_repo_json], queue=False, show_api=False) tagger_generate_from_image.click( lambda: ("", "", ""), None, [v2_series, v2_character, prompt], queue=False, show_api=False, ).success( predict_tags_wd, [tagger_image, prompt, tagger_algorithms, tagger_general_threshold, tagger_character_threshold], [v2_series, v2_character, prompt, v2_copy], show_api=False, ).success( predict_tags_fl2_flux, [tagger_image, prompt, tagger_algorithms], [prompt], show_api=False, ).success( predict_tags_fl2_cog, [tagger_image, prompt, tagger_algorithms], [prompt], show_api=False, ).success( compose_prompt_to_copy, [v2_character, v2_series, prompt], [prompt], queue=False, show_api=False, ) with gr.Tab("FLUX Prompt Generator"): from prompt import (PromptGenerator, HuggingFaceInferenceNode, florence_caption, ARTFORM, PHOTO_TYPE, BODY_TYPES, DEFAULT_TAGS, ROLES, HAIRSTYLES, ADDITIONAL_DETAILS, PHOTOGRAPHY_STYLES, DEVICE, PHOTOGRAPHER, ARTIST, DIGITAL_ARTFORM, PLACE, LIGHTING, CLOTHING, COMPOSITION, POSE, BACKGROUND, pg_title) prompt_generator = PromptGenerator() huggingface_node = HuggingFaceInferenceNode() gr.HTML(pg_title) with gr.Row(): with gr.Column(scale=2): with gr.Accordion("Basic Settings"): pg_seed = gr.Slider(0, 30000, label='Seed', step=1, value=random.randint(0,30000)) pg_custom = gr.Textbox(label="Custom Input Prompt (optional)") pg_subject = gr.Textbox(label="Subject (optional)") # Add the radio button for global option selection pg_global_option = gr.Radio( ["Disabled", "Random", "No Figure Rand"], label="Set all options to:", value="Disabled" ) with gr.Accordion("Artform and Photo Type", open=False): pg_artform = gr.Dropdown(["disabled", "random"] + ARTFORM, label="Artform", value="disabled") pg_photo_type = gr.Dropdown(["disabled", "random"] + PHOTO_TYPE, label="Photo Type", value="disabled") with gr.Accordion("Character Details", open=False): pg_body_types = gr.Dropdown(["disabled", "random"] + BODY_TYPES, label="Body Types", value="disabled") pg_default_tags = gr.Dropdown(["disabled", "random"] + DEFAULT_TAGS, label="Default Tags", value="disabled") pg_roles = gr.Dropdown(["disabled", "random"] + ROLES, label="Roles", value="disabled") pg_hairstyles = gr.Dropdown(["disabled", "random"] + HAIRSTYLES, label="Hairstyles", value="disabled") pg_clothing = gr.Dropdown(["disabled", "random"] + CLOTHING, label="Clothing", value="disabled") with gr.Accordion("Scene Details", open=False): pg_place = gr.Dropdown(["disabled", "random"] + PLACE, label="Place", value="disabled") pg_lighting = gr.Dropdown(["disabled", "random"] + LIGHTING, label="Lighting", value="disabled") pg_composition = gr.Dropdown(["disabled", "random"] + COMPOSITION, label="Composition", value="disabled") pg_pose = gr.Dropdown(["disabled", "random"] + POSE, label="Pose", value="disabled") pg_background = gr.Dropdown(["disabled", "random"] + BACKGROUND, label="Background", value="disabled") with gr.Accordion("Style and Artist", open=False): pg_additional_details = gr.Dropdown(["disabled", "random"] + ADDITIONAL_DETAILS, label="Additional Details", value="disabled") pg_photography_styles = gr.Dropdown(["disabled", "random"] + PHOTOGRAPHY_STYLES, label="Photography Styles", value="disabled") pg_device = gr.Dropdown(["disabled", "random"] + DEVICE, label="Device", value="disabled") pg_photographer = gr.Dropdown(["disabled", "random"] + PHOTOGRAPHER, label="Photographer", value="disabled") pg_artist = gr.Dropdown(["disabled", "random"] + ARTIST, label="Artist", value="disabled") pg_digital_artform = gr.Dropdown(["disabled", "random"] + DIGITAL_ARTFORM, label="Digital Artform", value="disabled") pg_generate_button = gr.Button("Generate Prompt") with gr.Column(scale=2): with gr.Accordion("Image and Caption", open=False): pg_input_image = gr.Image(label="Input Image (optional)") pg_caption_output = gr.Textbox(label="Generated Caption", lines=3) pg_create_caption_button = gr.Button("Create Caption") pg_add_caption_button = gr.Button("Add Caption to Prompt") with gr.Accordion("Prompt Generation", open=True): pg_output = gr.Textbox(label="Generated Prompt / Input Text", lines=4) pg_t5xxl_output = gr.Textbox(label="T5XXL Output", visible=True) pg_clip_l_output = gr.Textbox(label="CLIP L Output", visible=True) pg_clip_g_output = gr.Textbox(label="CLIP G Output", visible=True) with gr.Column(scale=2): with gr.Accordion("Prompt Generation with LLM", open=False): pg_model = gr.Dropdown(["Mixtral", "Mistral", "Llama 3", "Mistral-Nemo"], label="Model", value="Llama 3") pg_happy_talk = gr.Checkbox(label="Happy Talk", value=True) pg_compress = gr.Checkbox(label="Compress", value=True) pg_compression_level = gr.Radio(["soft", "medium", "hard"], label="Compression Level", value="hard") pg_poster = gr.Checkbox(label="Poster", value=False) pg_custom_base_prompt = gr.Textbox(label="Custom Base Prompt", lines=5) pg_generate_text_button = gr.Button("Generate Prompt with LLM") pg_text_output = gr.Textbox(label="Generated Text", lines=10) description_ui() def create_caption(image): if image is not None: return florence_caption(image) return "" pg_create_caption_button.click( create_caption, inputs=[pg_input_image], outputs=[pg_caption_output] ) pg_generate_button.click( prompt_generator.generate_prompt, inputs=[pg_seed, pg_custom, pg_subject, pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_additional_details, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform, pg_place, pg_lighting, pg_clothing, pg_composition, pg_pose, pg_background], outputs=[pg_output, gr.Number(visible=False), pg_t5xxl_output, pg_clip_l_output, pg_clip_g_output] ) pg_add_caption_button.click( prompt_generator.add_caption_to_prompt, inputs=[pg_output, pg_caption_output], outputs=[pg_output] ) pg_generate_text_button.click( huggingface_node.generate, inputs=[pg_model, pg_output, pg_happy_talk, pg_compress, pg_compression_level, pg_poster, pg_custom_base_prompt], outputs=pg_text_output ) def update_all_options(choice): updates = {} if choice == "Disabled": for dropdown in [ pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing, pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform ]: updates[dropdown] = gr.update(value="disabled") elif choice == "Random": for dropdown in [ pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing, pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform ]: updates[dropdown] = gr.update(value="random") else: # No Figure Random for dropdown in [pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing, pg_pose, pg_additional_details]: updates[dropdown] = gr.update(value="disabled") for dropdown in [pg_artform, pg_place, pg_lighting, pg_composition, pg_background, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform]: updates[dropdown] = gr.update(value="random") return updates pg_global_option.change( update_all_options, inputs=[pg_global_option], outputs=[ pg_artform, pg_photo_type, pg_body_types, pg_default_tags, pg_roles, pg_hairstyles, pg_clothing, pg_place, pg_lighting, pg_composition, pg_pose, pg_background, pg_additional_details, pg_photography_styles, pg_device, pg_photographer, pg_artist, pg_digital_artform ] ) app.queue() app.launch()