"""The training run: the set's media become latents, its captions become conditioning, ComfyUI's own trainer does the learning.

Blocking, and run on ComfyUI's worker. Everything that needs torch or comfy is imported inside the
functions, so the rest of the module (sets, tags, the editor's routes) loads without them.
"""

from __future__ import annotations

import os
from pathlib import Path

from ..._core import log
from . import captions, jobs, loop, media_io, store, swap

POSITIVE = ("own", "liked")
NEGATIVE = ("disliked", "bad_composition", "bad_image")
USES = {"normal": "liked and your own media", "avoid": "disliked results, to train an 'avoid' LoRA", "all": "everything in the set"}
ALGORITHMS = ("LoRA", "LoHa", "LoKr", "OFT")


class Stopped(Exception):
    """The person pressed Stop."""


def pick(doc: dict, use: str) -> list[dict]:
    """The items a run trains on. 'normal' leaves out what was disliked, 'avoid' keeps only that, 'all' keeps everything."""
    keep = {"normal": POSITIVE, "avoid": NEGATIVE, "all": POSITIVE + NEGATIVE}.get(use)
    if keep is None:
        raise ValueError(f"use is one of {sorted(USES)}, not {use!r}")
    return [i for i in doc["items"] if i["label"] in keep]


def roots() -> dict:
    import folder_paths
    return {"output": folder_paths.get_output_directory(), "temp": folder_paths.get_temp_directory(),
            "input": folder_paths.get_input_directory()}


def bin_path(media_id: str):
    from ..._core import media
    try:
        return media.path_for(media_id)
    except Exception:                                    # noqa: BLE001 -- a bin entry that is gone is a missing file, said by the caller
        return None


def plan_items(set_name: str, use: str):
    """(items with their files, problems). Nothing here touches torch."""
    doc = store.get(set_name)
    if doc is None:
        return [], [f"there is no training set called {set_name!r}"]
    problems, found = [], []
    chosen = pick(doc, use)
    if not chosen:
        problems.append(f"the set {set_name!r} has nothing to train on for 'use: {use}' ({USES[use]})")
    for item in chosen:
        path = store.resolve(item["src"], roots(), bin_path)
        if path is None:
            problems.append(f"{item['name'] or item['id']} is no longer there")
            continue
        found.append((item, path))
    return found, problems


def _fingerprint(obj) -> str:
    """Tells one set of VAE or text-encoder weights from another cheaply: class, size and a few values."""
    import hashlib
    if obj is None:
        return ""
    inner = getattr(obj, "first_stage_model", None) or getattr(obj, "cond_stage_model", None) or getattr(obj, "model", None) or obj
    h = hashlib.sha256(type(inner).__name__.encode())
    try:
        params = list(inner.parameters())
        h.update(str(sum(p.numel() for p in params)).encode())
        for p in (params[0], params[-1]) if params else ():
            h.update(p.detach().flatten()[:16].float().cpu().numpy().tobytes())
    except Exception:                                      # noqa: BLE001 -- an object with no parameters: its class has to do
        pass
    return h.hexdigest()[:16]


def _pair(result):
    return result[0].clone(), None if result[1] is None else result[1].clone()


def _cached(kind, parts, make, torch):
    """What an item encodes to, kept on disk for the next run: same file, same settings, same weights give the same answer. Returns (value, was_cached)."""
    import hashlib
    import json
    from ..._core import config
    folder = config.TRAINING_DIR / "cache"
    path = folder / f"{kind}-{hashlib.sha256(json.dumps(parts, default=str).encode()).hexdigest()[:24]}.pt"
    if path.is_file():
        try:
            return torch.load(path, weights_only=False), True
        except Exception:                                  # noqa: BLE001 -- a damaged or foreign-device file: encode again
            pass
    value = make()
    try:
        folder.mkdir(parents=True, exist_ok=True)
        tmp = path.with_suffix(".tmp")
        torch.save(value, tmp)
        tmp.replace(path)
    except Exception as exc:                               # noqa: BLE001 -- the cache is only a speed-up
        log.warning("FunUI LoRA training", f"could not keep {kind} in the cache: {exc}")
    return value, False


def _encode_text(clip, text: str):
    tokens = clip.tokenize(text)
    return clip.encode_from_tokens_scheduled(tokens)


def _encode_pixels(vae, pixels):
    """Pixels [N,H,W,3] -> latent tensor, in the VAE's own layout."""
    return vae.encode(pixels[:, :, :, :3])


