From 94c99ba7a20306e196f6fa0ce658380d54a2f2b9 Mon Sep 17 00:00:00 2001 From: hndrr Date: Wed, 16 Jul 2025 17:18:08 +0900 Subject: [PATCH] Add WebUI functionality for training with GPU support and enhance main function for config file handling --- run_modal_docker.py | 119 +++++++++++++++++++++++-- run_modal_docker_bk.py | 193 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 306 insertions(+), 6 deletions(-) create mode 100644 run_modal_docker_bk.py diff --git a/run_modal_docker.py b/run_modal_docker.py index 06fb9f84..233269b2 100644 --- a/run_modal_docker.py +++ b/run_modal_docker.py @@ -70,9 +70,9 @@ def print_end_message(jobs_completed, jobs_failed): timeout=7200, # 2 hours, increase or decrease if needed secrets=[modal.Secret.from_name("huggingface-secret")] ) -def main(config_file_list_str: str, recover: bool = False, name: str = None): +def main(config_file_list_str: str, recover: bool = False, name: str = None): # convert the config file list from a string to a list - config_file_list = config_file_list_str.split(",") + config_file_list = str(config_file_list_str).split(",") jobs_completed = 0 jobs_failed = 0 @@ -104,6 +104,102 @@ def main(config_file_list_str: str, recover: bool = False, name: str = None): print_end_message(jobs_completed, jobs_failed) +# WebUIアクセス用の新しい関数を追加（GPU対応でトレーニング可能） +@app.function( + image=image, + gpu="A100", # ← GPUを有効化！WebUIからトレーニングできるように + cpu=4, + memory=32768, # 32GB + timeout=7200, # 2時間 + volumes={MOUNT_DIR: model_volume}, # ← ボリュームもマウント + secrets=[modal.Secret.from_name("huggingface-secret")] +) +def webui(): + import subprocess + import time + import os + + # Hugging Faceにログイン + try: + from huggingface_hub import login + hf_token = os.environ.get("HF_TOKEN") + if hf_token: + login(token=hf_token) + print("✅ Successfully logged in to Hugging Face") + else: + print("⚠️ No HF_TOKEN found") + except Exception as e: + print(f"⚠️ Failed to login to Hugging Face: {e}") + + # AI-toolkit UIディレクトリに移動 + os.chdir("/app/ai-toolkit/ui") + + # トレーニング結果保存先を作成 + os.makedirs(MOUNT_DIR, exist_ok=True) + + # modal.forwardでポート8675を公開 + with modal.forward(8675) as tunnel: + print(f"🌐 AI Toolkit WebUI is accessible at: {tunnel.url}") + print(f"🔧 Starting UI server...") + print(f"💾 Training outputs will be saved to: {MOUNT_DIR}") + + # npm run startでWebUIを起動 + env_vars = os.environ.copy() + if hf_token: + env_vars.update({ + "HF_TOKEN": hf_token, + "HUGGINGFACE_HUB_TOKEN": hf_token, + "HF_API_TOKEN": hf_token + }) + + process = subprocess.Popen( + ["npm", "run", "start"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd="/app/ai-toolkit/ui", + env=env_vars # 環境変数を明示的に渡す + ) + + # サーバーが起動するまで待機 + print("⏳ Waiting for server to start...") + time.sleep(30) + print(f"✅ WebUI should now be accessible at: {tunnel.url}") + print(f"🚀 You can now start training jobs from the WebUI!") + + # 2時間維持（トレーニング時間を考慮） + try: + while process.poll() is None: + # プロセスのログを定期的に出力 + try: + stdout_line = process.stdout.readline() + if stdout_line: + print(f"[WebUI] {stdout_line.decode().strip()}") + + stderr_line = process.stderr.readline() + if stderr_line: + print(f"[WebUI ERROR] {stderr_line.decode().strip()}") + except: + pass + + time.sleep(10) + # 定期的にボリュームをコミット + model_volume.commit() + except KeyboardInterrupt: + print("🛑 Shutting down WebUI...") + process.terminate() + finally: + # プロセス終了時の最終ログ + if process.poll() is not None: + stdout, stderr = process.communicate() + if stdout: + print(f"[WebUI Final STDOUT] {stdout.decode()}") + if stderr: + print(f"[WebUI Final STDERR] {stderr.decode()}") + + # 最終コミット + model_volume.commit() + print("💾 Final volume commit completed") + if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -129,9 +225,20 @@ if __name__ == "__main__": default=None, help='Name to replace [name] tag in config file, useful for shared config file' ) + + # WebUIオプションを追加 + parser.add_argument( + '--webui', + action='store_true', + help='Launch WebUI instead of training' + ) + args = parser.parse_args() - # convert list of config files to a comma-separated string for Modal compatibility - config_file_list_str = ",".join(args.config_file_list) - - main.call(config_file_list_str=config_file_list_str, recover=args.recover, name=args.name) + if args.webui: + # WebUI起動 + webui.remote() + else: + # 従来通りのトレーニング実行 + config_file_list_str = ",".join(args.config_file_list) + main.remote(config_file_list_str=config_file_list_str, recover=args.recover, name=args.name) \ No newline at end of file diff --git a/run_modal_docker_bk.py b/run_modal_docker_bk.py new file mode 100644 index 00000000..1ff62c6a --- /dev/null +++ b/run_modal_docker_bk.py @@ -0,0 +1,193 @@ +''' + +ostris/ai-toolkit on https://modal.com +Run training with the following command: +modal run run_modal.py --config-file-list-str=/app/ai-toolkit/config/whatever_you_want.yml + +''' + +import os +os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" +import sys +import modal +from dotenv import load_dotenv +# Load the .env file if it exists +load_dotenv() + +sys.path.insert(0, "/app/ai-toolkit") +# must come before ANY torch or fastai imports +# import toolkit.cuda_malloc + +# turn off diffusers telemetry until I can figure out how to make it opt-in +os.environ['DISABLE_TELEMETRY'] = 'YES' + +# define the volume for storing model outputs, using "creating volumes lazily": https://modal.com/docs/guide/volumes +# you will find your model, samples and optimizer stored in: https://modal.com/storage/your-username/main/flux-lora-models +model_volume = modal.Volume.from_name("flux-lora-models", create_if_missing=True) + +# modal_output, due to "cannot mount volume on non-empty path" requirement +MOUNT_DIR = "/app/ai-toolkit/modal_output" # modal_output, due to "cannot mount volume on non-empty path" requirement + +# define modal app +image = modal.Image.from_dockerfile("docker/Dockerfile").add_local_dir("config", "/app/ai-toolkit/config").add_local_dir("input/images", '/app/ai-toolkit/input/images') + +# mount for the entire ai-toolkit directory +# example: "/Users/username/ai-toolkit" is the local directory, "/root/ai-toolkit" is the remote directory +# code_mount = modal.Mount.from_local_dir("D:/ai-toolkit", remote_path="/root/ai-toolkit") + +# create the Modal app with the necessary mounts and volumes +# app = modal.App(name="flux-lora-training", image=image, mounts=[code_mount], volumes={MOUNT_DIR: model_volume}) +app = modal.App(name="flux-lora-training", image=image, volumes={MOUNT_DIR: model_volume}) + +# Check if we have DEBUG_TOOLKIT in env +if os.environ.get("DEBUG_TOOLKIT", "0") == "1": + # Set torch to trace mode + import torch + torch.autograd.set_detect_anomaly(True) + +import argparse +from toolkit.job import get_job + +def print_end_message(jobs_completed, jobs_failed): + failure_string = f"{jobs_failed} failure{'' if jobs_failed == 1 else 's'}" if jobs_failed > 0 else "" + completed_string = f"{jobs_completed} completed job{'' if jobs_completed == 1 else 's'}" + + print("") + print("========================================") + print("Result:") + if len(completed_string) > 0: + print(f" - {completed_string}") + if len(failure_string) > 0: + print(f" - {failure_string}") + print("========================================") + + +@app.function( + # request a GPU with at least 24GB VRAM + # more about modal GPU's: https://modal.com/docs/guide/gpu + gpu="A100", # gpu="H100" + # more about modal timeouts: https://modal.com/docs/guide/timeouts + timeout=7200, # 2 hours, increase or decrease if needed + secrets=[modal.Secret.from_name("huggingface-secret")] +) + +def main(config_file_list_str: str, recover: bool = False, name: str = None): + # convert the config file list from a string to a list + config_file_list = config_file_list_str.split(",") + + jobs_completed = 0 + jobs_failed = 0 + + print(f"Running {len(config_file_list)} job{'' if len(config_file_list) == 1 else 's'}") + + for config_file in config_file_list: + try: + job = get_job(config_file, name) + + job.config['process'][0]['training_folder'] = MOUNT_DIR + os.makedirs(MOUNT_DIR, exist_ok=True) + print(f"Training outputs will be saved to: {MOUNT_DIR}") + + # run the job + job.run() + + # commit the volume after training + model_volume.commit() + + job.cleanup() + jobs_completed += 1 + except Exception as e: + print(f"Error running job: {e}") + jobs_failed += 1 + if not recover: + print_end_message(jobs_completed, jobs_failed) + raise e + + print_end_message(jobs_completed, jobs_failed) + +# WebUIアクセス用の新しい関数を追加（GPU対応でトレーニング可能） +@app.function( + image=image, + gpu="A100", + cpu=4, + memory=32768, # 32GB + timeout=7200, # 2時間 + volumes={MOUNT_DIR: model_volume} +) +def webui(): + import subprocess + import time + import os + + # AI-toolkit UIディレクトリに移動 + os.chdir("/app/ai-toolkit/ui") + + # トレーニング結果保存先を作成 + os.makedirs(MOUNT_DIR, exist_ok=True) + + # modal.forwardでポート8675を公開 + with modal.forward(8675) as tunnel: + print(f"🌐 AI Toolkit WebUI is accessible at: {tunnel.url}") + print(f"🔧 Starting UI server...") + print(f"💾 Training outputs will be saved to: {MOUNT_DIR}") + + # npm run startでWebUIを起動 + process = subprocess.Popen( + ["npm", "run", "start"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + cwd="/app/ai-toolkit/ui", + env={**os.environ, "TRAINING_FOLDER": MOUNT_DIR} + ) + + # サーバーが起動するまで待機 + print("⏳ Waiting for server to start...") + time.sleep(30) + print(f"✅ WebUI should now be accessible at: {tunnel.url}") + print(f"🚀 You can now start training jobs from the WebUI!") + + # 2時間維持（トレーニング時間を考慮） + try: + while process.poll() is None: + time.sleep(60) + # 定期的にボリュームをコミット + model_volume.commit() + except KeyboardInterrupt: + print("🛑 Shutting down WebUI...") + process.terminate() + finally: + # 最終コミット + model_volume.commit() + print("💾 Final volume commit completed") + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + # require at least one config file + parser.add_argument( + 'config_file_list', + nargs='+', + type=str, + help='Name of config file (eg: person_v1 for config/person_v1.json/yaml), or full path if it is not in config folder, you can pass multiple config files and run them all sequentially' + ) + + # flag to continue if a job fails + parser.add_argument( + '-r', '--recover', + action='store_true', + help='Continue running additional jobs even if a job fails' + ) + + # optional name replacement for config file + parser.add_argument( + '-n', '--name', + type=str, + default=None, + help='Name to replace [name] tag in config file, useful for shared config file' + ) + args = parser.parse_args() + + # convert list of config files to a comma-separated string for Modal compatibility + config_file_list_str = ",".join(args.config_file_list) + + main.call(config_file_list_str=config_file_list_str, recover=args.recover, name=args.name)