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.
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()
to start using private models and skip the queue