def encode_item(item, path, size, cfg, vae, adapt, audio_vae):
    """One item's latents. For video models a clip stays a clip (cfg['frames'] frames); for picture models a clip gives one frame."""
    import torch
    width, height = size
    if adapt and adapt.get("video"):
        frames = adapt["frames"](cfg["frames"])
        if item["media_kind"] == "video":
            pixels, rate = media_io.clip_frames(path, frames, cfg["fps"] or adapt.get("fps", 24), cfg["start"], width, height)
        else:
            from PIL import Image, ImageOps
            with Image.open(path) as img:
                still = media_io.cover(ImageOps.exif_transpose(img).convert("RGB"), width, height)
            pixels, rate = media_io.to_tensor([still] * frames), float(adapt.get("fps", 24))
        video = _encode_pixels(vae, pixels)
        if video.ndim == 4:                               # a 2D VAE gave a stack of frames: [T,C,h,w] -> [1,C,T,h,w]
            video = video.permute(1, 0, 2, 3)[None]
        sound = None
        if audio_vae is not None and adapt.get("silence"):
            sound = _sound_latent(path, item["media_kind"], frames / rate, audio_vae, adapt)
        return video, sound
    img = media_io_picture(path, item["media_kind"], width, height)
    return _encode_pixels(vae, media_io.to_tensor([img])), None


def media_io_picture(path, kind, width, height):
    from . import frames
    return media_io.cover(frames.picture(path), width, height)


def _sound_latent(path, kind, seconds, audio_vae, adapt):
    import comfy.audio
    wave = media_io.read_audio(path, seconds) if kind == "video" else None
    if wave is None:
        return adapt["silence"](audio_vae, seconds)
    waveform, rate = wave
    target = int(getattr(audio_vae, "audio_sample_rate", 44100))
    if rate != target:
        waveform = comfy.audio.resample(waveform, rate, target)
    return audio_vae.encode(waveform.movedim(1, -1))


def run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = "") -> str:
    """Train and save. Returns the saved LoRA's file name inside the loras folder. Raises with a reason a person can act on.

    A failed run lets go of everything it held: ComfyUI keeps the error (and so every frame's locals, the model copy
    among them) alive after it is shown, and the next run then starts with that much of the card already taken.
    """
    import gc
    import traceback
    import torch
    import comfy.model_management as mm
    _free_the_card(mm, torch)                             # whatever a generation or an earlier run left on the card goes first, cached training data or not
    _card("at the start of the run", mm, torch)
    try:
        result = _run(model, clip, vae, audio_vae, set_name, cfg, prompt_id)
        model = clip = vae = audio_vae = None             # a paused or finished run lets the card go: the pipeline's preview or the next run needs it
        _free_the_card(mm, torch)
        return result
    except BaseException as exc:
        traceback.clear_frames(exc.__traceback__)
        model = clip = vae = audio_vae = None
        _free_the_card(mm, torch)
        raise


