|
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, |
|
) |
|
|
