Claude
3adebb731f
Old project/workflow names and browser settings migrate on load; /funpack redirects to /fun. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_016acRss69eH1235eRdFTZHY
322 lines
17 KiB
Python
322 lines
17 KiB
Python
"""The graph a run is queued as, and the module's routes."""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _needs_comfy(comfyui):
|
|
"""Needs ComfyUI's node schemas."""
|
|
|
|
|
|
def pipeline():
|
|
return [
|
|
{"id": "model", "node": "FunUIDiffusionModelLoader", "group": "Loaders", "inputs": {"model_name": "x.safetensors"}},
|
|
{"id": "clip", "node": "FunUICLIPLoader", "group": "Loaders", "inputs": {}},
|
|
{"id": "vae", "node": "FunUIVAELoader", "group": "Loaders", "inputs": {}},
|
|
{"id": "audio_vae", "node": "FunUIVAELoader", "group": "Loaders", "inputs": {}},
|
|
{"id": "lora", "node": "FunUILoraLoader", "group": "Loaders", "inputs": {"model": ["model", 0]}},
|
|
]
|
|
|
|
|
|
def test_the_trainer_is_fed_by_the_pipelines_loaders(monkeypatch):
|
|
from core import graph
|
|
from modules.training.lora import plan
|
|
monkeypatch.setattr(graph, "build", lambda slots, schemas=None, optional=None: ({s["id"]: s for s in slots}, []))
|
|
|
|
class Schemas:
|
|
def outputs(self, node):
|
|
return {"FunUIDiffusionModelLoader": ["MODEL"], "FunUICLIPLoader": ["CLIP", "STRING"], "FunUIVAELoader": ["VAE"],
|
|
"FunUILoraLoader": ["MODEL", "CLIP"]}[node]
|
|
|
|
def inputs(self, node):
|
|
return {"FunUILoraLoader": {"model": "MODEL"}}.get(node, {})
|
|
built, problems = plan.assemble(pipeline(), "mine", {"steps": 10}, Schemas())
|
|
assert problems == []
|
|
inputs = built["train_lora"]["inputs"]
|
|
assert "model" not in inputs and inputs["model_name"] == "x.safetensors" and inputs["weight_dtype"] == "default" # the file, not a loader node
|
|
assert "model" not in built # and no loader node runs in front of the trainer
|
|
assert inputs["clip"] == ["clip", 0] and inputs["audio_vae"] == ["audio_vae", 0] and inputs["set_name"] == "mine" and inputs["steps"] == 10
|
|
|
|
|
|
def test_a_model_slot_that_is_not_a_plain_loader_is_still_linked(monkeypatch):
|
|
from core import graph
|
|
from modules.training.lora import plan
|
|
monkeypatch.setattr(graph, "build", lambda slots, schemas=None, optional=None: ({s["id"]: s for s in slots}, []))
|
|
slots = pipeline()
|
|
slots[0]["inputs"] = {} # no file named: the loader node itself has to run
|
|
class Schemas:
|
|
def outputs(self, node):
|
|
return {"FunUIDiffusionModelLoader": ["MODEL"], "FunUICLIPLoader": ["CLIP"], "FunUIVAELoader": ["VAE"], "FunUILoraLoader": ["MODEL"]}[node]
|
|
|
|
def inputs(self, node):
|
|
return {"FunUILoraLoader": {"model": "MODEL"}}.get(node, {})
|
|
built, _ = plan.assemble(slots, "mine", {}, Schemas())
|
|
assert built["train_lora"]["inputs"]["model"] == ["model", 0] and "model" in built
|
|
assert "lora" not in built, "a patch on the model is not part of the training graph"
|
|
|
|
|
|
def test_a_pipeline_without_a_loader_is_refused_with_which(monkeypatch):
|
|
from modules.training.lora import plan
|
|
|
|
class Schemas:
|
|
def outputs(self, node):
|
|
return ["MODEL"]
|
|
|
|
def inputs(self, node):
|
|
return {}
|
|
built, problems = plan.assemble(pipeline()[:1], "mine", {}, Schemas())
|
|
assert built == {} and "clip, vae" in problems[0]
|
|
|
|
|
|
def test_routes_make_and_edit_a_set(tmp_path, monkeypatch):
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from core import config
|
|
from modules.training import lora
|
|
from core import addons
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
monkeypatch.setattr(addons, "enabled", lambda _id: False) # the add-on is off: a plan is refused
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
bad = await c.post("/m/sets", json={"name": "../x"})
|
|
made = await c.post("/m/sets", json={"name": "cats"})
|
|
added = await (await c.post("/m/sets/cats/add", json={"entries": [
|
|
{"src": {"kind": "output", "type": "output", "subfolder": "", "filename": "a.png"}, "caption": "cat, dog"}]})).json()
|
|
bulk = await (await c.post("/m/sets/cats/bulk", json={"op": "remove", "tag": "dog"})).json()
|
|
missing = await c.get("/m/sets/nope")
|
|
plan_off = await c.post("/m/plan", json={"set": "cats", "slots": []})
|
|
gone = await (await c.delete("/m/sets/cats")).json()
|
|
return bad.status, made.status, added, bulk, missing.status, plan_off.status, gone
|
|
bad, made, added, bulk, missing, plan_off, gone = asyncio.run(go())
|
|
assert bad == 400 and made == 200 and added["added"] == 1
|
|
assert bulk["changed"] == 1 and bulk["set"]["items"][0]["caption"] == "cat"
|
|
assert missing == 404 and plan_off == 409 and gone == {"deleted": True}
|
|
|
|
|
|
def test_rated_results_list_only_the_videos_that_are_still_here(tmp_path, monkeypatch):
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from core import config, projects
|
|
from modules.training import lora
|
|
from modules.training.lora import train
|
|
(tmp_path / "here.mp4").write_bytes(b"x")
|
|
doc = {"scenes": [{"id": "a", "rating": "10"}, {"id": "b", "rating": "10"}],
|
|
"scene_renders": {"a": {"media": {"type": "output", "subfolder": "", "filename": "here.mp4"}},
|
|
"b": {"media": {"type": "output", "subfolder": "", "filename": "gone.mp4"}}}}
|
|
|
|
class Project:
|
|
def to_dict(self):
|
|
return doc
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
monkeypatch.setattr(projects, "listing", lambda: [{"id": "p"}])
|
|
monkeypatch.setattr(projects, "get", lambda _id: Project())
|
|
monkeypatch.setattr(train, "roots", lambda: {"output": str(tmp_path)})
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
return await (await c.get("/m/rated")).json()
|
|
got = asyncio.run(go())
|
|
assert [r["media"]["filename"] for r in got["results"]] == ["here.mp4"] and got["counts"]["liked"] == 1
|
|
|
|
|
|
def test_a_finished_run_can_be_downloaded_and_only_its_own_file(tmp_path, monkeypatch):
|
|
import folder_paths
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from core import config
|
|
from modules.training import lora
|
|
from modules.training.lora import jobs
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
loras = tmp_path / "loras"
|
|
loras.mkdir()
|
|
(loras / "cats.safetensors").write_bytes(b"weights")
|
|
(tmp_path / "secret.safetensors").write_bytes(b"no")
|
|
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [str(loras)])
|
|
runs = jobs._dir()
|
|
import json
|
|
for run, name in (("aaaaaaaaaaaa", "cats.safetensors"), ("bbbbbbbbbbbb", "../secret.safetensors"), ("cccccccccccc", "gone.safetensors"), ("dddddddddddd", "")):
|
|
(runs / f"{run}.json").write_text(json.dumps({"run": run, "lora": name}))
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
ok = await c.get("/m/runs/aaaaaaaaaaaa/lora")
|
|
return ok.status, ok.headers.get("Content-Disposition"), await ok.read(), [(await c.get(f"/m/runs/{r}/lora")).status for r in ("bbbbbbbbbbbb", "cccccccccccc", "dddddddddddd", "eeeeeeeeeeee")]
|
|
status, disposition, body, refused = asyncio.run(go())
|
|
assert status == 200 and body == b"weights" and "attachment" in disposition and "cats.safetensors" in disposition
|
|
assert refused == [404, 404, 404, 404]
|
|
|
|
|
|
def test_a_run_lists_its_saved_copies_and_preview_and_serves_each(tmp_path, monkeypatch):
|
|
import json
|
|
import folder_paths
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from core import config
|
|
from modules.training import lora
|
|
from modules.training.lora import jobs
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
loras = tmp_path / "loras"
|
|
loras.mkdir()
|
|
for f in ("cats.safetensors", "cats_step200.safetensors", "cats_step50.safetensors", "cats_step50_2.safetensors", "cats_preview.safetensors", "dogs_step50.safetensors", "cats_stepx.safetensors"):
|
|
(loras / f).write_bytes(f.encode())
|
|
monkeypatch.setattr(folder_paths, "get_folder_paths", lambda kind: [str(loras)])
|
|
(jobs._dir() / "aaaaaaaaaaaa.json").write_text(json.dumps({"run": "aaaaaaaaaaaa", "name": "cats", "lora": "cats.safetensors"}))
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
rec = await (await c.get("/m/runs/aaaaaaaaaaaa")).json()
|
|
one = await c.get("/m/runs/aaaaaaaaaaaa/lora?file=cats_step200.safetensors")
|
|
other = await c.get("/m/runs/aaaaaaaaaaaa/lora?file=dogs_step50.safetensors")
|
|
return rec["files"], await one.read(), other.status
|
|
files, body, other = asyncio.run(go())
|
|
assert files == ["cats.safetensors", "cats_step50.safetensors", "cats_step50_2.safetensors", "cats_step200.safetensors", "cats_preview.safetensors"]
|
|
assert body == b"cats_step200.safetensors" and other == 404
|
|
|
|
|
|
def _png_with_prompt(path, texts, kind="X"):
|
|
import json
|
|
from PIL import Image, PngImagePlugin
|
|
info = PngImagePlugin.PngInfo()
|
|
info.add_text("prompt", json.dumps({str(i): {"class_type": kind, "inputs": {"text": t, "n": 3, "link": ["1", 0]}} for i, t in enumerate(texts)}))
|
|
Image.new("RGB", (4, 4)).save(path, pnginfo=info)
|
|
|
|
|
|
def test_a_file_made_from_another_prompt_is_not_the_one_that_was_rated(tmp_path):
|
|
from modules.training.lora import rated
|
|
mine, other, bare = tmp_path / "mine.png", tmp_path / "other.png", tmp_path / "bare.png"
|
|
_png_with_prompt(mine, ["a red fox runs through the snow, 35mm", "ugly, blurry"])
|
|
_png_with_prompt(other, ["a city street at night with neon signs"])
|
|
from PIL import Image
|
|
Image.new("RGB", (4, 4)).save(bare)
|
|
caption = "a red fox runs through snow"
|
|
assert rated.same_video(caption, mine)
|
|
assert not rated.same_video(caption, other), "same name, another video"
|
|
assert rated.same_video(caption, bare), "a file that carries no prompt gets the benefit of the doubt"
|
|
assert rated.same_video("", other), "a rating with no words has nothing to check"
|
|
assert rated.same_video(caption, tmp_path / "gone.png"), "unreadable: not judged"
|
|
|
|
|
|
def test_the_rated_list_keeps_every_file_but_marks_those_that_may_be_other_videos(tmp_path, monkeypatch):
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from core import config, projects
|
|
from modules.training import lora
|
|
from modules.training.lora import train
|
|
_png_with_prompt(tmp_path / "same.png", ["a red fox runs through the snow"])
|
|
_png_with_prompt(tmp_path / "clash.png", ["a city street at night"])
|
|
doc = {"scenes": [{"id": "a", "rating": "10", "text": "a red fox runs through the snow"}, {"id": "b", "rating": "10", "text": "a red fox runs through the snow"}],
|
|
"scene_renders": {"a": {"media": {"type": "output", "subfolder": "", "filename": "same.png"}, "text": "a red fox runs through the snow"},
|
|
"b": {"media": {"type": "output", "subfolder": "", "filename": "clash.png"}, "text": "a red fox runs through the snow"}}}
|
|
|
|
class Project:
|
|
def to_dict(self):
|
|
return doc
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
monkeypatch.setattr(projects, "listing", lambda: [{"id": "p"}])
|
|
monkeypatch.setattr(projects, "get", lambda _id: Project())
|
|
monkeypatch.setattr(train, "roots", lambda: {"output": str(tmp_path)})
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
return await (await c.get("/m/rated")).json()
|
|
got = asyncio.run(go())
|
|
assert [(r["media"]["filename"], r["doubtful"]) for r in got["results"]] == [("same.png", False), ("clash.png", True)] and got["doubtful"] == 1 and got["counts"]["liked"] == 2
|
|
|
|
|
|
def test_a_file_added_without_words_gets_the_prompt_it_was_made_from(tmp_path, monkeypatch):
|
|
import json
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from PIL import Image, PngImagePlugin
|
|
from core import config
|
|
from modules.training import lora
|
|
from modules.training.lora import train
|
|
info = PngImagePlugin.PngInfo()
|
|
info.add_text("prompt", json.dumps({"enhance": {"class_type": "FunUIEnhancePrompt", "inputs": {"text": "a red fox,\nrunning", "enabled": True}},
|
|
"negative": {"class_type": "CLIPTextEncode", "inputs": {"text": "ugly"}}}))
|
|
Image.new("RGB", (4, 4)).save(tmp_path / "made_here.png", pnginfo=info)
|
|
Image.new("RGB", (4, 4)).save(tmp_path / "plain.png")
|
|
Image.new("RGB", (4, 4)).save(tmp_path / "typed.png")
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
monkeypatch.setattr(train, "roots", lambda: {"output": str(tmp_path)})
|
|
out = lambda name: {"src": {"kind": "output", "type": "output", "subfolder": "", "filename": name}, "label": "own"}
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
await c.post("/m/sets", json={"name": "cats"})
|
|
kept = {**out("typed.png"), "caption": "typed by hand"}
|
|
r = await c.post("/m/sets/cats/add", json={"entries": [out("made_here.png"), out("plain.png"), kept]})
|
|
return (await r.json())["set"]["items"]
|
|
items = asyncio.run(go())
|
|
by_name = {i["src"]["filename"]: i["caption"] for i in items}
|
|
assert by_name["made_here.png"] == "a red fox, running"
|
|
assert "typed by hand" in [i["caption"] for i in items] and "" in [i["caption"] for i in items], "no prompt in the file: left empty; a caption already there: kept"
|
|
|
|
|
|
def test_a_rated_result_is_captioned_from_its_own_file_and_an_older_rating_on_a_reused_name_is_marked(tmp_path, monkeypatch):
|
|
import os
|
|
import time
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
from core import config, projects
|
|
from modules.training import lora
|
|
from modules.training.lora import train
|
|
now = time.time() * 1000
|
|
_png_with_prompt(tmp_path / "now.png", ["a red fox runs through the snow"], "FunUIEnhancePrompt")
|
|
_png_with_prompt(tmp_path / "reused.png", ["a city street at night"], "FunUIEnhancePrompt")
|
|
os.utime(tmp_path / "now.png", ((now - 60_000) / 1000,) * 2)
|
|
_png_with_prompt(tmp_path / "bare.png", [], "FunUIEnhancePrompt")
|
|
os.utime(tmp_path / "bare.png", ((now - 60_000) / 1000,) * 2)
|
|
doc = {"scenes": [{"id": "c", "rating": "10", "rated": {"text": "the project's words", "at": now}}, {"id": "a", "rating": "10", "rated": {"text": "something the project kept that is wrong", "at": now}},
|
|
{"id": "b", "rating": "10", "rated": {"text": "a red fox runs through the snow", "at": now - 86_400_000}}],
|
|
"scene_renders": {"c": {"media": {"type": "output", "subfolder": "", "filename": "bare.png"}},
|
|
"a": {"media": {"type": "output", "subfolder": "", "filename": "now.png"}},
|
|
"b": {"media": {"type": "output", "subfolder": "", "filename": "reused.png"}}}}
|
|
|
|
class Project:
|
|
def to_dict(self):
|
|
return doc
|
|
monkeypatch.setattr(config, "TRAINING_DIR", tmp_path)
|
|
monkeypatch.setattr(projects, "listing", lambda: [{"id": "p"}])
|
|
monkeypatch.setattr(projects, "get", lambda _id: Project())
|
|
monkeypatch.setattr(train, "roots", lambda: {"output": str(tmp_path)})
|
|
|
|
async def go():
|
|
app, table = web.Application(), web.RouteTableDef()
|
|
lora.routes(table, "/m", web)
|
|
app.add_routes(table)
|
|
async with TestClient(TestServer(app)) as c:
|
|
return await (await c.get("/m/rated")).json()
|
|
got = asyncio.run(go())
|
|
assert {r["media"]["filename"]: r["caption"] for r in got["results"]}["bare.png"] == "", "a file with no prompt of its own gets none: the project's words are never used"
|
|
got["results"] = [r for r in got["results"] if r["media"]["filename"] != "bare.png"]
|
|
assert [(r["media"]["filename"], r["caption"], r["doubtful"]) for r in got["results"]] == [("now.png", "a red fox runs through the snow", False), ("reused.png", "a city street at night", True)]
|
|
assert got["doubtful"] == 1
|
|
|
|
|
|
def test_the_training_node_is_never_answered_from_comfys_cache():
|
|
from modules.training.lora.nodes import FunUITrainLora
|
|
a, b = FunUITrainLora.fingerprint_inputs(resume="r1"), FunUITrainLora.fingerprint_inputs(resume="r1")
|
|
assert a != b, "NaN never equals itself: ComfyUI runs the node again even with identical inputs"
|