LoRAMaster/main.py at main · AIMixer/LoRAMaster · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
from nicegui import ui
from layout import with_layout  # 如果封装在 layout.py
import  wan_lora_train
import about
import framepack_lora_train
import  hunyuan_lora_train
import kontext_lora_train
import qwen_image_lora_train
import flux_lora_train
import  sys,os
sys.path.append(os.path.join(os.path.dirname(__file__), 'musubi-tuner'))
import subprocess
import atexit
from dataset_manager import AutoCaptioning
from dataset_manager import image_convert
from dataset_manager import image_rename

# 启动 tensorboard（后台运行）
tensorboard_process = subprocess.Popen([
    'tensorboard',
    '--logdir', './logs',     # 你的训练日志路径
    '--port', '6006',
    '--host', '0.0.0.0'
])

@atexit.register
def cleanup():
    tensorboard_process.terminate()


@ui.page('/')
def dashboard_page():
    def content():
        about.draw_ui()
    with_layout(content)


@ui.page('/Wan')
def wan_train_page():
    def content():
        wan_lora_train.draw_ui()
    with_layout(content)

@ui.page('/FramePack')
def framepack_train_page():
    def content():
        framepack_lora_train.draw_ui()
    with_layout(content)

@ui.page('/HunyuanVideo')
def hunyuan_train_page():
    def content():
        hunyuan_lora_train.draw_ui()
    with_layout(content)

@ui.page('/FluxKontext')
def kontext_train_page():
    def content():
        kontext_lora_train.draw_ui()
    with_layout(content)

@ui.page('/Flux')
def flux_train_page():
    def content():
        flux_lora_train.draw_ui()
    with_layout(content)
@ui.page('/QwenImage')
def qwen_image_page():
    def content():
        qwen_image_lora_train.draw_ui()
    with_layout(content)

@ui.page('/Tensorboard')
def settings_page():
    def content():
        ui.html(
            '<iframe src="http://127.0.0.1:6006" '
            'style="width:100%; height:100%; border:none;"></iframe>'
        ).style('width:100%;height:100vh;')
    with_layout(content)

@ui.page('/AutoCaptioning')
def auto_captioning_page():
    def content():
        AutoCaptioning.draw_ui()
    with_layout(content)
@ui.page('/ImageConvert')
def ImageConvert():
    def content():
        image_convert.draw_ui()
    with_layout(content)
@ui.page('/ImageRename')
def ImageRename():
    def content():
        image_rename.draw_ui()
    with_layout(content)
# --------------------------
# 启动默认首页
# --------------------------
ui.run(reload=False,title='LoRA训练大师 - by AI搅拌手',port=8080)