Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import yaml | |
| import os | |
| import shutil | |
| from functools import lru_cache | |
| from core.settings import * | |
| from utils.app_utils import * | |
| from core.generation_logic import * | |
| from comfy_integration.nodes import SAMPLER_CHOICES, SCHEDULER_CHOICES | |
| from core.pipelines.controlnet_preprocessor import CPU_ONLY_PREPROCESSORS | |
| from utils.app_utils import PREPROCESSOR_MODEL_MAP, PREPROCESSOR_PARAMETER_MAP, save_uploaded_file_with_hash | |
| from ui.shared.ui_components import RESOLUTION_MAP, MAX_CONTROLNETS, MAX_IPADAPTERS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_LORAS | |
| def load_controlnet_config(): | |
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| _CN_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'controlnet_models.yaml') | |
| try: | |
| print("--- Loading controlnet_models.yaml ---") | |
| with open(_CN_MODEL_LIST_PATH, 'r', encoding='utf-8') as f: | |
| config = yaml.safe_load(f) | |
| print("--- ✅ controlnet_models.yaml loaded successfully ---") | |
| return config.get("ControlNet", {}).get("SDXL", []) | |
| except Exception as e: | |
| print(f"Error loading controlnet_models.yaml: {e}") | |
| return [] | |
| def load_ipadapter_config(): | |
| _PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| _IPA_MODEL_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'ipadapter.yaml') | |
| try: | |
| print("--- Loading ipadapter.yaml ---") | |
| with open(_IPA_MODEL_LIST_PATH, 'r', encoding='utf-8') as f: | |
| config = yaml.safe_load(f) | |
| print("--- ✅ ipadapter.yaml loaded successfully ---") | |
| return config | |
| except Exception as e: | |
| print(f"Error loading ipadapter.yaml: {e}") | |
| return {} | |
| def apply_data_to_ui(data, prefix, ui_components): | |
| final_sampler = data.get('sampler') if data.get('sampler') in SAMPLER_CHOICES else SAMPLER_CHOICES[0] | |
| default_scheduler = 'normal' if 'normal' in SCHEDULER_CHOICES else SCHEDULER_CHOICES[0] | |
| final_scheduler = data.get('scheduler') if data.get('scheduler') in SCHEDULER_CHOICES else default_scheduler | |
| updates = {} | |
| base_model_name = data.get('base_model') | |
| model_map = MODEL_MAP_CHECKPOINT | |
| if f'base_model_{prefix}' in ui_components: | |
| model_dropdown_component = ui_components[f'base_model_{prefix}'] | |
| if base_model_name and base_model_name in model_map: | |
| updates[model_dropdown_component] = base_model_name | |
| else: | |
| updates[model_dropdown_component] = gr.update() | |
| common_params = { | |
| f'prompt_{prefix}': data.get('prompt', ''), | |
| f'neg_prompt_{prefix}': data.get('negative_prompt', ''), | |
| f'seed_{prefix}': data.get('seed', -1), | |
| f'cfg_{prefix}': data.get('cfg_scale', 7.5), | |
| f'steps_{prefix}': data.get('steps', 28), | |
| f'sampler_{prefix}': final_sampler, | |
| f'scheduler_{prefix}': final_scheduler, | |
| } | |
| for comp_name, value in common_params.items(): | |
| if comp_name in ui_components: | |
| updates[ui_components[comp_name]] = value | |
| if prefix == 'txt2img': | |
| if f'width_{prefix}' in ui_components: | |
| updates[ui_components[f'width_{prefix}']] = data.get('width', 1024) | |
| if f'height_{prefix}' in ui_components: | |
| updates[ui_components[f'height_{prefix}']] = data.get('height', 1024) | |
| tab_indices = {"txt2img": 0, "img2img": 1, "inpaint": 2, "outpaint": 3, "hires_fix": 4} | |
| tab_index = tab_indices.get(prefix, 0) | |
| updates[ui_components['tabs']] = gr.Tabs(selected=0) | |
| updates[ui_components['image_gen_tabs']] = gr.Tabs(selected=tab_index) | |
| return updates | |
| def send_info_to_tab(image, prefix, ui_components): | |
| if not image or not image.info.get('parameters', ''): | |
| all_comps = [comp for comp_or_list in ui_components.values() for comp in (comp_or_list if isinstance(comp_or_list, list) else [comp_or_list])] | |
| return {comp: gr.update() for comp in all_comps} | |
| data = parse_parameters(image.info['parameters']) | |
| image_input_map = { | |
| "img2img": 'input_image_img2img', | |
| "inpaint": 'input_image_dict_inpaint', | |
| "outpaint": 'input_image_outpaint', | |
| "hires_fix": 'input_image_hires_fix' | |
| } | |
| updates = apply_data_to_ui(data, prefix, ui_components) | |
| if prefix in image_input_map and image_input_map[prefix] in ui_components: | |
| component_key = image_input_map[prefix] | |
| updates[ui_components[component_key]] = gr.update(value=image) | |
| return updates | |
| def send_info_by_hash(image, ui_components): | |
| if not image or not image.info.get('parameters', ''): | |
| all_comps = [comp for comp_or_list in ui_components.values() for comp in (comp_or_list if isinstance(comp_or_list, list) else [comp_or_list])] | |
| return {comp: gr.update() for comp in all_comps} | |
| data = parse_parameters(image.info['parameters']) | |
| return apply_data_to_ui(data, "txt2img", ui_components) | |
| def attach_event_handlers(ui_components, demo): | |
| def update_cn_input_visibility(choice): | |
| return { | |
| ui_components["cn_image_input"]: gr.update(visible=choice == "Image"), | |
| ui_components["cn_video_input"]: gr.update(visible=choice == "Video") | |
| } | |
| ui_components["cn_input_type"].change( | |
| fn=update_cn_input_visibility, | |
| inputs=[ui_components["cn_input_type"]], | |
| outputs=[ui_components["cn_image_input"], ui_components["cn_video_input"]] | |
| ) | |
| def update_preprocessor_models_dropdown(preprocessor_name): | |
| models = PREPROCESSOR_MODEL_MAP.get(preprocessor_name) | |
| if models: | |
| model_filenames = [m[1] for m in models] | |
| return gr.update(choices=model_filenames, value=model_filenames[0], visible=True) | |
| else: | |
| return gr.update(choices=[], value=None, visible=False) | |
| def update_preprocessor_settings_ui(preprocessor_name): | |
| from ui.layout import MAX_DYNAMIC_CONTROLS | |
| params = PREPROCESSOR_PARAMETER_MAP.get(preprocessor_name, []) | |
| slider_updates, dropdown_updates, checkbox_updates = [], [], [] | |
| s_idx, d_idx, c_idx = 0, 0, 0 | |
| for param in params: | |
| if s_idx + d_idx + c_idx >= MAX_DYNAMIC_CONTROLS: break | |
| name = param["name"] | |
| ptype = param["type"] | |
| config = param["config"] | |
| label = name.replace('_', ' ').title() | |
| if ptype == "INT" or ptype == "FLOAT": | |
| if s_idx < MAX_DYNAMIC_CONTROLS: | |
| slider_updates.append(gr.update( | |
| label=label, | |
| minimum=config.get('min', 0), | |
| maximum=config.get('max', 255), | |
| step=config.get('step', 0.1 if ptype == "FLOAT" else 1), | |
| value=config.get('default', 0), | |
| visible=True | |
| )) | |
| s_idx += 1 | |
| elif isinstance(ptype, list): | |
| if d_idx < MAX_DYNAMIC_CONTROLS: | |
| dropdown_updates.append(gr.update( | |
| label=label, | |
| choices=ptype, | |
| value=config.get('default', ptype[0] if ptype else None), | |
| visible=True | |
| )) | |
| d_idx += 1 | |
| elif ptype == "BOOLEAN": | |
| if c_idx < MAX_DYNAMIC_CONTROLS: | |
| checkbox_updates.append(gr.update( | |
| label=label, | |
| value=config.get('default', False), | |
| visible=True | |
| )) | |
| c_idx += 1 | |
| for _ in range(s_idx, MAX_DYNAMIC_CONTROLS): slider_updates.append(gr.update(visible=False)) | |
| for _ in range(d_idx, MAX_DYNAMIC_CONTROLS): dropdown_updates.append(gr.update(visible=False)) | |
| for _ in range(c_idx, MAX_DYNAMIC_CONTROLS): checkbox_updates.append(gr.update(visible=False)) | |
| return slider_updates + dropdown_updates + checkbox_updates | |
| def update_run_button_for_cpu(preprocessor_name): | |
| if preprocessor_name in CPU_ONLY_PREPROCESSORS: | |
| return gr.update(value="Run Preprocessor CPU Only", variant="primary"), gr.update(visible=False) | |
| else: | |
| return gr.update(value="Run Preprocessor", variant="primary"), gr.update(visible=True) | |
| ui_components["preprocessor_cn"].change( | |
| fn=update_preprocessor_models_dropdown, | |
| inputs=[ui_components["preprocessor_cn"]], | |
| outputs=[ui_components["preprocessor_model_cn"]] | |
| ).then( | |
| fn=update_preprocessor_settings_ui, | |
| inputs=[ui_components["preprocessor_cn"]], | |
| outputs=ui_components["cn_sliders"] + ui_components["cn_dropdowns"] + ui_components["cn_checkboxes"] | |
| ).then( | |
| fn=update_run_button_for_cpu, | |
| inputs=[ui_components["preprocessor_cn"]], | |
| outputs=[ui_components["run_cn"], ui_components["zero_gpu_cn"]] | |
| ) | |
| all_dynamic_inputs = ( | |
| ui_components["cn_sliders"] + | |
| ui_components["cn_dropdowns"] + | |
| ui_components["cn_checkboxes"] | |
| ) | |
| ui_components["run_cn"].click( | |
| fn=run_cn_preprocessor_entry, | |
| inputs=[ | |
| ui_components["cn_input_type"], | |
| ui_components["cn_image_input"], | |
| ui_components["cn_video_input"], | |
| ui_components["preprocessor_cn"], | |
| ui_components["preprocessor_model_cn"], | |
| ui_components["zero_gpu_cn"], | |
| ] + all_dynamic_inputs, | |
| outputs=[ui_components["output_gallery_cn"]] | |
| ) | |
| def create_lora_event_handlers(prefix): | |
| lora_rows = ui_components[f'lora_rows_{prefix}'] | |
| lora_ids = ui_components[f'lora_ids_{prefix}'] | |
| lora_scales = ui_components[f'lora_scales_{prefix}'] | |
| lora_uploads = ui_components[f'lora_uploads_{prefix}'] | |
| count_state = ui_components[f'lora_count_state_{prefix}'] | |
| add_button = ui_components[f'add_lora_button_{prefix}'] | |
| del_button = ui_components[f'delete_lora_button_{prefix}'] | |
| def add_lora_row(c): | |
| updates = {} | |
| if c < MAX_LORAS: | |
| c += 1 | |
| updates[lora_rows[c - 1]] = gr.update(visible=True) | |
| updates[count_state] = c | |
| updates[add_button] = gr.update(visible=c < MAX_LORAS) | |
| updates[del_button] = gr.update(visible=c > 1) | |
| return updates | |
| def del_lora_row(c): | |
| updates = {} | |
| if c > 1: | |
| updates[lora_rows[c - 1]] = gr.update(visible=False) | |
| updates[lora_ids[c - 1]] = "" | |
| updates[lora_scales[c - 1]] = 0.0 | |
| updates[lora_uploads[c - 1]] = None | |
| c -= 1 | |
| updates[count_state] = c | |
| updates[add_button] = gr.update(visible=True) | |
| updates[del_button] = gr.update(visible=c > 1) | |
| return updates | |
| add_outputs = [count_state, add_button, del_button] + lora_rows | |
| del_outputs = [count_state, add_button, del_button] + lora_rows + lora_ids + lora_scales + lora_uploads | |
| add_button.click(add_lora_row, [count_state], add_outputs, show_progress=False) | |
| del_button.click(del_lora_row, [count_state], del_outputs, show_progress=False) | |
| def create_controlnet_event_handlers(prefix): | |
| cn_rows = ui_components[f'controlnet_rows_{prefix}'] | |
| cn_types = ui_components[f'controlnet_types_{prefix}'] | |
| cn_series = ui_components[f'controlnet_series_{prefix}'] | |
| cn_filepaths = ui_components[f'controlnet_filepaths_{prefix}'] | |
| cn_images = ui_components[f'controlnet_images_{prefix}'] | |
| cn_strengths = ui_components[f'controlnet_strengths_{prefix}'] | |
| count_state = ui_components[f'controlnet_count_state_{prefix}'] | |
| add_button = ui_components[f'add_controlnet_button_{prefix}'] | |
| del_button = ui_components[f'delete_controlnet_button_{prefix}'] | |
| accordion = ui_components[f'controlnet_accordion_{prefix}'] | |
| def add_cn_row(c): | |
| c += 1 | |
| updates = { | |
| count_state: c, | |
| cn_rows[c-1]: gr.update(visible=True), | |
| add_button: gr.update(visible=c < MAX_CONTROLNETS), | |
| del_button: gr.update(visible=True) | |
| } | |
| return updates | |
| def del_cn_row(c): | |
| c -= 1 | |
| updates = { | |
| count_state: c, | |
| cn_rows[c]: gr.update(visible=False), | |
| cn_images[c]: None, | |
| cn_strengths[c]: 1.0, | |
| add_button: gr.update(visible=True), | |
| del_button: gr.update(visible=c > 0) | |
| } | |
| return updates | |
| add_outputs = [count_state, add_button, del_button] + cn_rows | |
| del_outputs = [count_state, add_button, del_button] + cn_rows + cn_images + cn_strengths | |
| add_button.click(fn=add_cn_row, inputs=[count_state], outputs=add_outputs, show_progress=False) | |
| del_button.click(fn=del_cn_row, inputs=[count_state], outputs=del_outputs, show_progress=False) | |
| def on_cn_type_change(selected_type): | |
| cn_config = load_controlnet_config() | |
| series_choices = [] | |
| if selected_type: | |
| series_choices = sorted(list(set( | |
| model.get("Series", "Default") for model in cn_config | |
| if selected_type in model.get("Type", []) | |
| ))) | |
| default_series = series_choices[0] if series_choices else None | |
| filepath = "None" | |
| if default_series: | |
| for model in cn_config: | |
| if model.get("Series") == default_series and selected_type in model.get("Type", []): | |
| filepath = model.get("Filepath") | |
| break | |
| return gr.update(choices=series_choices, value=default_series), filepath | |
| def on_cn_series_change(selected_series, selected_type): | |
| cn_config = load_controlnet_config() | |
| filepath = "None" | |
| if selected_series and selected_type: | |
| for model in cn_config: | |
| if model.get("Series") == selected_series and selected_type in model.get("Type", []): | |
| filepath = model.get("Filepath") | |
| break | |
| return filepath | |
| for i in range(MAX_CONTROLNETS): | |
| cn_types[i].change( | |
| fn=on_cn_type_change, | |
| inputs=[cn_types[i]], | |
| outputs=[cn_series[i], cn_filepaths[i]], | |
| show_progress=False | |
| ) | |
| cn_series[i].change( | |
| fn=on_cn_series_change, | |
| inputs=[cn_series[i], cn_types[i]], | |
| outputs=[cn_filepaths[i]], | |
| show_progress=False | |
| ) | |
| def on_accordion_expand(*images): | |
| return [gr.update() for _ in images] | |
| accordion.expand( | |
| fn=on_accordion_expand, | |
| inputs=cn_images, | |
| outputs=cn_images, | |
| show_progress=False | |
| ) | |
| def create_ipadapter_event_handlers(prefix): | |
| ipa_rows = ui_components[f'ipadapter_rows_{prefix}'] | |
| ipa_lora_strengths = ui_components[f'ipadapter_lora_strengths_{prefix}'] | |
| ipa_final_preset = ui_components[f'ipadapter_final_preset_{prefix}'] | |
| ipa_final_lora_strength = ui_components[f'ipadapter_final_lora_strength_{prefix}'] | |
| count_state = ui_components[f'ipadapter_count_state_{prefix}'] | |
| add_button = ui_components[f'add_ipadapter_button_{prefix}'] | |
| del_button = ui_components[f'delete_ipadapter_button_{prefix}'] | |
| accordion = ui_components[f'ipadapter_accordion_{prefix}'] | |
| def add_ipa_row(c): | |
| c += 1 | |
| return { | |
| count_state: c, | |
| ipa_rows[c - 1]: gr.update(visible=True), | |
| add_button: gr.update(visible=c < MAX_IPADAPTERS), | |
| del_button: gr.update(visible=True), | |
| } | |
| def del_ipa_row(c): | |
| c -= 1 | |
| return { | |
| count_state: c, | |
| ipa_rows[c]: gr.update(visible=False), | |
| add_button: gr.update(visible=True), | |
| del_button: gr.update(visible=c > 0), | |
| } | |
| add_outputs = [count_state, add_button, del_button] + ipa_rows | |
| del_outputs = [count_state, add_button, del_button] + ipa_rows | |
| add_button.click(fn=add_ipa_row, inputs=[count_state], outputs=add_outputs, show_progress=False) | |
| del_button.click(fn=del_ipa_row, inputs=[count_state], outputs=del_outputs, show_progress=False) | |
| def on_preset_change(preset_value): | |
| config = load_ipadapter_config() | |
| faceid_presets = [] | |
| if isinstance(config, list): | |
| faceid_presets = [ | |
| p.get('preset_name', '') for p in config | |
| if 'FACE' in p.get('preset_name', '') or 'FACEID' in p.get('preset_name', '') | |
| ] | |
| is_visible = preset_value in faceid_presets | |
| updates = [gr.update(visible=is_visible)] * (MAX_IPADAPTERS + 1) | |
| return updates | |
| all_lora_strength_sliders = [ipa_final_lora_strength] + ipa_lora_strengths | |
| ipa_final_preset.change(fn=on_preset_change, inputs=[ipa_final_preset], outputs=all_lora_strength_sliders, show_progress=False) | |
| accordion.expand(fn=lambda *imgs: [gr.update() for _ in imgs], inputs=ui_components[f'ipadapter_images_{prefix}'], outputs=ui_components[f'ipadapter_images_{prefix}'], show_progress=False) | |
| def create_embedding_event_handlers(prefix): | |
| rows = ui_components[f'embedding_rows_{prefix}'] | |
| ids = ui_components[f'embeddings_ids_{prefix}'] | |
| files = ui_components[f'embeddings_files_{prefix}'] | |
| count_state = ui_components[f'embedding_count_state_{prefix}'] | |
| add_button = ui_components[f'add_embedding_button_{prefix}'] | |
| del_button = ui_components[f'delete_embedding_button_{prefix}'] | |
| def add_row(c): | |
| c += 1 | |
| return { | |
| count_state: c, | |
| rows[c - 1]: gr.update(visible=True), | |
| add_button: gr.update(visible=c < MAX_EMBEDDINGS), | |
| del_button: gr.update(visible=True) | |
| } | |
| def del_row(c): | |
| c -= 1 | |
| return { | |
| count_state: c, | |
| rows[c]: gr.update(visible=False), | |
| ids[c]: "", | |
| files[c]: None, | |
| add_button: gr.update(visible=True), | |
| del_button: gr.update(visible=c > 0) | |
| } | |
| add_outputs = [count_state, add_button, del_button] + rows | |
| del_outputs = [count_state, add_button, del_button] + rows + ids + files | |
| add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False) | |
| del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False) | |
| def create_conditioning_event_handlers(prefix): | |
| rows = ui_components[f'conditioning_rows_{prefix}'] | |
| prompts = ui_components[f'conditioning_prompts_{prefix}'] | |
| count_state = ui_components[f'conditioning_count_state_{prefix}'] | |
| add_button = ui_components[f'add_conditioning_button_{prefix}'] | |
| del_button = ui_components[f'delete_conditioning_button_{prefix}'] | |
| def add_row(c): | |
| c += 1 | |
| return { | |
| count_state: c, | |
| rows[c - 1]: gr.update(visible=True), | |
| add_button: gr.update(visible=c < MAX_CONDITIONINGS), | |
| del_button: gr.update(visible=True), | |
| } | |
| def del_row(c): | |
| c -= 1 | |
| return { | |
| count_state: c, | |
| rows[c]: gr.update(visible=False), | |
| prompts[c]: "", | |
| add_button: gr.update(visible=True), | |
| del_button: gr.update(visible=c > 0), | |
| } | |
| add_outputs = [count_state, add_button, del_button] + rows | |
| del_outputs = [count_state, add_button, del_button] + rows + prompts | |
| add_button.click(fn=add_row, inputs=[count_state], outputs=add_outputs, show_progress=False) | |
| del_button.click(fn=del_row, inputs=[count_state], outputs=del_outputs, show_progress=False) | |
| def on_vae_upload(file_obj): | |
| if not file_obj: | |
| return gr.update(), gr.update(), None | |
| hashed_filename = save_uploaded_file_with_hash(file_obj, VAE_DIR) | |
| return hashed_filename, "File", file_obj | |
| def on_lora_upload(file_obj): | |
| if not file_obj: | |
| return gr.update(), gr.update() | |
| hashed_filename = save_uploaded_file_with_hash(file_obj, LORA_DIR) | |
| return hashed_filename, "File" | |
| def on_embedding_upload(file_obj): | |
| if not file_obj: | |
| return gr.update(), gr.update(), None | |
| hashed_filename = save_uploaded_file_with_hash(file_obj, EMBEDDING_DIR) | |
| return hashed_filename, "File", file_obj | |
| def create_run_event(prefix: str, task_type: str): | |
| run_inputs_map = { | |
| 'model_display_name': ui_components[f'base_model_{prefix}'], | |
| 'positive_prompt': ui_components[f'prompt_{prefix}'], | |
| 'negative_prompt': ui_components[f'neg_prompt_{prefix}'], | |
| 'seed': ui_components[f'seed_{prefix}'], | |
| 'batch_size': ui_components[f'batch_size_{prefix}'], | |
| 'guidance_scale': ui_components[f'cfg_{prefix}'], | |
| 'num_inference_steps': ui_components[f'steps_{prefix}'], | |
| 'sampler': ui_components[f'sampler_{prefix}'], | |
| 'scheduler': ui_components[f'scheduler_{prefix}'], | |
| 'zero_gpu_duration': ui_components[f'zero_gpu_{prefix}'], | |
| 'civitai_api_key': ui_components.get(f'civitai_api_key_{prefix}'), | |
| 'clip_skip': ui_components[f'clip_skip_{prefix}'], | |
| 'task_type': gr.State(task_type) | |
| } | |
| if task_type not in ['img2img', 'inpaint']: | |
| run_inputs_map.update({'width': ui_components[f'width_{prefix}'], 'height': ui_components[f'height_{prefix}']}) | |
| task_specific_map = { | |
| 'img2img': {'img2img_image': f'input_image_{prefix}', 'img2img_denoise': f'denoise_{prefix}'}, | |
| 'inpaint': {'inpaint_image_dict': f'input_image_dict_{prefix}'}, | |
| 'outpaint': {'outpaint_image': f'input_image_{prefix}', 'outpaint_left': f'outpaint_left_{prefix}', 'outpaint_top': f'outpaint_top_{prefix}', 'outpaint_right': f'outpaint_right_{prefix}', 'outpaint_bottom': f'outpaint_bottom_{prefix}'}, | |
| 'hires_fix': {'hires_image': f'input_image_{prefix}', 'hires_upscaler': f'hires_upscaler_{prefix}', 'hires_scale_by': f'hires_scale_by_{prefix}', 'hires_denoise': f'denoise_{prefix}'} | |
| } | |
| if task_type in task_specific_map: | |
| for key, comp_name in task_specific_map[task_type].items(): | |
| run_inputs_map[key] = ui_components[comp_name] | |
| lora_data_components = ui_components.get(f'all_lora_components_flat_{prefix}', []) | |
| controlnet_data_components = ui_components.get(f'all_controlnet_components_flat_{prefix}', []) | |
| ipadapter_data_components = ui_components.get(f'all_ipadapter_components_flat_{prefix}', []) | |
| embedding_data_components = ui_components.get(f'all_embedding_components_flat_{prefix}', []) | |
| conditioning_data_components = ui_components.get(f'all_conditioning_components_flat_{prefix}', []) | |
| run_inputs_map['vae_source'] = ui_components.get(f'vae_source_{prefix}') | |
| run_inputs_map['vae_id'] = ui_components.get(f'vae_id_{prefix}') | |
| run_inputs_map['vae_file'] = ui_components.get(f'vae_file_{prefix}') | |
| input_keys = list(run_inputs_map.keys()) | |
| input_list_flat = [v for v in run_inputs_map.values() if v is not None] | |
| input_list_flat += lora_data_components + controlnet_data_components + ipadapter_data_components + embedding_data_components + conditioning_data_components | |
| def create_ui_inputs_dict(*args): | |
| valid_keys = [k for k in input_keys if run_inputs_map[k] is not None] | |
| ui_dict = dict(zip(valid_keys, args[:len(valid_keys)])) | |
| arg_idx = len(valid_keys) | |
| ui_dict['lora_data'] = list(args[arg_idx : arg_idx + len(lora_data_components)]) | |
| arg_idx += len(lora_data_components) | |
| ui_dict['controlnet_data'] = list(args[arg_idx : arg_idx + len(controlnet_data_components)]) | |
| arg_idx += len(controlnet_data_components) | |
| ui_dict['ipadapter_data'] = list(args[arg_idx : arg_idx + len(ipadapter_data_components)]) | |
| arg_idx += len(ipadapter_data_components) | |
| ui_dict['embedding_data'] = list(args[arg_idx : arg_idx + len(embedding_data_components)]) | |
| arg_idx += len(embedding_data_components) | |
| ui_dict['conditioning_data'] = list(args[arg_idx : arg_idx + len(conditioning_data_components)]) | |
| return ui_dict | |
| ui_components[f'run_{prefix}'].click( | |
| fn=lambda *args, progress=gr.Progress(track_tqdm=True): generate_image_wrapper(create_ui_inputs_dict(*args), progress), | |
| inputs=input_list_flat, | |
| outputs=[ui_components[f'result_{prefix}']] | |
| ) | |
| for prefix, task_type in [ | |
| ("txt2img", "txt2img"), ("img2img", "img2img"), ("inpaint", "inpaint"), | |
| ("outpaint", "outpaint"), ("hires_fix", "hires_fix"), | |
| ]: | |
| if f'add_lora_button_{prefix}' in ui_components: | |
| create_lora_event_handlers(prefix) | |
| lora_uploads = ui_components[f'lora_uploads_{prefix}'] | |
| lora_ids = ui_components[f'lora_ids_{prefix}'] | |
| lora_sources = ui_components[f'lora_sources_{prefix}'] | |
| for i in range(MAX_LORAS): | |
| lora_uploads[i].upload( | |
| fn=on_lora_upload, | |
| inputs=[lora_uploads[i]], | |
| outputs=[lora_ids[i], lora_sources[i]], | |
| show_progress=False | |
| ) | |
| if f'add_controlnet_button_{prefix}' in ui_components: create_controlnet_event_handlers(prefix) | |
| if f'add_ipadapter_button_{prefix}' in ui_components: create_ipadapter_event_handlers(prefix) | |
| if f'add_embedding_button_{prefix}' in ui_components: | |
| create_embedding_event_handlers(prefix) | |
| if f'embeddings_uploads_{prefix}' in ui_components: | |
| emb_uploads = ui_components[f'embeddings_uploads_{prefix}'] | |
| emb_ids = ui_components[f'embeddings_ids_{prefix}'] | |
| emb_sources = ui_components[f'embeddings_sources_{prefix}'] | |
| emb_files = ui_components[f'embeddings_files_{prefix}'] | |
| for i in range(MAX_EMBEDDINGS): | |
| emb_uploads[i].upload( | |
| fn=on_embedding_upload, | |
| inputs=[emb_uploads[i]], | |
| outputs=[emb_ids[i], emb_sources[i], emb_files[i]], | |
| show_progress=False | |
| ) | |
| if f'add_conditioning_button_{prefix}' in ui_components: create_conditioning_event_handlers(prefix) | |
| if f'vae_source_{prefix}' in ui_components: | |
| upload_button = ui_components.get(f'vae_upload_button_{prefix}') | |
| if upload_button: | |
| upload_button.upload( | |
| fn=on_vae_upload, | |
| inputs=[upload_button], | |
| outputs=[ | |
| ui_components[f'vae_id_{prefix}'], | |
| ui_components[f'vae_source_{prefix}'], | |
| ui_components[f'vae_file_{prefix}'] | |
| ] | |
| ) | |
| create_run_event(prefix, task_type) | |
| ui_components["info_get_button"].click( | |
| get_png_info, | |
| [ui_components["info_image_input"]], | |
| [ui_components["info_prompt_output"], ui_components["info_neg_prompt_output"], ui_components["info_params_output"]] | |
| ) | |
| flat_ui_list = [comp for comp_or_list in ui_components.values() for comp in (comp_or_list if isinstance(comp_or_list, list) else [comp_or_list])] | |
| ui_components["send_to_txt2img_button"].click(lambda img: send_info_by_hash(img, ui_components), [ui_components["info_image_input"]], flat_ui_list) | |
| ui_components["send_to_img2img_button"].click(lambda img: send_info_to_tab(img, "img2img", ui_components), [ui_components["info_image_input"]], flat_ui_list) | |
| ui_components["send_to_inpaint_button"].click(lambda img: send_info_to_tab(img, "inpaint", ui_components), [ui_components["info_image_input"]], flat_ui_list) | |
| ui_components["send_to_outpaint_button"].click(lambda img: send_info_to_tab(img, "outpaint", ui_components), [ui_components["info_image_input"]], flat_ui_list) | |
| ui_components["send_to_hires_fix_button"].click(lambda img: send_info_to_tab(img, "hires_fix", ui_components), [ui_components["info_image_input"]], flat_ui_list) | |
| def on_aspect_ratio_change(ratio_key, model_display_name): | |
| model_type = MODEL_TYPE_MAP.get(model_display_name, 'sdxl').lower() | |
| res_map = RESOLUTION_MAP.get(model_type, RESOLUTION_MAP.get("sdxl", {})) | |
| w, h = res_map.get(ratio_key, (1024, 1024)) | |
| return w, h | |
| for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]: | |
| if f'aspect_ratio_{prefix}' in ui_components: | |
| aspect_ratio_dropdown = ui_components[f'aspect_ratio_{prefix}'] | |
| width_component = ui_components[f'width_{prefix}'] | |
| height_component = ui_components[f'height_{prefix}'] | |
| model_dropdown = ui_components[f'base_model_{prefix}'] | |
| aspect_ratio_dropdown.change(fn=on_aspect_ratio_change, inputs=[aspect_ratio_dropdown, model_dropdown], outputs=[width_component, height_component], show_progress=False) | |
| if 'view_mode_inpaint' in ui_components: | |
| def toggle_inpaint_fullscreen_view(view_mode): | |
| is_fullscreen = (view_mode == "Fullscreen View") | |
| other_elements_visible = not is_fullscreen | |
| editor_height = 800 if is_fullscreen else 272 | |
| return { | |
| ui_components['model_and_run_row_inpaint']: gr.update(visible=other_elements_visible), | |
| ui_components['prompts_column_inpaint']: gr.update(visible=other_elements_visible), | |
| ui_components['params_and_gallery_row_inpaint']: gr.update(visible=other_elements_visible), | |
| ui_components['accordion_wrapper_inpaint']: gr.update(visible=other_elements_visible), | |
| ui_components['input_image_dict_inpaint']: gr.update(height=editor_height), | |
| } | |
| output_components = [ | |
| ui_components['model_and_run_row_inpaint'], ui_components['prompts_column_inpaint'], | |
| ui_components['params_and_gallery_row_inpaint'], ui_components['accordion_wrapper_inpaint'], | |
| ui_components['input_image_dict_inpaint'] | |
| ] | |
| ui_components['view_mode_inpaint'].change(fn=toggle_inpaint_fullscreen_view, inputs=[ui_components['view_mode_inpaint']], outputs=output_components, show_progress=False) | |
| def initialize_all_cn_dropdowns(): | |
| cn_config = load_controlnet_config() | |
| if not cn_config: return {} | |
| all_types = sorted(list(set(t for model in cn_config for t in model.get("Type", [])))) | |
| default_type = all_types[0] if all_types else None | |
| series_choices = [] | |
| if default_type: | |
| series_choices = sorted(list(set(model.get("Series", "Default") for model in cn_config if default_type in model.get("Type", [])))) | |
| default_series = series_choices[0] if series_choices else None | |
| filepath = "None" | |
| if default_series and default_type: | |
| for model in cn_config: | |
| if model.get("Series") == default_series and default_type in model.get("Type", []): | |
| filepath = model.get("Filepath") | |
| break | |
| updates = {} | |
| for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]: | |
| if f'controlnet_types_{prefix}' in ui_components: | |
| for type_dd in ui_components[f'controlnet_types_{prefix}']: | |
| updates[type_dd] = gr.update(choices=all_types, value=default_type) | |
| for series_dd in ui_components[f'controlnet_series_{prefix}']: | |
| updates[series_dd] = gr.update(choices=series_choices, value=default_series) | |
| for filepath_state in ui_components[f'controlnet_filepaths_{prefix}']: | |
| updates[filepath_state] = filepath | |
| return updates | |
| def initialize_all_ipa_dropdowns(): | |
| config = load_ipadapter_config() | |
| if not config or not isinstance(config, list): return {} | |
| unified_presets = [] | |
| faceid_presets = [] | |
| for preset_info in config: | |
| name = preset_info.get("preset_name") | |
| if not name: | |
| continue | |
| if "FACEID" in name or "FACE" in name: | |
| faceid_presets.append(name) | |
| else: | |
| unified_presets.append(name) | |
| all_presets = unified_presets + faceid_presets | |
| default_preset = all_presets[0] if all_presets else None | |
| is_faceid_default = default_preset in faceid_presets | |
| lora_strength_update = gr.update(visible=is_faceid_default) | |
| updates = {} | |
| for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]: | |
| if f'ipadapter_final_preset_{prefix}' in ui_components: | |
| for lora_strength_slider in ui_components[f'ipadapter_lora_strengths_{prefix}']: | |
| updates[lora_strength_slider] = lora_strength_update | |
| updates[ui_components[f'ipadapter_final_preset_{prefix}']] = gr.update(choices=all_presets, value=default_preset) | |
| updates[ui_components[f'ipadapter_final_lora_strength_{prefix}']] = lora_strength_update | |
| return updates | |
| def run_on_load(): | |
| cn_updates = initialize_all_cn_dropdowns() | |
| ipa_updates = initialize_all_ipa_dropdowns() | |
| all_updates = {**cn_updates, **ipa_updates} | |
| default_preprocessor = "Canny Edge" | |
| model_update = update_preprocessor_models_dropdown(default_preprocessor) | |
| all_updates[ui_components["preprocessor_model_cn"]] = model_update | |
| settings_outputs = update_preprocessor_settings_ui(default_preprocessor) | |
| dynamic_outputs = ui_components["cn_sliders"] + ui_components["cn_dropdowns"] + ui_components["cn_checkboxes"] | |
| for i, comp in enumerate(dynamic_outputs): | |
| all_updates[comp] = settings_outputs[i] | |
| run_button_update, zero_gpu_update = update_run_button_for_cpu(default_preprocessor) | |
| all_updates[ui_components["run_cn"]] = run_button_update | |
| all_updates[ui_components["zero_gpu_cn"]] = zero_gpu_update | |
| return all_updates | |
| all_load_outputs = [] | |
| for prefix in ["txt2img", "img2img", "inpaint", "outpaint", "hires_fix"]: | |
| if f'controlnet_types_{prefix}' in ui_components: | |
| all_load_outputs.extend(ui_components[f'controlnet_types_{prefix}']) | |
| all_load_outputs.extend(ui_components[f'controlnet_series_{prefix}']) | |
| all_load_outputs.extend(ui_components[f'controlnet_filepaths_{prefix}']) | |
| if f'ipadapter_final_preset_{prefix}' in ui_components: | |
| all_load_outputs.extend(ui_components[f'ipadapter_lora_strengths_{prefix}']) | |
| all_load_outputs.append(ui_components[f'ipadapter_final_preset_{prefix}']) | |
| all_load_outputs.append(ui_components[f'ipadapter_final_lora_strength_{prefix}']) | |
| all_load_outputs.extend([ | |
| ui_components["preprocessor_model_cn"], | |
| *ui_components["cn_sliders"], | |
| *ui_components["cn_dropdowns"], | |
| *ui_components["cn_checkboxes"], | |
| ui_components["run_cn"], | |
| ui_components["zero_gpu_cn"] | |
| ]) | |
| demo.load( | |
| fn=run_on_load, | |
| outputs=all_load_outputs | |
| ) |