examples/dev_tools/webui_train.py

import os, importlib.util, argparse, pkgutil, inspect
from dataclasses import dataclass
import streamlit as st
st.set_page_config(layout="wide")

available_data_file_keys = ["animate_face_video", "animate_pose_video", "audio", "blockwise_controlnet_image", "blockwise_controlnet_inpaint_mask", "context_image", "control_video", "controlnet_image", "controlnet_inpaint_mask", "edit_image", "eligen_entity_masks", "image", "in_context_videos", "infinityou_id_image", "input_audio", "ipadapter_images", "kontext_images", "layer_input_image", "nexus_gen_reference_image", "reference_image", "s2v_pose_video", "step1x_reference_image", "vace_reference_image", "vace_video", "vap_video", "video", "wantodance_keyframes", "wantodance_music_path", "wantodance_reference_image"]
available_extra_inputs = ["animate_face_video", "animate_pose_video", "blockwise_controlnet_image", "blockwise_controlnet_inpaint_mask", "camera_control_direction", "camera_control_speed", "cfg_scale", "context_image", "control_video", "controlnet_image", "controlnet_inpaint_mask", "controlnet_processor_id", "edit_image", "eligen_entity_masks", "eligen_entity_prompts", "end_image", "frame_rate", "framewise_decoding", "in_context_downsample_factor", "in_context_videos", "infinityou_guidance", "infinityou_id_image", "input_audio", "input_image", "ipadapter_images", "kontext_images", "layer_input_image", "layer_num", "lora_encoder_inputs", "motion_bucket_id", "nexus_gen_reference_image", "num_inference_steps", "rand_device", "reference_image", "s2v_pose_video", "seed", "step1x_reference_image", "template_inputs", "vace_reference_image", "vace_video", "value_controller_inputs", "vap_video", "wantodance_fps", "wantodance_keyframes", "wantodance_keyframes_mask", "wantodance_music_path", "wantodance_reference_image"]
available_model_components = ["animate_adapter", "audio_dit", "audio_encoder", "audio_vae", "audio_vae_decoder", "audio_vae_encoder", "audio_vocoder", "blockwise_controlnet", "conditioner", "controlnet", "dinov3_image_encoder", "dit", "dit2", "dual_tower_bridge", "image2lora_coarse", "image2lora_fine", "image2lora_style", "image_encoder", "image_proj_model", "infinityou_processor", "ipadapter", "ipadapter_image_encoder", "lora_encoder", "lora_patcher", "motion_controller", "nexus_gen", "nexus_gen_editing_adapter", "nexus_gen_generation_adapter", "qwenvl", "siglip2_image_encoder", "step1x_connector", "text_encoder", "text_encoder_1", "text_encoder_2", "text_encoder_post_modules", "text_encoder_qwen3", "tokenizer_model", "tokenizer_t5xxl", "unet", "upsampler", "vace", "vace2", "vae", "vae_decoder", "vae_encoder", "value_controller", "vap", "video_dit", "video_dit2", "video_vae", "video_vae_decoder", "video_vae_encoder"]

@dataclass
class Parameter:
    name: str = None
    dtype: type = None
    value: any = None
    required: bool = False
    choices: list = None
    help: str = None

def parse_available_pipelines():
    from diffsynth.diffusion.base_pipeline import BasePipeline
    import diffsynth.pipelines as _pipelines_pkg
    available_pipelines = {}
    for _, name, _ in pkgutil.iter_modules(_pipelines_pkg.__path__):
        mod = importlib.import_module(f"diffsynth.pipelines.{name}")
        classes = {
            cls_name: cls for cls_name, cls in inspect.getmembers(mod, inspect.isclass)
            if issubclass(cls, BasePipeline) and cls is not BasePipeline and cls.__module__ == mod.__name__
        }
        available_pipelines.update(classes)
    return available_pipelines

def search_for_options(name):
    files = search_for_files("examples", ".sh")
    params = set()
    for file in files:
        with open(file, "r", encoding="utf-8") as f:
            for line in f.readlines():
                if f"--{name}" in line and not line.startswith("#"):
                    line = line.strip()
                    line = line.replace(" \\", "").replace(f"--{name}", "").replace('"', "").replace(" ", "")
                    for param in line.split(","):
                        params.add(param)
    for param in sorted(list(params)):
        print(f'"{param}", ', end="")
    print()

