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'''

Note: you might want to use private custom LoRa models

To do so, duplicate the Space and run it on your own profile using your own access token and eventually a GPU (T4-small or A10G-small) for faster inference without waiting in the queue.

Duplicate this Space to start using private models and skip the queue

''', elem_id="warning-duplicate") title = gr.HTML( '''

SD-XL LoRA Fusion

Fuse 2 custom StableDiffusion-XL LoRa models
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

''' ) # 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()