def _run(model, clip, vae, audio_vae, set_name: str, cfg: dict, prompt_id: str = "") -> str:
    import torch
    import comfy.model_management as mm
    import comfy_extras.nodes_train as nt
    import folder_paths
    import safetensors.torch

    from ..._core import registry
    resume, failed = None, None
    if cfg.get("resume"):
        try:
            resume = loop.load_checkpoint(cfg["resume"], torch)
        except ValueError as exc:
            failed = exc
    run_id = jobs.begin(set_name, cfg["name"], cfg["steps"], cfg, prompt_id, {"run": cfg["resume"], "loss": resume["loss"], "times": resume.get("times"), "seconds": resume.get("seconds")} if resume else None)        # first, so a refusal below is shown in the window and not only in ComfyUI's log
    try:
        if failed:
            raise failed
        if model is None:
            model = load_model(cfg.get("model_name"), cfg.get("weight_dtype", "default"), cfg.get("compute_dtype", "default"))
        adapt = registry.current().ask("training_adapt", model)
        items, problems = plan_items(set_name, cfg["use"])
        if problems and not items:
            raise ValueError("; ".join(problems))
        if problems:
            log.warning("FunUI LoRA training", f"left out of the run: {'; '.join(problems[:5])}{' ...' if len(problems) > 5 else ''}")
        if getattr(model, "cached_patcher_init", None) is None:
            raise ValueError("this model was loaded in a way the trainer cannot copy (MXFP8 or int8 conversion in the loader, or a file the loader cannot read again); "
                             "untick those in the loader and wire it again, or leave the model unwired and pick the file in the trainer")
        batch = cfg["batch_size"]
        if adapt and batch > adapt.get("batch_limit", batch):
            log.warning("FunUI LoRA training", f"{adapt['name']} trains one sample at a time; batch size {batch} is lowered to {adapt['batch_limit']}")
            batch = adapt["batch_limit"]
        jobs.say(f"Reading {len(items)} items")
        step = (adapt or {}).get("step") or 64
        sizes = []
        for item, path in items:
            try:
                sizes.append(media_io.size_of(path, item["media_kind"]))
            except Exception as exc:                      # noqa: BLE001
                raise ValueError(f"{item['name'] or item['id']} cannot be read: {exc}") from exc
        single = bool(adapt) or not cfg["aspect_buckets"]
        buckets = media_io.buckets_for([w / h for w, h in sizes], cfg["resolution"], step, single)
        texts = [captions.with_parts(item["caption"], cfg["trigger"], cfg.get("postfix", "")) for item, _ in items]

        if adapt and adapt.get("video"):
            cfg = _fit_frames(cfg, adapt, items)
        videos, sounds, conds = [], [], []
        hits = 0
        with torch.inference_mode():
            vae_id, audio_id, clip_id = _fingerprint(vae), _fingerprint(audio_vae), _fingerprint(clip)
            for number, ((item, path), size, text) in enumerate(zip(items, buckets, texts), 1):
                if jobs.stop_requested():
                    raise Stopped()
                jobs.say(f"Preparing {number} of {len(items)}")
                shape = [str(Path(path).resolve()), Path(path).stat().st_mtime_ns, Path(path).stat().st_size, item["media_kind"], list(size), cfg["frames"], cfg["fps"], cfg["start"],
                         (adapt or {}).get("name"), vae_id, audio_id if audio_vae is not None else None]
                (latent, sound), hit = _cached("latent", shape, lambda: _pair(encode_item(item, path, size, cfg, vae, adapt, audio_vae)), torch)
                cond, hit_text = _cached("text", [text, clip_id], lambda: _encode_text(clip, text), torch)
                hits += hit and hit_text
                videos.append(latent.clone())
                sounds.append(None if sound is None else sound.clone())
                conds.append(cond)
        log.info("FunUI LoRA training", f"prepared {len(items)} items, {hits} entirely from the cache")
        videos, sounds, conds = _normal(videos), _normal(sounds), _normal(conds)         # made in inference mode: copies a backward pass can use
        if adapt and adapt.get("pack") and audio_vae is not None and all(s is not None for s in sounds):
            latents = adapt["pack"](videos, sounds)
            payload = [latents]
        elif adapt and adapt.get("pack"):
            raise ValueError("this model trains picture and sound together and needs its audio VAE: add the audio VAE loader to the pipeline")
        elif len({tuple(v.shape) for v in videos}) == 1:
            payload = [{"samples": torch.cat(videos, 0)}]
        else:
            payload = [{"samples": v} for v in videos]

        _sizes(payload, conds)
        _card("before freeing", mm, torch)
        del vae, clip, audio_vae                          # their work is done; the trainer copies the model next and needs the room
        _free_the_card(mm, torch)
        _card("after freeing", mm, torch)
        _train_on_a_plain_copy(model)
        jobs.say("Training", "training")
        shim = adapt["shims"]() if adapt and adapt.get("shims") else _nothing()
        watch = _watch(nt, mm, torch, _swapper(adapt, cfg, mm, torch))
        if adapt and adapt.get("lora_targets"):
            def keep(state, step):                        # a copy of the LoRA so far, in the loras folder, never over an earlier file
                name = _save(state, f"{cfg['name']}_step{step}", folder_paths, safetensors.torch, set_name, cfg)
                log.info("FunUI LoRA training", f"saved {name} at step {step}")
            lora, paused = loop.train(model, payload, conds, cfg, adapt, nt, mm, torch, watch, shim, resume, run_id, keep)       # our own loop, ComfyUI carries the batches
            if paused:
                preview = _save(lora, f"{cfg['name']}_preview", folder_paths, safetensors.torch, set_name, cfg, keep=True)
                jobs.pause(paused.reason, paused.step, preview)
                log.info("FunUI LoRA training", f"paused at step {paused.step} ({paused.reason}); the LoRA so far is {preview}")
                return ""
        else:
            losses = _loss_hook(nt, cfg["steps"] * cfg["grad_accumulation"])
            with shim, losses, watch:
                out = nt.TrainLoraNode.execute(
                    model=[model], latents=payload, positive=conds, batch_size=[batch], steps=[cfg["steps"]],
                    grad_accumulation_steps=[cfg["grad_accumulation"]], learning_rate=[cfg["learning_rate"]], rank=[cfg["rank"]],
                    optimizer=[cfg["optimizer"]], loss_function=[cfg["loss"]], seed=[cfg["seed"]], training_dtype=[cfg["training_dtype"]],
                    lora_dtype=[cfg["lora_dtype"]], quantized_backward=[False], algorithm=[cfg["algorithm"]],
                    gradient_checkpointing=[cfg["gradient_checkpointing"]], checkpoint_depth=[(adapt or {}).get("checkpoint_depth", 1)], offloading=[cfg["offloading"]],
                    existing_lora=["[None]"], bucket_mode=[False], bypass_mode=[cfg["bypass"]])
            lora = (out.result if hasattr(out, "result") else out)[0]
        file = _save(lora, cfg["name"], folder_paths, safetensors.torch, set_name, cfg)
        discard(run_id, cfg["name"])
        jobs.finish("done", lora=file)
        return file
    except Stopped:
        jobs.finish("stopped")
        raise
    except mm.InterruptProcessingException:
        jobs.finish("stopped")
        raise
    except Exception as exc:                              # noqa: BLE001 -- recorded for the window, then raised for ComfyUI to show
        jobs.finish("error", error=f"{type(exc).__name__}: {exc}")
        raise


