LoRA Studio · 端到端 LoRA 训练管线
桌面应用

# launcher 被强结时把所有子进程一起带走 (KILL_ON_JOB_CLOSE)
def _job_bind_children() -> None:
    global _job_handle
    JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x2000
    JobObjectExtendedLimitInformation = 9

    h = ctypes.windll.kernel32.CreateJobObjectW(None, None)
    if not h:
        return

    # 设置 ExtendedLimitInfo, 让 launcher 进程退出时 OS 清理子进程
    info = JOBOBJECT_EXTENDED_LIMIT_INFORMATION()
    info.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE
    ctypes.windll.kernel32.SetInformationJobObject(
        h, JobObjectExtendedLimitInformation,
        ctypes.byref(info), ctypes.sizeof(info))

    # 把当前进程绑进 job, 后续 spawn 的子进程自动继承
    ctypes.windll.kernel32.AssignProcessToJobObject(
        h, ctypes.windll.kernel32.GetCurrentProcess())
    _job_handle = h

# ----- SSE 实时日志 -----
@app.get("/api/events")
async def api_events(request: Request):
    async def gen():
        last = 0
        yield "retry: 2000\n\n"          # 客户端断线 2s 重连
        while True:
            if await request.is_disconnected():
                break
            new, nxt, phase, stages, error = STATE.logs_since(last)
            last = nxt
            payload = {
                "logs": new, "phase": phase, "stages": stages,
                "error": error,
                "has_report": bool(STATE.quality_report_path),
                "result": dict(STATE.result),
            }
            yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
            await asyncio.sleep(0.4)

    return StreamingResponse(gen(), media_type="text/event-stream")

def compute_train_params(img_count, gpu_vram_gb, arch_family,
                         subject_type, quality_tier,
                         overrides=None, single_image=False, photoreal=False):
    img_count = max(1, int(img_count))
    vram = int(gpu_vram_gb or 0)
    is_style = (subject_type == "style")

    # 分辨率 × batch 阶梯, 防 OOM (SDXL 24G 报 23 → 阈值 22)
    if arch_family == "flux":
        reso, batch = 1024, 1
    elif arch_family == "sdxl":
        reso, batch = 640, 1
        for mn, r, b in [(48,1024,6), (40,1024,4),
                          (30,1024,3), (22,1024,2),
                          (14,768,1), (0,640,1)]:
            if vram >= mn:
                reso, batch = r, b
                break

    # 步数 · dim · 学习率: 风格 vs 主体 / 通用 vs photoreal 分档
    spi = {"fast": 10, "balanced": 20, "hq": 35}[quality_tier]
    lo, hi = {"hq": (1500, 3500)}.get(quality_tier, (500, 2500))
    steps = max(lo, min(hi, spi * img_count))
    dim   = 32 if quality_tier == "hq" else 16
    unet_lr = "1e-4" if photoreal else "3e-4"   # 防写实过拟合
    return {"resolution": reso, "batch": batch, "steps": steps,
            "network_dim": dim, "unet_lr": unet_lr, ...}

# 9 个工具中, 烧钱/动 pod/删除 的 4 个 action 不进白名单
# LLM 永远只能 "request_confirm", 由前端弹原生确认框, 用户亲点才执行

def _exec_tool(name, args, ctx):
    if name == "get_pipeline_state":        # 读 · 安全
        return json.dumps(STATE.snapshot(), ensure_ascii=False)

    if name == "preview_train_params":      # 读 · 安全
        params = param_engine.compute_train_params(...)
        return json.dumps({"params": params}, ensure_ascii=False)

    if name == "set_form_fields":           # 改 UI · 字段白名单
        applied, rejected = {}, []
        for u in args.get("updates") or []:
            field = str((u or {}).get("field", ""))
            if field not in SETTABLE_FIELDS:
                rejected.append(field); continue
            ctx["form_actions"].append({"field": field, "value": ...})
        return json.dumps({"applied": applied, "rejected": rejected})

    if name == "request_confirm":           # 危险动作 · 不执行只登记
        action = str(args.get("action", ""))
        if action not in ("start", "test_pod",
                          "confirm_train", "abort"):
            return json.dumps({"ok": False, "error": "unknown action"})
        ctx["confirm"] = {"action": action, "summary": args.get("summary")}
        return json.dumps({"ok": True,
            "note": "确认请求已登记并弹给用户。你不能自己执行。"})