Spaces:
Runtime error
Runtime error
| 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 | |
| from backend.models.lcmdiffusion_setting import LCMDiffusionSetting | |
| _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 | |
| ) | |
| # Load a new LoRA | |
| settings = app_settings.settings.lcm_diffusion_setting | |
| settings.lora.fuse = False | |
| settings.lora.enabled = False | |
| 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, | |
| ) | |
| # Update Gradio LoRA UI | |
| 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, | |
| ) | |