app.py

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(
            """<h1><img src="https://huggingface.co/spaces/multimodalart/flux-lora-the-explorer/resolve/main/flux_lora.png" alt="LoRA">FLUX LoRA the Explorer Mod</h1>""",
            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()