import os import sys import argparse import subprocess import contextlib # Apply PyTorch compatibility before Transformers imports TorchAO. import musubi_tuner_gui.torch_compat # noqa: F401 import gradio as gr from musubi_tuner_gui.lora_gui import lora_tab from musubi_tuner_gui.qwen_image_lora_gui import qwen_image_lora_tab from musubi_tuner_gui.wan_lora_gui import wan_lora_tab from musubi_tuner_gui.flux_lora_gui import flux_lora_tab from musubi_tuner_gui.zimage_lora_gui import zimage_lora_tab from musubi_tuner_gui.ideogram4_lora_gui import ideogram4_lora_tab from musubi_tuner_gui.krea2_lora_gui import krea2_lora_tab from musubi_tuner_gui.image_captioning_gui import image_captioning_tab from musubi_tuner_gui.model_quantizer_gui import model_quantizer_tab from musubi_tuner_gui.image_preprocessing_gui import image_preprocessing_tab from musubi_tuner_gui.changelog_gui import version_history_tab from musubi_tuner_gui.lora_extractor_gui import lora_extractor_tab from musubi_tuner_gui.lora_merge_gui import lora_merge_tab from musubi_tuner_gui.lora_convert_gui import lora_convert_tab from musubi_tuner_gui.custom_logging import setup_logging from musubi_tuner_gui.class_gui_config import GUIConfig from musubi_tuner_gui.class_tab_config_manager import TabConfigManager import toml PYTHON = sys.executable project_dir = os.path.dirname(os.path.abspath(__file__)) # Function to read file content, suppressing any FileNotFoundError def read_file_content(file_path): with contextlib.suppress(FileNotFoundError): with open(file_path, "r", encoding="utf8") as file: return file.read() return "" # Function to initialize the Gradio UI interface def initialize_ui_interface(config_manager, headless, release_info, readme_content): # Load custom CSS if available css = read_file_content("./assets/style.css") # Create the main Gradio Blocks interface ui_interface = gr.Blocks(css=css, title="SECourses Musubi Trainer V30.2", theme=gr.themes.Soft()) with ui_interface: # Add title with Patreon link gr.Markdown("# SECourses Musubi Trainer V30.2 : [https://www.patreon.com/posts/137551634](https://www.patreon.com/posts/137551634)") # Create tabs for different functionalities with gr.Tab("Qwen Image Training"): qwen_config = config_manager.get_config_for_tab("qwen_image") qwen_image_lora_tab(headless=headless, config=qwen_config) with gr.Tab("Wan Models Training"): wan_config = config_manager.get_config_for_tab("wan") wan_lora_tab(headless=headless, config=wan_config) with gr.Tab("FLUX Training"): flux_config = config_manager.get_config_for_tab("flux") flux_lora_tab(headless=headless, config=flux_config) with gr.Tab("Z Image Training"): zimage_config = config_manager.get_config_for_tab("zimage") zimage_lora_tab(headless=headless, config=zimage_config) with gr.Tab("Ideogram 4 Training"): ideogram4_config = config_manager.get_config_for_tab("ideogram4") ideogram4_lora_tab(headless=headless, config=ideogram4_config) with gr.Tab("Krea 2 Training"): krea2_config = config_manager.get_config_for_tab("krea2") krea2_lora_tab(headless=headless, config=krea2_config) with gr.Tab("Image Captioning"): captioning_config = config_manager.get_config_for_tab("image_captioning") image_captioning_tab(headless=headless, config=captioning_config) with gr.Tab("Model Quantizer"): quant_config = config_manager.get_config_for_tab("model_quantizer") model_quantizer_tab(headless=headless, config=quant_config) with gr.Tab("LoRA Extractor"): lora_extractor_tab(headless=headless, config=None) with gr.Tab("LoRA Merger"): lora_merge_tab(headless=headless, config=None) with gr.Tab("LoRA Converter"): lora_convert_tab(headless=headless, config=None) with gr.Tab("Image Preprocessing"): preprocessing_config = config_manager.get_config_for_tab("image_preprocessing") image_preprocessing_tab(headless=headless, config=preprocessing_config) with gr.Tab("Version History"): version_history_tab(headless=headless, config=None) with gr.Tab("Musubi Tuner (Deprecated)"): musubi_config = config_manager.get_config_for_tab("musubi_tuner") lora_tab(headless=headless, config=musubi_config) return ui_interface # Function to configure and launch the UI def UI(**kwargs): # Add custom JavaScript if specified log.info(f"headless: {kwargs.get('headless', False)}") # Load release and README information release_info = "v18.0" # Hardcoded version since pyproject.toml is not needed readme_content = read_file_content("./README.md") # Initialize tab-aware configuration manager - default to qwen_image_defaults.toml config_manager = TabConfigManager(config_file_path=kwargs.get("config", "./qwen_image_defaults.toml")) if config_manager.user_loaded_config: log.info(f"Loaded user configuration from '{kwargs.get('config', './qwen_image_defaults.toml')}'...") else: log.info("No user config loaded - will use tab-specific defaults") # Initialize the Gradio UI interface ui_interface = initialize_ui_interface(config_manager, kwargs.get("headless", False), release_info, readme_content) # Construct launch parameters using dictionary comprehension launch_params = { "server_name": kwargs.get("listen"), "auth": (kwargs["username"], kwargs["password"]) if kwargs.get("username") and kwargs.get("password") else None, "server_port": kwargs.get("server_port", 0) if kwargs.get("server_port", 0) > 0 else None, "inbrowser": kwargs.get("inbrowser", True), "share": kwargs.get("share", False), "root_path": kwargs.get("root_path", None), "debug": kwargs.get("debug", False), } # This line filters out any key-value pairs from `launch_params` where the value is `None`, ensuring only valid parameters are passed to the `launch` function. launch_params = {k: v for k, v in launch_params.items() if v is not None} # Launch the Gradio interface with the specified parameters ui_interface.launch(**launch_params) # Function to initialize argument parser for command-line arguments def initialize_arg_parser(): parser = argparse.ArgumentParser() parser.add_argument("--config", type=str, default="./qwen_image_defaults.toml", help="Path to the toml config file for interface defaults") parser.add_argument("--debug", action="store_true", help="Debug on") parser.add_argument("--listen", type=str, default="127.0.0.1", help="IP to listen on for connections to Gradio") parser.add_argument("--username", type=str, default="", help="Username for authentication") parser.add_argument("--password", type=str, default="", help="Password for authentication") parser.add_argument("--server_port", type=int, default=0, help="Port to run the server listener on") parser.add_argument("--inbrowser", action="store_true", default=True, help="Open in browser") parser.add_argument("--share", action="store_true", help="Share the gradio UI") parser.add_argument("--headless", action="store_true", help="Is the server headless") parser.add_argument("--language", type=str, default=None, help="Set custom language") parser.add_argument("--use-ipex", action="store_true", help="Use IPEX environment") parser.add_argument("--use-rocm", action="store_true", help="Use ROCm environment") parser.add_argument("--do_not_use_shell", action="store_true", help="Enforce not to use shell=True when running external commands") parser.add_argument("--do_not_share", action="store_true", help="Do not share the gradio UI") parser.add_argument("--requirements", type=str, default=None, help="requirements file to use for validation") parser.add_argument("--root_path", type=str, default=None, help="`root_path` for Gradio to enable reverse proxy support. e.g. /kohya_ss") parser.add_argument("--noverify", action="store_true", help="Disable requirements verification") return parser if __name__ == "__main__": # Initialize argument parser and parse arguments parser = initialize_arg_parser() args = parser.parse_args() # Set up logging based on the debug flag log = setup_logging(debug=args.debug) # Launch the UI with the provided arguments UI(**vars(args))