app.py

import gradio as gr
from huggingface_hub import login, HfFileSystem, HfApi, ModelCard

from diffusers import DiffusionPipeline, StableDiffusionXLPipeline
import torch
import copy
import os
import spaces
import random

import user_history

is_shared_ui = True if "fffiloni/sd-xl-lora-fusion" in os.environ['SPACE_ID'] else False
hf_token = os.environ.get("HF_TOKEN")
login(token = hf_token)

fs = HfFileSystem(token=hf_token)
api = HfApi()

original_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)

def get_files(file_paths):
    last_files = {}  # Dictionary to store the last file for each path

    for file_path in file_paths:
        # Split the file path into directory and file components
        directory, file_name = file_path.rsplit('/', 1)
    
        # Update the last file for the current path
        last_files[directory] = file_name
    
    # Extract the last files from the dictionary
    result = list(last_files.values())

    return result

def load_sfts(repo_1_id, repo_2_id):

    card_1 = ModelCard.load(repo_1_id)
    
    repo_1_data = card_1.data.to_dict()
    instance_prompt_1 = repo_1_data.get("instance_prompt")
    if instance_prompt_1 is not None:
        print(f"Trigger word 1: {instance_prompt_1}")
    else:
        instance_prompt_1 = "no trigger word needed"
        print(f"Trigger word 1: no trigger word needed")

    card_2 = ModelCard.load(repo_2_id)
    
    repo_2_data = card_2.data.to_dict()
    instance_prompt_2 = repo_2_data.get("instance_prompt")
    if instance_prompt_2 is not None:
        print(f"Trigger word 2: {instance_prompt_2}")
    else:
        instance_prompt_2 = "no trigger word needed"
        print(f"Trigger word 2: no trigger word needed")

      
    # List all ".safetensors" files in repos
    
    sfts_available_files_1 = fs.glob(f"{repo_1_id}/*.safetensors")
    sfts_available_files_1 = get_files(sfts_available_files_1)
    
    if sfts_available_files_1 == []:
        sfts_available_files_1 = ["NO SAFETENSORS FILE"]
    
    print(f"sfts 1: {sfts_available_files_1}")

    
    sfts_available_files_2 = fs.glob(f"{repo_2_id}/*.safetensors")
    sfts_available_files_2 = get_files(sfts_available_files_2)
    
    if sfts_available_files_2 == []:
        sfts_available_files_2 = ["NO SAFETENSORS FILE"]
    
    return gr.update(choices=sfts_available_files_1, value=sfts_available_files_1[0], visible=True), gr.update(choices=sfts_available_files_2, value=sfts_available_files_2[0], visible=True), gr.update(value=instance_prompt_1, visible=True), gr.update(value=instance_prompt_2, visible=True)
    
@spaces.GPU
def infer(lora_1_id, lora_1_sfts, lora_2_id, lora_2_sfts, prompt, negative_prompt, lora_1_scale, lora_2_scale, seed, profile: gr.OAuthProfile | None):

    unet = copy.deepcopy(original_pipe.unet)
    text_encoder = copy.deepcopy(original_pipe.text_encoder)
    text_encoder_2 = copy.deepcopy(original_pipe.text_encoder_2)

    pipe = StableDiffusionXLPipeline(
        vae = original_pipe.vae,
        text_encoder = text_encoder,
        text_encoder_2 = text_encoder_2,
        scheduler = original_pipe.scheduler,
        tokenizer = original_pipe.tokenizer,
        tokenizer_2 = original_pipe.tokenizer_2,
        unet = unet
    )

    pipe.to("cuda")

    if lora_1_sfts == "NO SAFETENSORS FILE": 
        pipe.load_lora_weights(
            lora_1_id,     
            low_cpu_mem_usage = True,
            use_auth_token = True
        )

    else:
        pipe.load_lora_weights(
            lora_1_id,
            weight_name = lora_1_sfts,        
            low_cpu_mem_usage = True,
            use_auth_token = True
        )

    

    pipe.fuse_lora(lora_1_scale)

    if lora_2_sfts == "NO SAFETENSORS FILE": 
        pipe.load_lora_weights(
            lora_2_id,     
            low_cpu_mem_usage = True,
            use_auth_token = True
        )

    else:
        pipe.load_lora_weights(
            lora_2_id,
            weight_name = lora_2_sfts,        
            low_cpu_mem_usage = True,
            use_auth_token = True
        )


    pipe.fuse_lora(lora_2_scale)

    if negative_prompt == "" :
        negative_prompt = None
    
    if seed < 0 :
        seed = random.randint(0, 423538377342)
    
    generator = torch.Generator(device="cuda").manual_seed(seed)

    image = pipe(
        prompt = prompt,
        negative_prompt = negative_prompt,
        num_inference_steps = 25,
        width = 1024,
        height = 1024,
        generator = generator
    ).images[0]

    pipe.unfuse_lora()

    # save generated images (if logged in)
    user_history.save_image(label=prompt, image=image, profile=profile, metadata={
        "prompt": prompt,
        "negative_prompt": negative_prompt,
        "lora_1_repo_id": lora_1_id,
        "lora_2_repo_id": lora_2_id,
        "lora_1_scale": lora_1_scale,
        "lora_2_scale": lora_2_scale,
        "seed": seed,
    })

    return image, seed