def search_for_available_pipeline_options():
    from diffsynth.diffusion.base_pipeline import BasePipeline
    pipeline_classes = parse_available_pipelines()
    base_attrs = set(vars(BasePipeline()))
    black_list = ["tokenizer", "processor", "tokenizer_1", "tokenizer_2", "audio_processor"]
    options = []
    for pipeline_class in pipeline_classes:
        pipe = pipeline_classes[pipeline_class]()
        members = [attr for attr in vars(pipe) if not attr.startswith("__") and attr not in base_attrs and getattr(pipe, attr) is None]
        members = [attr for attr in members if attr not in black_list]
        options.extend(members)
    options = sorted(list(set(options)))
    for option in sorted(list(options)):
        print(f'"{option}", ', end="")
    print()

def parse_available_training_scripts(path):
    training_scripts = {}
    for folder in os.listdir(path):
        if os.path.isfile(f"{path}/{folder}/model_training/train.py"):
            training_scripts[folder] = f"{path}/{folder}/model_training/train.py"
    return training_scripts

def search_for_files(path, suffix):
    if os.path.isfile(path):
        if path.endswith(suffix): return [path]
        return []
    else:
        files = []
        for sub_path in os.listdir(path):
            files.extend(search_for_files(os.path.join(path, sub_path), suffix))
        return files

def parse_available_examples(path):
    path = os.path.dirname(path)
    examples = search_for_files(path, ".sh")
    return examples

def parse_example(example_path):
    value_dict = {}
    with open(example_path, "r", encoding="utf-8") as f:
        for line in f.readlines():
            line = line.strip()
            if line.startswith("#"):
                continue
            if not line.startswith("--"):
                continue
            line = line.replace("\\", "").strip()
            if " " in line:
                name, value = line[2:line.index(" ")], line[line.index(" ") + 1:]
                if value.startswith('"') and value.endswith('"'):
                    value = value[1:-1]
                if value.startswith("'") and value.endswith("'"):
                    value = value[1:-1]
            else:
                name, value = line[2:], True
            value_dict[name] = value
    return value_dict

