473 lines
21 KiB
Python
473 lines
21 KiB
Python
"""A whole training run through the module's own entry point, on a tiny real H3 on CPU.
|
|
|
|
The model is ComfyUI's own MiniMaxH3 class, a few hundred thousand parameters wide, with a stand-in text
|
|
encoder and VAEs that return latents of the right shape. What this proves is OUR side: the set, the
|
|
captions, the latents packed the way H3 wants, the two H3 repairs, the loss stream, and the saved file.
|
|
It does not prove anything about real weights, a real VAE or real hardware.
|
|
"""
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Needs ComfyUI's source for the trainer."""
|
|
|
|
|
|
def tiny_h3_patcher():
|
|
import torch
|
|
import comfy.supported_models as sm
|
|
from comfy.model_patcher import ModelPatcher
|
|
unet = dict(image_model="minimax_h3", hidden_size=64, num_layers=2, token_refiner_num_layers=1, num_attention_heads=2,
|
|
attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, time_embed_hidden_size=64,
|
|
time_embed_dim=32, rope_inv_freq_len=4)
|
|
cfg = sm.MiniMaxH3(dict(unet))
|
|
cfg.set_inference_dtype(torch.float32, None)
|
|
|
|
def make(disable_dynamic=False):
|
|
base = cfg.get_model({}, "", device=torch.device("cpu"))
|
|
torch.manual_seed(0)
|
|
for p in base.parameters():
|
|
torch.nn.init.normal_(p, std=0.05)
|
|
base.diffusion_model.rope.inv_freq.copy_(torch.rand(4))
|
|
return ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
|
|
patcher = make()
|
|
patcher.cached_patcher_init = (make, ())
|
|
return patcher
|
|
|
|
|
|
class FakeClip:
|
|
def tokenize(self, text):
|
|
return text
|
|
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
import torch
|
|
return [[torch.randn(1, 6, 48), {}]]
|
|
|
|
|
|
class FakeVae:
|
|
def encode(self, pixels):
|
|
import torch
|
|
from comfy_extras.nodes_minimax_h3 import video_latent_t
|
|
frames, h, w = pixels.shape[0], pixels.shape[1], pixels.shape[2]
|
|
return torch.randn(1, 24, video_latent_t(frames), h // 16, w // 16)
|
|
|
|
|
|
class FakeAudioVae:
|
|
audio_sample_rate = 44100
|
|
|
|
def encode(self, wave):
|
|
import torch
|
|
seconds = wave.shape[1] / 44100
|
|
return torch.randn(1, 32, 2, round(seconds * 40))
|
|
|
|
|
|
@pytest.fixture
|
|
def run_env(tmp_path, monkeypatch):
|
|
import folder_paths
|
|
from PIL import Image
|
|
from modules.training.lora import jobs, store, train
|
|
from core import config
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path / "training")
|
|
monkeypatch.setattr(train, "roots", lambda: {"input": str(tmp_path)})
|
|
loras = tmp_path / "loras"
|
|
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [str(loras)])
|
|
from modules.models.minimax_h3 import training as h3
|
|
monkeypatch.setattr(train, "registry", None, raising=False)
|
|
import modules.training.lora.train as t
|
|
class Reg:
|
|
def ask(self, capability, *a, **k):
|
|
return h3.adapt(*a, **k) if capability == "training_adapt" else None
|
|
from modules import _core
|
|
monkeypatch.setattr(_core.registry, "current", lambda: Reg())
|
|
names = []
|
|
for i in range(2):
|
|
Image.new("RGB", (80, 80), (40 * i, 90, 160)).save(tmp_path / f"p{i}.png")
|
|
names.append(f"p{i}.png")
|
|
store.create("tiny")
|
|
store.add("tiny", [{"src": {"kind": "output", "type": "input", "subfolder": "", "filename": n}, "media_kind": "image",
|
|
"label": "own", "caption": "a blue square, flat"} for n in names])
|
|
return tmp_path, jobs
|
|
|
|
|
|
def settings(**over):
|
|
base = {"name": "tiny_lora", "use": "normal", "steps": 6, "learning_rate": 3e-3, "rank": 4, "algorithm": "LoRA", "optimizer": "AdamW",
|
|
"loss": "MSE", "batch_size": 4, "grad_accumulation": 1, "resolution": 64, "aspect_buckets": True, "frames": 5, "fps": 0.0,
|
|
"start": "start", "trigger": "zq", "seed": 0, "training_dtype": "bf16", "lora_dtype": "fp32",
|
|
"gradient_checkpointing": False, "offloading": False, "bypass": False}
|
|
base.update(over)
|
|
return base
|
|
|
|
|
|
def test_a_tiny_h3_trains_and_saves_a_lora(run_env):
|
|
import safetensors.torch
|
|
from modules.training.lora import train
|
|
tmp_path, jobs = run_env
|
|
file = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings())
|
|
saved = safetensors.torch.load_file(str(tmp_path / "loras" / file))
|
|
assert file == "tiny_lora.safetensors" and len(saved) > 20
|
|
snap = jobs.snapshot()
|
|
assert snap["state"] == "done" and snap["step"] == 6 and len(snap["loss"]) == 6
|
|
assert all(x == x for x in snap["loss"]), "a NaN loss"
|
|
|
|
|
|
def test_a_second_run_never_overwrites_the_first(run_env):
|
|
from modules.training.lora import train
|
|
tmp_path, _ = run_env
|
|
first = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
|
|
second = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
|
|
assert first != second and (tmp_path / "loras" / first).is_file()
|
|
|
|
|
|
def test_h3_without_its_audio_vae_says_what_to_add(run_env):
|
|
from modules.training.lora import train
|
|
with pytest.raises(ValueError, match="audio VAE"):
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), None, "tiny", settings())
|
|
|
|
|
|
def test_a_model_the_trainer_cannot_copy_is_refused_before_training_with_the_reason(run_env):
|
|
"""A loader that converted the weights (MXFP8, int8) leaves no way to re-read the file: ComfyUI's own error must not be what is shown."""
|
|
from modules.training.lora import train
|
|
patcher = tiny_h3_patcher()
|
|
patcher.cached_patcher_init = None
|
|
with pytest.raises(ValueError, match="MXFP8"):
|
|
train.run(patcher, FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
|
|
|
|
|
|
def test_stop_ends_the_run_and_records_it(run_env):
|
|
from modules.training.lora import train
|
|
_, jobs = run_env
|
|
real = jobs.step
|
|
|
|
def stop_after_two(loss, seconds=None):
|
|
real(loss, seconds)
|
|
if jobs.snapshot()["step"] >= 2:
|
|
jobs.request_stop()
|
|
jobs.step = stop_after_two
|
|
try:
|
|
with pytest.raises(train.Stopped):
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=50))
|
|
finally:
|
|
jobs.step = real
|
|
snap = jobs.snapshot()
|
|
assert snap["state"] == "stopped" and snap["step"] < 50
|
|
|
|
|
|
def test_a_set_with_nothing_for_the_chosen_use_is_refused(run_env):
|
|
from modules.training.lora import train
|
|
with pytest.raises(ValueError, match="nothing to train on"):
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(use="avoid"))
|
|
|
|
|
|
# --- a picture model: no adapter, plain latent tensors ------------------------------------------------
|
|
|
|
def tiny_sd_patcher():
|
|
import torch
|
|
import comfy.supported_models as sm
|
|
from comfy.model_patcher import ModelPatcher
|
|
unet = {"context_dim": 32, "model_channels": 32, "use_linear_in_transformer": False, "adm_in_channels": None, "use_temporal_attention": False,
|
|
"in_channels": 4, "out_channels": 4, "image_size": 32, "num_res_blocks": [1, 1], "transformer_depth": [1, 1], "channel_mult": [1, 2],
|
|
"transformer_depth_output": [1, 1, 1, 1], "transformer_depth_middle": 1, "use_spatial_transformer": True, "num_head_channels": 8}
|
|
cfg = sm.SD15(dict(unet))
|
|
cfg.set_inference_dtype(torch.float32, None)
|
|
|
|
def make(disable_dynamic=False):
|
|
base = cfg.get_model({}, "", device=torch.device("cpu"))
|
|
torch.manual_seed(0)
|
|
for p in base.parameters():
|
|
torch.nn.init.normal_(p, std=0.05)
|
|
return ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu"))
|
|
patcher = make()
|
|
patcher.cached_patcher_init = (make, ())
|
|
return patcher
|
|
|
|
|
|
class FakeClip32(FakeClip):
|
|
def encode_from_tokens_scheduled(self, tokens):
|
|
import torch
|
|
return [[torch.randn(1, 6, 32), {}]]
|
|
|
|
|
|
class FakeVae4:
|
|
def encode(self, pixels):
|
|
import torch
|
|
return torch.randn(pixels.shape[0], 4, pixels.shape[1] // 8, pixels.shape[2] // 8)
|
|
|
|
|
|
@pytest.fixture
|
|
def plain_env(run_env, monkeypatch):
|
|
from modules import _core
|
|
class NoAdapter:
|
|
def ask(self, capability, *a, **k):
|
|
return None
|
|
monkeypatch.setattr(_core.registry, "current", lambda: NoAdapter())
|
|
return run_env
|
|
|
|
|
|
def test_a_picture_model_trains_from_a_set_of_one_shape(plain_env):
|
|
from modules.training.lora import train
|
|
_, jobs = plain_env
|
|
file = train.run(tiny_sd_patcher(), FakeClip32(), FakeVae4(), None, "tiny", settings(aspect_buckets=False, batch_size=2, steps=4))
|
|
snap = jobs.snapshot()
|
|
assert file.endswith(".safetensors") and snap["state"] == "done" and len(snap["loss"]) == 4
|
|
assert all(x == x for x in snap["loss"])
|
|
|
|
|
|
def test_a_picture_model_trains_from_pictures_of_different_shapes(plain_env):
|
|
from PIL import Image
|
|
from modules.training.lora import store, train
|
|
tmp_path, jobs = plain_env
|
|
Image.new("RGB", (160, 64), (200, 30, 30)).save(tmp_path / "wide.png")
|
|
store.add("tiny", [{"src": {"kind": "output", "type": "input", "subfolder": "", "filename": "wide.png"}, "media_kind": "image", "label": "own", "caption": "wide"}])
|
|
train.run(tiny_sd_patcher(), FakeClip32(), FakeVae4(), None, "tiny", settings(aspect_buckets=True, batch_size=1, steps=3))
|
|
assert jobs.snapshot()["state"] == "done"
|
|
|
|
|
|
def test_a_dynamic_model_is_copied_as_a_plain_one_for_the_trainer():
|
|
from modules.training.lora import train
|
|
|
|
class Dynamic:
|
|
calls = []
|
|
|
|
def is_dynamic(self):
|
|
return True
|
|
|
|
def clone(self, disable_dynamic=False, force_deepcopy=False):
|
|
self.calls.append((disable_dynamic, force_deepcopy))
|
|
m = Dynamic()
|
|
train._train_on_a_plain_copy(m)
|
|
m.clone(force_deepcopy=True)
|
|
assert m.calls == [(True, True)]
|
|
|
|
class Plain(Dynamic):
|
|
def is_dynamic(self):
|
|
return False
|
|
p = Plain()
|
|
train._train_on_a_plain_copy(p)
|
|
assert "clone" not in vars(p)
|
|
|
|
|
|
def test_the_trainer_loads_the_model_file_once_and_its_copy_is_that_model(monkeypatch):
|
|
import comfy.sd
|
|
import folder_paths
|
|
import torch
|
|
from modules.training.lora import train
|
|
seen = {}
|
|
|
|
class Patcher:
|
|
def __init__(self):
|
|
self.weight = torch.zeros(2) # made where the loader runs: inference mode must not reach it
|
|
|
|
def set_model_compute_dtype(self, dtype):
|
|
seen["compute"] = dtype
|
|
|
|
def load(path, model_options=None, disable_dynamic=False):
|
|
seen.update(path=path, options=model_options, dynamic=disable_dynamic, loads=seen.get("loads", 0) + 1)
|
|
return Patcher()
|
|
monkeypatch.setattr(comfy.sd, "load_diffusion_model", load)
|
|
monkeypatch.setattr(folder_paths, "get_full_path_or_raise", lambda kind, name: f"/models/{name}")
|
|
with torch.inference_mode():
|
|
model = train.load_model("m.safetensors", "bf16", "fp32")
|
|
assert not model.weight.is_inference()
|
|
assert seen["path"] == "/models/m.safetensors" and seen["options"] == {"dtype": torch.bfloat16} and seen["dynamic"] and seen["compute"] == torch.float32
|
|
factory, args = model.cached_patcher_init
|
|
assert factory(*args, disable_dynamic=True) is model and seen["loads"] == 1 # the trainer's copy reads nothing again
|
|
with pytest.raises(ValueError, match="GGUF"):
|
|
train.load_model("m.gguf")
|
|
with pytest.raises(ValueError, match="no model"):
|
|
train.load_model("")
|
|
|
|
|
|
def test_h3_is_checkpointed_one_block_at_a_time_not_as_a_whole(run_env):
|
|
import comfy_extras.nodes_train as nt
|
|
from modules.models.minimax_h3 import training
|
|
patcher = tiny_h3_patcher()
|
|
depth = training.adapt(patcher)["checkpoint_depth"]
|
|
net = patcher.model.diffusion_model
|
|
pieces = nt.find_modules_at_depth(net, depth=depth)
|
|
assert depth == 2 and len(pieces) > len(net.blocks) and all(b in pieces for b in net.blocks)
|
|
assert net not in pieces
|
|
|
|
|
|
def test_what_was_made_in_inference_mode_is_copied_for_the_backward_pass():
|
|
import torch
|
|
from modules.training.lora import train
|
|
with torch.inference_mode():
|
|
made = [(torch.ones(2), {"k": torch.zeros(1), "n": 3})], [torch.ones(1)]
|
|
assert made[1][0].is_inference()
|
|
got = train._normal(made)
|
|
assert not got[1][0].is_inference() and not got[0][0][0].is_inference() and not got[0][0][1]["k"].is_inference() and got[0][0][1]["n"] == 3
|
|
x = got[1][0].requires_grad_(False)
|
|
(x * torch.ones(1, requires_grad=True)).sum().backward() # a saved-for-backward use no longer fails
|
|
|
|
|
|
def test_a_tiny_h3_trains_when_run_the_way_comfyui_runs_a_node(run_env):
|
|
"""ComfyUI runs every node in inference mode, and the models, text and latents here are made inside it."""
|
|
import torch
|
|
from modules.training.lora import train
|
|
_, jobs = run_env
|
|
with torch.inference_mode():
|
|
patcher = tiny_h3_patcher()
|
|
train.run(patcher, FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
|
|
assert jobs.snapshot()["state"] == "done"
|
|
|
|
|
|
def test_block_swap_moves_the_last_blocks_to_system_memory_and_leaves_trainable_weights():
|
|
import torch
|
|
from modules.training.lora import swap
|
|
|
|
class Layer(torch.nn.Linear):
|
|
comfy_cast_weights = False
|
|
|
|
def block():
|
|
b = torch.nn.Sequential(Layer(4, 4), Layer(4, 4))
|
|
for p in b[1].parameters():
|
|
p.requires_grad_(True) # a trainable weight is never moved
|
|
return b
|
|
|
|
blocks = [block() for _ in range(4)]
|
|
done = swap.move(blocks, 2, torch)
|
|
assert done["blocks"] == 2
|
|
assert all(l.comfy_cast_weights for l in blocks[3].modules() if isinstance(l, Layer)), "the swapped layers fetch their weights"
|
|
assert not any(l.comfy_cast_weights for l in blocks[0].modules() if isinstance(l, Layer))
|
|
assert swap.move(blocks, 2, torch)["blocks"] == 0, "asking twice moves nothing again"
|
|
assert swap.move(blocks, 0, torch)["blocks"] == 0
|
|
|
|
|
|
def test_a_tiny_h3_trains_with_block_swap_on(run_env):
|
|
from modules.training.lora import train
|
|
tmp_path, jobs = run_env
|
|
file = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2, gradient_checkpointing=True, swap_blocks=1))
|
|
snap = jobs.snapshot()
|
|
assert file.endswith(".safetensors") and snap["state"] == "done" and all(x == x for x in snap["loss"])
|
|
|
|
|
|
def test_block_modulation_is_out_of_place_while_training_and_put_back_after():
|
|
import torch
|
|
from comfy.ldm.minimax import model as m
|
|
from modules.models.minimax_h3 import training
|
|
before = (m._mod_scale_shift, m._mod_gate)
|
|
segs = [(0, 3, 0), (3, 5, 1)]
|
|
shift, scale, gate = (torch.randn(2, 4) for _ in range(3))
|
|
x, other = torch.randn(5, 4), torch.randn(5, 4)
|
|
want_h = m._mod_scale_shift(x.clone(), shift, scale, segs)
|
|
want_x = m._mod_gate(x.clone(), gate, other, segs)
|
|
with training.shims():
|
|
assert m._mod_scale_shift is not before[0]
|
|
x0 = x.clone()
|
|
assert torch.allclose(m._mod_scale_shift(x0, shift, scale, segs), want_h)
|
|
assert torch.equal(x0, x) # the input is left alone
|
|
assert torch.allclose(m._mod_gate(x0, gate, other, segs), want_x)
|
|
assert torch.equal(x0, x)
|
|
assert (m._mod_scale_shift, m._mod_gate) == before
|
|
|
|
|
|
def test_a_run_pauses_for_a_preview_and_carries_on_to_the_same_end(run_env):
|
|
import safetensors.torch
|
|
from modules.training.lora import train
|
|
tmp_path, jobs = run_env
|
|
first = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2))
|
|
snap = jobs.snapshot()
|
|
assert first == "" and snap["state"] == "paused"
|
|
assert snap["paused"]["reason"] == "preview" and snap["paused"]["step"] == 2 and snap["step"] == 2
|
|
preview = tmp_path / "loras" / snap["paused"]["preview"]
|
|
assert preview.name == "tiny_lora_preview.safetensors"
|
|
assert any(k.endswith(".lora_up.weight") for k in safetensors.torch.load_file(str(preview)))
|
|
run = snap["run"]
|
|
# the person looks at the preview, and training carries on from where it stopped
|
|
second = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
|
|
snap = jobs.snapshot()
|
|
assert second == "" and snap["run"] == run and snap["paused"]["step"] == 4 and len(snap["loss"]) == 4
|
|
done = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
|
|
snap = jobs.snapshot()
|
|
assert done == "tiny_lora.safetensors" and snap["state"] == "done" and len(snap["loss"]) == 6
|
|
times = snap["times"]
|
|
assert len(times) == 6 and [t is None for t in times] == [True, False, True, False, True, False], "each restart's first step also loads the model, so it is not timed"
|
|
assert snap["step_seconds"] and snap["elapsed"] >= sum(t for t in times if t), "the time of every part of the run is counted"
|
|
import json
|
|
from modules.training.lora import jobs as j
|
|
saved = json.loads((j._dir() / f"{run}.json").read_text())
|
|
assert saved["seconds"] >= snap["elapsed"] - 1 and saved["step_seconds"] == snap["step_seconds"] and len(saved["times"]) == 6
|
|
assert not preview.exists(), "the preview file stays after the run is over"
|
|
assert not (tmp_path / "training" / "checkpoints" / f"{run}.pt").exists()
|
|
|
|
|
|
def test_a_failed_run_is_restored_from_its_saved_state_in_comfyuis_inference_mode(run_env, monkeypatch):
|
|
"""The optimizer's memory is read back in inference mode there; updating it in place used to fail on the first step."""
|
|
import torch
|
|
from modules.training.lora import jobs as j, loop, train
|
|
tmp_path, jobs = run_env
|
|
with torch.inference_mode():
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2))
|
|
run = jobs.snapshot()["run"]
|
|
assert loop.has_checkpoint(run)
|
|
real = j.step
|
|
|
|
def fail(loss, seconds=None):
|
|
real(loss, seconds)
|
|
raise RuntimeError("the card fell over")
|
|
monkeypatch.setattr(j, "step", fail)
|
|
with torch.inference_mode(), pytest.raises(RuntimeError, match="fell over"):
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
|
|
assert jobs.snapshot()["state"] == "error" and loop.has_checkpoint(run), "a failed run keeps its saved state"
|
|
monkeypatch.setattr(j, "step", real)
|
|
with torch.inference_mode():
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
|
|
done = train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, preview_every=2, resume=run))
|
|
snap = jobs.snapshot()
|
|
assert done == "tiny_lora.safetensors" and snap["state"] == "done" and len(snap["loss"]) == 6
|
|
assert not loop.has_checkpoint(run)
|
|
|
|
|
|
def test_a_second_run_on_the_same_data_does_not_encode_it_again(run_env, monkeypatch):
|
|
from modules.training.lora import train
|
|
tmp_path, jobs = run_env
|
|
calls = {"video": 0, "text": 0}
|
|
real_item, real_text = train.encode_item, train._encode_text
|
|
|
|
def item(*a, **k):
|
|
calls["video"] += 1
|
|
return real_item(*a, **k)
|
|
|
|
def text(*a, **k):
|
|
calls["text"] += 1
|
|
return real_text(*a, **k)
|
|
monkeypatch.setattr(train, "encode_item", item)
|
|
monkeypatch.setattr(train, "_encode_text", text)
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
|
|
first = dict(calls)
|
|
assert first["video"] > 0 and first["text"] > 0 and any((tmp_path / "training" / "cache").glob("latent-*.pt"))
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=2))
|
|
assert calls == first, "the clips and captions came from the cache"
|
|
|
|
|
|
def test_pause_button_waits_at_the_next_step_and_stop_ends_a_paused_run(run_env, monkeypatch):
|
|
from modules.training.lora import jobs as j, train
|
|
tmp_path, jobs = run_env
|
|
real = j.step
|
|
|
|
def press(loss, seconds=None):
|
|
real(loss, seconds)
|
|
j.request_pause()
|
|
monkeypatch.setattr(j, "step", press)
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6))
|
|
snap = jobs.snapshot()
|
|
assert snap["state"] == "paused" and snap["paused"]["reason"] == "user" and snap["paused"]["step"] == 1
|
|
assert j.request_stop() and jobs.snapshot()["state"] == "stopped"
|
|
|
|
|
|
def test_a_run_that_cannot_be_resumed_says_why(run_env):
|
|
from modules.training.lora import train
|
|
_, jobs = run_env
|
|
with pytest.raises(ValueError, match="saved state is gone"):
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(resume="0123456789ab"))
|
|
assert jobs.snapshot()["state"] == "error"
|
|
|
|
|
|
def test_copies_of_the_lora_so_far_are_saved_as_name_stepN(run_env):
|
|
tmp_path, jobs = run_env
|
|
from modules.training.lora import train
|
|
train.run(tiny_h3_patcher(), FakeClip(), FakeVae(), FakeAudioVae(), "tiny", settings(steps=6, save_every=2))
|
|
names = sorted(p.name for p in (tmp_path / "loras").glob("*.safetensors"))
|
|
assert names == ["tiny_lora.safetensors", "tiny_lora_step2.safetensors", "tiny_lora_step4.safetensors"]
|