class _nothing:
    def __enter__(self):
        return self

    def __exit__(self, *exc):
        return False


_WEIGHT = {"fp8_e4m3fn": "float8_e4m3fn", "fp8_e5m2": "float8_e5m2", "fp16": "float16", "bf16": "bfloat16", "fp32": "float32"}


def load_model(name, weight_dtype="default", compute_dtype="default"):
    """The model file, read once and never moved to the card here. The trainer's own copy shares it, so one model is in memory."""
    import torch
    import comfy.sd
    import folder_paths
    if not name:
        raise ValueError("no model to train on: the pipeline's model loader names no file")
    if str(name).lower().endswith(".gguf"):
        raise ValueError("a GGUF model cannot be trained: pick a .safetensors file in the model loader")
    path = folder_paths.get_full_path_or_raise("diffusion_models", name)
    options = {}
    if weight_dtype in _WEIGHT:
        options["dtype"] = getattr(torch, _WEIGHT[weight_dtype])
    with torch.inference_mode(False):                    # ComfyUI runs nodes in inference mode; weights made there cannot take part in a backward pass
        model = comfy.sd.load_diffusion_model(path, model_options=options, disable_dynamic=True)
    if compute_dtype in _WEIGHT and hasattr(model, "set_model_compute_dtype"):
        model.set_model_compute_dtype(getattr(torch, _WEIGHT[compute_dtype]))
    model.cached_patcher_init = (lambda *args, **kwargs: model, ())      # the trainer's "fresh copy" is this model: no second read of the file
    return model


def _normal(value):
    """Tensors made in inference mode, copied so autograd may save them; lists, tuples and dicts are walked."""
    import torch
    if isinstance(value, torch.Tensor):
        with torch.inference_mode(False):
            return value.clone() if value.is_inference() else value
    if isinstance(value, list):
        return [_normal(v) for v in value]
    if isinstance(value, tuple):
        return tuple(_normal(v) for v in value)
    if isinstance(value, dict):
        return {k: _normal(v) for k, v in value.items()}
    return value


def _fit_frames(cfg, adapt, items):
    """Every clip is read for the same number of frames, so the shortest clip sets the limit: fewer than asked is used, and said."""
    rate = cfg["fps"] or adapt.get("fps", 24)
    shortest = min((media_io.frames_at(path, rate) for item, path in items if item["media_kind"] == "video"), default=None)
    wanted = adapt["frames"](cfg["frames"])
    if shortest is None or shortest >= wanted:
        return cfg
    fit = adapt["fit"](shortest) if adapt.get("fit") else shortest
    if fit < 1:
        raise ValueError(f"the shortest clip has only {shortest} frames at {rate:g} fps: too short to train on")
    log.warning("FunUI LoRA training", f"the shortest clip gives {shortest} frames at {rate:g} fps, so every clip is read for {fit} frames and not {wanted}")
    return {**cfg, "frames": fit}


