Spaces:
Sleeping
Sleeping
| 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) | |
| 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() | |