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