def _sizes(payload, conds):
    """One Log line: the shape of the first item's latent and of its text, to tell a run that is too big from one that is not."""
    def shape(x):
        tensors = getattr(x, "tensors", None)
        if tensors is not None:
            return " + ".join(str(list(t.shape)) for t in tensors)
        return str(list(x.shape)) if hasattr(x, "shape") else type(x).__name__
    try:
        first = payload[0]["samples"] if isinstance(payload[0], dict) else payload[0]
        text = conds[0][0] if isinstance(conds[0], (list, tuple)) else conds[0]
        log.info("FunUI LoRA training", f"{len(conds)} text(s); first latent {shape(first)}, first text {shape(text)}")
    except Exception as exc:                              # noqa: BLE001 -- a report must never stop a run
        log.info("FunUI LoRA training", f"sizes: not readable ({type(exc).__name__})")


def _card(when, mm, torch):
    """One line in the Log: what the card holds and which models ComfyUI still has loaded."""
    try:
        device = mm.get_torch_device()
        gib = 1024 ** 3
        loaded = ", ".join(f"{type(getattr(m.model, 'model', m.model)).__name__} {m.model_loaded_memory() / gib:.1f} GiB" for m in mm.current_loaded_models) or "none"
        line = f"card {when}: tensors {torch.cuda.memory_allocated(device) / gib:.1f} GiB, free {mm.get_free_memory(device) / gib:.1f} GiB, loaded models: {loaded}"
        jobs.trail(line)
    except Exception as exc:                              # noqa: BLE001 -- a report must never stop a run
        log.info("FunUI LoRA training", f"card {when}: not readable ({type(exc).__name__})")


class _watch:
    """While open, a Log line after each stage of the trainer's setup, so an out-of-memory says which stage filled the card."""
    STAGES = ("_prepare_latents_and_count", "_setup_lora_adapters", "_setup_lora_adapters_bypass", "_create_optimizer")

    def __init__(self, nt, mm, torch, after_load=None):
        self.nt, self.mm, self.torch, self.put, self.after_load = nt, mm, torch, [], after_load
        self.state = {"backward": False, "recomputes": 0}

    def __enter__(self):
        def wrap(owner, name, label):
            original = getattr(owner, name, None)
            if original is None:
                return

            def watched(*args, **kwargs):
                if label == "loading the model":
                    jobs.trail("loading the model onto the card now")
                out = original(*args, **kwargs)
                if label == "loading the model" and self.after_load:
                    self.after_load(args[0] if args else kwargs.get("models"))
                _card(f"after {label}", self.mm, self.torch)
                return out
            setattr(owner, name, watched)
            self.put.append((owner, name, original))
        for name in self.STAGES:
            wrap(self.nt, name, name.strip("_").replace("_", " "))
        wrap(self.mm, "load_models_gpu", "loading the model")
        self._backward()
        self._recomputes()
        return self

    def _recomputes(self):
        """A Log line for the first block recomputes of the backward pass: how the card fills block by block."""
        nt, mm, torch, state = self.nt, self.mm, self.torch, self.state
        original, counter = nt.patch, [0]

        def patch(m, offloading=False):
            idx, forward = counter[0], getattr(m, "forward", None)
            counter[0] += 1
            if forward is not None:
                def logged(*args, **kwargs):
                    if state["backward"]:
                        state["recomputes"] += 1
                        n = state["recomputes"]
                        if n <= 8 or n % 10 == 0:
                            _card(f"backward recompute {n} (module {idx}, peak {torch.cuda.max_memory_allocated() / 1024 ** 3:.1f} GiB)", mm, torch)
                    return forward(*args, **kwargs)
                m.forward = logged
            return original(m, offloading)
        nt.patch = patch
        self.put.append((nt, "patch", original))

    def _backward(self):
        """A Log line just before the first backward passes: what the forward pass left on the card (the part that grows with clip length)."""
        autograd, torch, seen = self.torch.autograd, self.torch, [0]
        original = autograd.backward

        def watched(*args, **kwargs):
            self.state["backward"] = True
            if seen[0] < 2:
                seen[0] += 1
                _card(f"before backward pass {seen[0]} (peak so far {torch.cuda.max_memory_allocated() / 1024 ** 3:.1f} GiB)", self.mm, torch)
            return original(*args, **kwargs)
        autograd.backward = watched
        self.put.append((autograd, "backward", original))

    def __exit__(self, *exc):
        for owner, name, original in reversed(self.put):
            setattr(owner, name, original)
        return False