css="""
#col-container{
    margin: 0 auto;
    max-width: 750px;
    text-align: left;
}
div#warning-duplicate {
    background-color: #ebf5ff;
    padding: 0 10px 5px;
    margin: 20px 0;
}
div#warning-duplicate > .gr-prose > h2, div#warning-duplicate > .gr-prose > p {
    color: #0f4592!important;
}
div#warning-duplicate strong {
    color: #0f4592;
}
p.actions {
    display: flex;
    align-items: center;
    margin: 20px 0;
}
div#warning-duplicate .actions a {
    display: inline-block;
    margin-right: 10px;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):

        if is_shared_ui:
            top_description = gr.HTML(f'''
                <div class="gr-prose">
                    <h2><svg xmlns="http://www.w3.org/2000/svg" width="18px" height="18px" style="margin-right: 0px;display: inline-block;"fill="none"><path fill="#fff" d="M7 13.2a6.3 6.3 0 0 0 4.4-10.7A6.3 6.3 0 0 0 .6 6.9 6.3 6.3 0 0 0 7 13.2Z"/><path fill="#fff" fill-rule="evenodd" d="M7 0a6.9 6.9 0 0 1 4.8 11.8A6.9 6.9 0 0 1 0 7 6.9 6.9 0 0 1 7 0Zm0 0v.7V0ZM0 7h.6H0Zm7 6.8v-.6.6ZM13.7 7h-.6.6ZM9.1 1.7c-.7-.3-1.4-.4-2.2-.4a5.6 5.6 0 0 0-4 1.6 5.6 5.6 0 0 0-1.6 4 5.6 5.6 0 0 0 1.6 4 5.6 5.6 0 0 0 4 1.7 5.6 5.6 0 0 0 4-1.7 5.6 5.6 0 0 0 1.7-4 5.6 5.6 0 0 0-1.7-4c-.5-.5-1.1-.9-1.8-1.2Z" clip-rule="evenodd"/><path fill="#000" fill-rule="evenodd" d="M7 2.9a.8.8 0 1 1 0 1.5A.8.8 0 0 1 7 3ZM5.8 5.7c0-.4.3-.6.6-.6h.7c.3 0 .6.2.6.6v3.7h.5a.6.6 0 0 1 0 1.3H6a.6.6 0 0 1 0-1.3h.4v-3a.6.6 0 0 1-.6-.7Z" clip-rule="evenodd"/></svg>
                    Note: you might want to use private custom LoRa models</h2>
                    <p class="main-message">
                        To do so, <strong>duplicate the Space</strong> and run it on your own profile using <strong>your own access token</strong> and eventually a GPU (T4-small or A10G-small) for faster inference without waiting in the queue.<br />
                    </p>
                    <p class="actions">
                        <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}?duplicate=true">
                            <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg-dark.svg" alt="Duplicate this Space" />
                        </a>
                        to start using private models and skip the queue
                    </p>
                </div>
            ''', elem_id="warning-duplicate")

        title = gr.HTML(
        '''
        <h1 style="text-align: center;">SD-XL LoRA Fusion</h1>
        <p style="text-align: center;">
        Fuse 2 custom StableDiffusion-XL LoRa models <br />
        If you are running this demo in a duplicated private space, all your private LoRa models tagged ["Diffusers", "stable-diffusion-sd-xl", "lora"] will be automatically listed in LoRa IDs dropdowns
        </p>
        '''
        )
        
        # PART 1 • MODELS
        if not is_shared_ui:
            your_username = api.whoami()["name"]
            my_models = api.list_models(author=your_username, filter=["diffusers", "stable-diffusion-xl", 'lora'])
            model_names = [item.modelId for item in my_models]
            
            #print(model_names)           
            
        with gr.Row():
                
            with gr.Column():

                if not is_shared_ui:
                    lora_1_id = gr.Dropdown(
                        label = "LoRa 1 ID",
                        choices = model_names,
                        allow_custom_value = True
                        #placeholder = "username/model_id"
                    )
                else:
                    lora_1_id = gr.Textbox(
                        label = "LoRa 1 ID",
                        placeholder = "username/model_id"
                    )
                    
                lora_1_sfts = gr.Dropdown(
                    label = "Safetensors file",
                    visible=False
                )

                instance_prompt_1 = gr.Textbox(
                    label = "Trigger Word 1",
                    visible = False,
                    interactive = False
                )
                
            with gr.Column():

                if not is_shared_ui:
                    lora_2_id = gr.Dropdown(
                        label = "LoRa 2 ID",
                        choices = model_names,
                        allow_custom_value = True
                        #placeholder = "username/model_id"
                    )
                else:
                    lora_2_id = gr.Textbox(
                        label = "LoRa 2 ID",
                        placeholder = "username/model_id"
                    )
    
                lora_2_sfts = gr.Dropdown(
                    label = "Safetensors file",
                    visible=False
                )

                instance_prompt_2 = gr.Textbox(
                    label = "Trigger Word 2",
                    visible = False,
                    interactive = False
                )
        
        load_models_btn = gr.Button("1. Load models and .safetensors")

        # PART 2 • INFERENCE
        with gr.Column():
            with gr.Row():
            
                prompt = gr.Textbox(
                    label = "Your prompt",
                    show_label = True,
                    info = "Use your trigger words into a coherent prompt",
                    placeholder = "e.g: a triggerWordOne portrait in triggerWord2 style"
                )
            # Advanced Settings
            with gr.Accordion("Advanced Settings", open=False):
                
                with gr.Row():
                    
                    lora_1_scale = gr.Slider(
                        label = "LoRa 1 scale",
                        minimum = 0,
                        maximum = 1,
                        step = 0.1,
                        value = 0.7
                    )
                    
                    lora_2_scale = gr.Slider(
                        label = "LoRa 2 scale",
                        minimum = 0,
                        maximum = 1,
                        step = 0.1,
                        value = 0.7
                    )
                
                negative_prompt = gr.Textbox(
                    label = "Negative prompt"
                )
    
                seed = gr.Slider(
                    label = "Seed",
                    info = "-1 denotes a random seed",
                    minimum = -1,
                    maximum = 423538377342,
                    value = -1
                )
    
                last_used_seed = gr.Number(
                    label = "Last used seed",
                    info = "the seed used in the last generation",
                )
                
            run_btn = gr.Button("2. Run", elem_id="run_button")
        
            output_image = gr.Image(
                label = "Output"
            )

        with gr.Accordion("Past generations", open=False):
            user_history.render()

        
    
    # ACTIONS
    load_models_btn.click(
        fn = load_sfts, 
        inputs = [
            lora_1_id,
            lora_2_id
        ],
        outputs = [
            lora_1_sfts,
            lora_2_sfts,
            instance_prompt_1,
            instance_prompt_2
        ],
        queue=False
    )
    run_btn.click(
        fn = infer,
        inputs = [
            lora_1_id,
            lora_1_sfts,
            lora_2_id,
            lora_2_sfts,
            prompt,
            negative_prompt,
            lora_1_scale,
            lora_2_scale,
            seed
        ],
        outputs = [
            output_image, 
            last_used_seed
        ]
    )

demo.queue(concurrency_count=2).launch()