| | import gradio as gr
|
| | from os import path
|
| | from backend.lora import (
|
| | get_lora_models,
|
| | get_active_lora_weights,
|
| | update_lora_weights,
|
| | load_lora_weight,
|
| | )
|
| | from state import get_settings, get_context
|
| | from frontend.utils import get_valid_lora_model
|
| | from models.interface_types import InterfaceType
|
| |
|
| |
|
| | _MAX_LORA_WEIGHTS = 5
|
| |
|
| | _custom_lora_sliders = []
|
| | _custom_lora_names = []
|
| | _custom_lora_columns = []
|
| |
|
| | app_settings = get_settings()
|
| |
|
| |
|
| | def on_click_update_weight(*lora_weights):
|
| | update_weights = []
|
| | active_weights = get_active_lora_weights()
|
| | if not len(active_weights):
|
| | gr.Warning("No active LoRAs, first you need to load LoRA model")
|
| | return
|
| | for idx, lora in enumerate(active_weights):
|
| | update_weights.append(
|
| | (
|
| | lora[0],
|
| | lora_weights[idx],
|
| | )
|
| | )
|
| | if len(update_weights) > 0:
|
| | update_lora_weights(
|
| | get_context(InterfaceType.WEBUI).lcm_text_to_image.pipeline,
|
| | app_settings.settings.lcm_diffusion_setting,
|
| | update_weights,
|
| | )
|
| |
|
| |
|
| | def on_click_load_lora(lora_name, lora_weight):
|
| | if app_settings.settings.lcm_diffusion_setting.use_openvino:
|
| | gr.Warning("Currently LoRA is not supported in OpenVINO.")
|
| | return
|
| | lora_models_map = get_lora_models(
|
| | app_settings.settings.lcm_diffusion_setting.lora.models_dir
|
| | )
|
| |
|
| |
|
| | settings = app_settings.settings.lcm_diffusion_setting
|
| | settings.lora.fuse = False
|
| | settings.lora.enabled = False
|
| | print(f"Selected Lora Model :{lora_name}")
|
| | print(f"Lora weight :{lora_weight}")
|
| | settings.lora.path = lora_models_map[lora_name]
|
| | settings.lora.weight = lora_weight
|
| | if not path.exists(settings.lora.path):
|
| | gr.Warning("Invalid LoRA model path!")
|
| | return
|
| | pipeline = get_context(InterfaceType.WEBUI).lcm_text_to_image.pipeline
|
| | if not pipeline:
|
| | gr.Warning("Pipeline not initialized. Please generate an image first.")
|
| | return
|
| | settings.lora.enabled = True
|
| | load_lora_weight(
|
| | get_context(InterfaceType.WEBUI).lcm_text_to_image.pipeline,
|
| | settings,
|
| | )
|
| |
|
| |
|
| | global _MAX_LORA_WEIGHTS
|
| | values = []
|
| | labels = []
|
| | rows = []
|
| | active_weights = get_active_lora_weights()
|
| | for idx, lora in enumerate(active_weights):
|
| | labels.append(f"{lora[0]}: ")
|
| | values.append(lora[1])
|
| | rows.append(gr.Row.update(visible=True))
|
| | for i in range(len(active_weights), _MAX_LORA_WEIGHTS):
|
| | labels.append(f"Update weight")
|
| | values.append(0.0)
|
| | rows.append(gr.Row.update(visible=False))
|
| | return labels + values + rows
|
| |
|
| |
|
| | def get_lora_models_ui() -> None:
|
| | with gr.Blocks() as ui:
|
| | gr.HTML(
|
| | "Download and place your LoRA model weights in <b>lora_models</b> folders and restart App"
|
| | )
|
| | with gr.Row():
|
| | with gr.Column():
|
| | with gr.Row():
|
| | lora_models_map = get_lora_models(
|
| | app_settings.settings.lcm_diffusion_setting.lora.models_dir
|
| | )
|
| | valid_model = get_valid_lora_model(
|
| | list(lora_models_map.values()),
|
| | app_settings.settings.lcm_diffusion_setting.lora.path,
|
| | app_settings.settings.lcm_diffusion_setting.lora.models_dir,
|
| | )
|
| | if valid_model != "":
|
| | valid_model_path = lora_models_map[valid_model]
|
| | app_settings.settings.lcm_diffusion_setting.lora.path = (
|
| | valid_model_path
|
| | )
|
| | else:
|
| | app_settings.settings.lcm_diffusion_setting.lora.path = ""
|
| |
|
| | lora_model = gr.Dropdown(
|
| | lora_models_map.keys(),
|
| | label="LoRA model",
|
| | info="LoRA model weight to load (You can use Lora models from Civitai or Hugging Face .safetensors format)",
|
| | value=valid_model,
|
| | interactive=True,
|
| | )
|
| |
|
| | lora_weight = gr.Slider(
|
| | 0.0,
|
| | 1.0,
|
| | value=app_settings.settings.lcm_diffusion_setting.lora.weight,
|
| | step=0.05,
|
| | label="Initial Lora weight",
|
| | interactive=True,
|
| | )
|
| | load_lora_btn = gr.Button(
|
| | "Load selected LoRA",
|
| | elem_id="load_lora_button",
|
| | scale=0,
|
| | )
|
| |
|
| | with gr.Row():
|
| | gr.Markdown(
|
| | "## Loaded LoRA models",
|
| | show_label=False,
|
| | )
|
| | update_lora_weights_btn = gr.Button(
|
| | "Update LoRA weights",
|
| | elem_id="load_lora_button",
|
| | scale=0,
|
| | )
|
| |
|
| | global _MAX_LORA_WEIGHTS
|
| | global _custom_lora_sliders
|
| | global _custom_lora_names
|
| | global _custom_lora_columns
|
| | for i in range(0, _MAX_LORA_WEIGHTS):
|
| | new_row = gr.Column(visible=False)
|
| | _custom_lora_columns.append(new_row)
|
| | with new_row:
|
| | lora_name = gr.Markdown(
|
| | "Lora Name",
|
| | show_label=True,
|
| | )
|
| | lora_slider = gr.Slider(
|
| | 0.0,
|
| | 1.0,
|
| | step=0.05,
|
| | label="LoRA weight",
|
| | interactive=True,
|
| | visible=True,
|
| | )
|
| |
|
| | _custom_lora_names.append(lora_name)
|
| | _custom_lora_sliders.append(lora_slider)
|
| |
|
| | load_lora_btn.click(
|
| | fn=on_click_load_lora,
|
| | inputs=[lora_model, lora_weight],
|
| | outputs=[
|
| | *_custom_lora_names,
|
| | *_custom_lora_sliders,
|
| | *_custom_lora_columns,
|
| | ],
|
| | )
|
| |
|
| | update_lora_weights_btn.click(
|
| | fn=on_click_update_weight,
|
| | inputs=[*_custom_lora_sliders],
|
| | outputs=None,
|
| | )
|
| |
|