def parse_parser(path):
    spec = importlib.util.spec_from_file_location("train", path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    for name in dir(module):
        if name.endswith("parser") and callable(getattr(module, name)):
            return getattr(module, name)
    return None

def parse_parser_action(action, value=None):
    if isinstance(action, argparse._StoreTrueAction) or isinstance(action, argparse._StoreFalseAction):
        dtype = bool
    else:
        dtype = action.type
    param = Parameter(
        name=action.dest,
        dtype=dtype,
        value=action.default if value is None else dtype(value),
        required=action.required,
        choices=action.choices,
        help=action.help,
    )
    return param

def parse_parser_actions(parser, example_path=None):
    value_dict = {} if example_path is None or example_path == "None" else parse_example(example_path)
    params = []
    for action in parser._actions:
        param = parse_parser_action(action, value=value_dict.get(action.dest))
        if param.name == "help":
            continue
        params.append(param)
    return params

def draw_model_id_with_origin_paths(param, disabled=False):
    with st.container(border=True):
        st.markdown(param.name, help=param.help)
        model_id_with_origin_paths = [] if param.value is None else param.value.split(",")
        num = st.number_input(f"Number of models", min_value=0, max_value=20, value=len(model_id_with_origin_paths), disabled=disabled)
        result = []
        for i in range(num):
            col1, col2 = st.columns(2)
            value = model_id_with_origin_paths[i].split(":") if i < len(model_id_with_origin_paths) else (None, None)
            with col1:
                model_id = st.text_input("model_id", value=value[0], key=f"model_id_{i}", disabled=disabled)
            with col2:
                origin_file_pattern = st.text_input("origin_file_pattern", value=value[1], key=f"origin_file_pattern_{i}", disabled=disabled)
            result.append(f"{model_id}:{origin_file_pattern}")
        result = ",".join(result)
    return result

def draw_parameter(param, check_enable=True, disabled=False):
    if check_enable and param.value is None:
        with st.container(border=True):
            enable_button = st.checkbox(f"Enable {param.name}", value=False, disabled=disabled)
            ui = draw_parameter(param, check_enable=False, disabled=disabled or not enable_button)
            if enable_button:
                return ui
            else:
                return None
    if param.name == "data_file_keys":
        ui = st.multiselect(param.name, options=available_data_file_keys, accept_new_options=True, default=param.value.split(","), disabled=disabled, help=param.help)
        ui = ",".join(ui)
    elif param.name == "model_paths":
        ui = st.text_area(param.name, value=param.value, height=3, disabled=disabled, help=param.help)
    elif param.name == "model_id_with_origin_paths":
        ui = draw_model_id_with_origin_paths(param, disabled=disabled)
    elif param.name == "extra_inputs":
        value = None if param.value is None else param.value.split(",")
        ui = st.multiselect(param.name, options=available_extra_inputs, accept_new_options=True, default=value, disabled=disabled, help=param.help)
        ui = ",".join(ui)
    elif param.name in ["fp8_models", "offload_models", "trainable_models", "lora_base_model", "preset_lora_model"]:
        value = None if param.value is None else param.value.split(",")
        ui = st.multiselect(param.name, options=available_model_components, accept_new_options=True, default=value, disabled=disabled, help=param.help)
        ui = ",".join(ui)
    elif param.name == "learning_rate":
        ui = st.number_input(param.name, value=param.value, format="%0.7f", step=1e-4, disabled=disabled, help=param.help)
    elif param.dtype == str:
        ui = st.text_input(param.name, value=param.value, disabled=disabled, help=param.help)
    elif param.dtype == int:
        ui = st.number_input(param.name, value=param.value, step=1, disabled=disabled, help=param.help)
    elif param.dtype == float:
        ui = st.number_input(param.name, value=param.value, disabled=disabled, help=param.help)
    elif param.dtype == bool:
        ui = st.checkbox(param.name, value=param.value, disabled=disabled, help=param.help)
    else:
        st.markdown(f"(`{param.name}` is not not configurable in WebUI). dtype: `{param.dtype}`.")
        ui = None
    return ui

def draw_dataset_configs(dataset_base_path, dataset_metadata_path, dataset_repeat, dataset_num_workers, data_file_keys):
    dataset_base_path = draw_parameter(dataset_base_path)
    dataset_metadata_path = draw_parameter(dataset_metadata_path)
    col_1, col_2 = st.columns(2)
    with col_1:
        dataset_repeat = draw_parameter(dataset_repeat)
    with col_2:
        dataset_num_workers = draw_parameter(dataset_num_workers)
    data_file_keys = draw_parameter(data_file_keys)
    params = {
        "dataset_base_path": dataset_base_path,
        "dataset_metadata_path": dataset_metadata_path,
        "dataset_repeat": dataset_repeat,
    }
    if dataset_num_workers > 0: params["dataset_num_workers"] = dataset_num_workers
    params["data_file_keys"] = data_file_keys
    return params

def draw_image_size(height, width, max_pixels):
    mode = st.selectbox("Image scaling and cropping", options=["Scale if pixel count exceeds threshold", "Resize to a fixed size and crop"], index=int(max_pixels.value is None))
    use_max_pixel = mode == "Scale if pixel count exceeds threshold"
    use_height_width = mode == "Resize to a fixed size and crop"
    col1, col2 = st.columns(2)
    with col1:
        with st.container(border=True):
            max_pixels = draw_parameter(max_pixels, disabled=not use_max_pixel)
    with col2:
        with st.container(border=True):
            height = draw_parameter(height, disabled=not use_height_width, check_enable=False)
            width = draw_parameter(width, disabled=not use_height_width, check_enable=False)
    if use_max_pixel:
        return {"max_pixels": max_pixels}
    else:
        return {"height": height, "width": width}

def draw_model_configs(model_paths, model_id_with_origin_paths, extra_inputs, fp8_models, offload_models):
    model_id_with_origin_paths = draw_parameter(model_id_with_origin_paths)
    model_paths = draw_parameter(model_paths)
    extra_inputs = draw_parameter(extra_inputs)
    fp8_models = draw_parameter(fp8_models)
    offload_models = draw_parameter(offload_models)
    params = {}
    if model_paths is not None: params["model_paths"] = model_paths
    if model_id_with_origin_paths is not None: params["model_id_with_origin_paths"] = model_id_with_origin_paths
    if extra_inputs is not None: params["extra_inputs"] = extra_inputs
    if fp8_models is not None: params["fp8_models"] = fp8_models
    if offload_models is not None: params["offload_models"] = offload_models
    return params

def draw_video_size(height, width, max_pixels, num_frames):
    mode = st.selectbox("Video scaling and cropping", options=["Scale if pixel count exceeds threshold", "Resize to a fixed size and crop"], index=int(max_pixels.value is not None))
    use_max_pixel = mode == "Scale if pixel count exceeds threshold"
    use_height_width = mode == "Resize to a fixed size and crop"
    col1, col2 = st.columns(2)
    with col1:
        with st.container(border=True):
            max_pixels = draw_parameter(max_pixels, disabled=not use_max_pixel)
    with col2:
        with st.container(border=True):
            height = draw_parameter(height, disabled=not use_height_width)
            width = draw_parameter(width, disabled=not use_height_width)
    num_frames = draw_parameter(num_frames)
    if use_max_pixel:
        return {"max_pixels": max_pixels, "num_frames": num_frames}
    else:
        return {"height": height, "width": width, "num_frames": num_frames}

def draw_training_configs(learning_rate, num_epochs, trainable_models, find_unused_parameters, weight_decay, task):
    learning_rate = draw_parameter(learning_rate)
    num_epochs = draw_parameter(num_epochs)
    trainable_models = draw_parameter(trainable_models)
    weight_decay = draw_parameter(weight_decay)
    task = draw_parameter(task)
    find_unused_parameters = draw_parameter(find_unused_parameters)
    params = {
        "task": task,
        "find_unused_parameters": find_unused_parameters,
        "num_epochs": num_epochs,
        "learning_rate": learning_rate,
    }
    if weight_decay != 0.01: params["weight_decay"] = weight_decay
    if trainable_models is not None: params["trainable_models"] = trainable_models
    return params

def draw_output_configs(output_path, remove_prefix_in_ckpt, save_steps):
    output_path = draw_parameter(output_path)
    remove_prefix_in_ckpt = draw_parameter(remove_prefix_in_ckpt)
    save_steps = draw_parameter(save_steps)
    params = {
        "output_path": output_path,
        "remove_prefix_in_ckpt": remove_prefix_in_ckpt,
    }
    if save_steps is not None: params["save_steps"] = save_steps
    return params

def draw_lora_configs(lora_base_model, lora_target_modules, lora_rank, lora_checkpoint, preset_lora_path, preset_lora_model):
    with st.container(border=True):
        train_lora = st.checkbox("Train LoRA", value=lora_base_model is not None)
        lora_base_model = draw_parameter(lora_base_model, check_enable=False, disabled=not train_lora)
        lora_target_modules = draw_parameter(lora_target_modules, check_enable=False, disabled=not train_lora)
        lora_rank = draw_parameter(lora_rank, check_enable=False, disabled=not train_lora)
        lora_checkpoint = draw_parameter(lora_checkpoint, check_enable=False, disabled=not train_lora)
    preset_lora_path = draw_parameter(preset_lora_path)
    preset_lora_model = draw_parameter(preset_lora_model)
    params = {}
    if train_lora:
        if lora_base_model is not None: params["lora_base_model"] = lora_base_model
        if lora_target_modules is not None: params["lora_target_modules"] = lora_target_modules
        if lora_rank is not None: params["lora_rank"] = lora_rank
        if lora_checkpoint is not None: params["lora_checkpoint"] = lora_checkpoint
    if preset_lora_path is not None: params["preset_lora_path"] = preset_lora_path
    if preset_lora_model is not None: params["preset_lora_model"] = preset_lora_model
    return params

def draw_gradient_configs(use_gradient_checkpointing, use_gradient_checkpointing_offload, gradient_accumulation_steps):
    use_gradient_checkpointing = draw_parameter(use_gradient_checkpointing)
    use_gradient_checkpointing_offload = draw_parameter(use_gradient_checkpointing_offload)
    gradient_accumulation_steps = draw_parameter(gradient_accumulation_steps)
    params = {
        "use_gradient_checkpointing": use_gradient_checkpointing,
        "use_gradient_checkpointing_offload": use_gradient_checkpointing_offload,
    }
    if gradient_accumulation_steps != 1:
        params["gradient_accumulation_steps"] = gradient_accumulation_steps
    return params

def draw_template_model_configs(template_model_id_or_path, enable_lora_hot_loading):
    template_model_id_or_path = draw_parameter(template_model_id_or_path)
    enable_lora_hot_loading = draw_parameter(enable_lora_hot_loading)
    params = {"enable_lora_hot_loading": enable_lora_hot_loading}
    if template_model_id_or_path is not None: params["template_model_id_or_path"] = template_model_id_or_path
    return params

def match_ui_groups(params, ui_groups):
    param_names = [param.name for param in params]
    for ui_group in ui_groups:
        if sum([name in param_names for name in ui_group["params"]]) == len(ui_group["params"]):
            group_params = {param.name: param for param in params if param.name in ui_group["params"]}
            other_params = [param for param in params if param.name not in ui_group["params"]]
            return group_params, other_params, ui_group
    return {}, params, None

def draw_other_params(params):
    results = {}
    for param in params:
        results[param.name] = draw_parameter(param)
    return results

def draw_all_params(params, ui_groups):
    matched_ui_groups = []
    while True:
        group_params, other_params, ui_group = match_ui_groups(params, ui_groups)
        if len(group_params) == 0:
            break
        matched_ui_groups.append((ui_group, group_params))
        params = other_params
    tabs = st.tabs([ui_group["name"] for ui_group, _ in matched_ui_groups] + ["Others"])
    inputs = {}
    for tab, (ui_group, group_params) in zip(tabs, matched_ui_groups):
        with tab:
            inputs.update(ui_group["fn"](**group_params))
    with tabs[-1]:
        inputs.update(draw_other_params(params))
    return inputs

def generate_training_script(script_path, inputs):
    cmd = f"accelerate launch {script_path}"
    for name, value in inputs.items():
        if value is not None:
            if isinstance(value, bool):
                if value == True:
                    cmd = f"{cmd} \\\n  --{name}"
            elif isinstance(value, str):
                cmd = f"{cmd} \\\n  --{name} \"{value}\""
            else:
                cmd = f"{cmd} \\\n  --{name} {value}"
    return cmd

ui_groups = [
    {
        "name": "Dataset",
        "params": ("dataset_base_path", "dataset_metadata_path", "dataset_repeat", "dataset_num_workers", "data_file_keys"),
        "fn": draw_dataset_configs,
    },
    {
        "name": "Video Size",
        "params": ("height", "width", "max_pixels", "num_frames"),
        "fn": draw_video_size,
    },
    {
        "name": "Image Size",
        "params": ("height", "width", "max_pixels"),
        "fn": draw_image_size,
    },
    {
        "name": "Model",
        "params": ("model_paths", "model_id_with_origin_paths", "extra_inputs", "fp8_models", "offload_models"),
        "fn": draw_model_configs,
    },
    {
        "name": "Training",
        "params": ("learning_rate", "num_epochs", "trainable_models", "find_unused_parameters", "weight_decay", "task"),
        "fn": draw_training_configs,
    },
    {
        "name": "Output",
        "params": ("output_path", "remove_prefix_in_ckpt", "save_steps"),
        "fn": draw_output_configs,
    },
    {
        "name": "LoRA",
        "params": ("lora_base_model", "lora_target_modules", "lora_rank", "lora_checkpoint", "preset_lora_path", "preset_lora_model"),
        "fn": draw_lora_configs,
    },
    {
        "name": "Gradient",
        "params": ("use_gradient_checkpointing", "use_gradient_checkpointing_offload", "gradient_accumulation_steps"),
        "fn": draw_gradient_configs,
    },
    {
        "name": "Templates",
        "params": ("template_model_id_or_path", "enable_lora_hot_loading"),
        "fn": draw_template_model_configs,
    },
]

def launch_webui():
    input_col, output_col = st.columns(2)
    with input_col:
        if "available_training_scripts" not in st.session_state:
            st.session_state["available_training_scripts"] = parse_available_training_scripts("examples")
        with st.container(border=True):
            script_path = st.selectbox(label="Script path", options=st.session_state["available_training_scripts"].values(), index=0)
            example_path = st.selectbox(label="Example path (Optional)", options=["None"] + parse_available_examples(script_path), index=0)
        if st.button("Step 1: Parse Training Script", type="primary"):
            st.session_state["script_path"] = script_path
        
        if "script_path" not in st.session_state:
            return
        with st.spinner("Fetching input parameters", show_time=False):
            parser = parse_parser(script_path)
        parser = parser()
        params = parse_parser_actions(parser, example_path)
        inputs = draw_all_params(params, ui_groups)
    with output_col:
        if st.button("Step 2: Generate training script", type="primary"):
            script = generate_training_script(script_path, inputs)
            st.code(script, language="shell")

launch_webui()