def _swapper(adapt, cfg, mm, torch):
    """After the model is on the card: the weights of the chosen blocks go to system memory (block swap), or None when not asked for."""
    count = int(cfg.get("swap_blocks") or 0)
    if count <= 0:
        return None
    if not (adapt and adapt.get("blocks")):
        log.warning("FunUI LoRA training", "block swap is asked for, but this model has no block list to swap: the whole model stays on the card")
        return None

    def after(patchers):
        for patcher in patchers or []:
            blocks = adapt["blocks"](patcher.model.diffusion_model)
            done = swap.move(blocks, count, torch)
            if done["blocks"]:
                log.info("FunUI LoRA training", f"block swap: {done['blocks']} of {len(blocks)} blocks ({done['bytes'] / 1024 ** 3:.1f} GiB) now wait in system memory")
                _card("after block swap", mm, torch)
            elif not swap.still_away(blocks, count):
                log.warning("FunUI LoRA training", "block swap: ComfyUI moved the weights back to the card")
    return after


def _free_the_card(mm, torch):
    import gc
    gc.collect()
    mm.unload_all_models()
    gc.collect()
    mm.soft_empty_cache()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _train_on_a_plain_copy(model):
    """The trainer copies the model with clone(force_deepcopy=True). On a dynamic-VRAM model that copy is dynamic too, and the
    trainer then loads it in full beside the original: two models on the card. Ask for a plain copy, loaded once."""
    is_dynamic = getattr(model, "is_dynamic", None)
    if not (callable(is_dynamic) and is_dynamic()):
        return
    clone = model.clone

    def plain(*args, **kwargs):
        return clone(*args, **{**kwargs, "disable_dynamic": True})
    model.clone = plain


class _loss_hook:
    """While open, every step's loss goes to the run record, and Stop (or ComfyUI's interrupt) ends the loop at the next step."""

    def __init__(self, nt, total):
        self.nt, self.total = nt, total

    def __enter__(self):
        import comfy.model_management as mm
        nt = self.nt
        self.original = nt.TrainSampler

        def sampler(*args, **kwargs):
            inner = kwargs.get("loss_callback")

            def callback(loss):
                if inner:
                    inner(loss)
                jobs.step(float(loss))
                mm.throw_exception_if_processing_interrupted()
                if jobs.stop_requested():
                    raise Stopped()
            kwargs["loss_callback"] = callback
            return self.original(*args, **kwargs)

        nt.TrainSampler = sampler
        return self

    def __exit__(self, *exc):
        self.nt.TrainSampler = self.original
        return False


def _preview_name(name) -> str:
    stem = "".join(c for c in f"{name}_preview" if c.isalnum() or c in "-_ .").strip() or "funui_lora_preview"
    return f"{stem}.safetensors"


def discard(run, name) -> None:
    """A run that is over (finished, stopped) keeps neither its saved state nor its preview file."""
    import folder_paths
    loop.forget_checkpoint(run)
    try:
        (Path(folder_paths.get_folder_paths("loras")[0]) / _preview_name(name.removesuffix("_preview"))).unlink()
    except OSError:
        pass


def _save(lora, name, folder_paths, safetensors_torch, set_name, cfg, keep=False) -> str:
    """Into the first loras folder, as <name>.safetensors, or <name>_2 and so on rather than over an earlier one. `keep`: replace the file of that name (the preview)."""
    folder = Path(folder_paths.get_folder_paths("loras")[0])
    folder.mkdir(parents=True, exist_ok=True)
    stem = "".join(c for c in name if c.isalnum() or c in "-_ .") .strip() or "funui_lora"
    n, path = 1, folder / f"{stem}.safetensors"
    while path.exists() and not keep:
        n += 1
        path = folder / f"{stem}_{n}.safetensors"
    meta = {"software": "FunUI", "training_set": set_name, "algorithm": cfg["algorithm"], "rank": str(cfg["rank"]),
            "steps": str(cfg["steps"]), "trigger": cfg["trigger"]}
    safetensors_torch.save_file(lora, str(path), metadata=meta)
    return path.name