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
79 lines
3.4 KiB
Python
79 lines
3.4 KiB
Python
"""The graph a training run is queued as: the pipeline's own loaders, and the training node fed by them."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from ..._core import graph
|
|
|
|
NODE = "FunUITrainLora"
|
|
WANTED = (("model", "MODEL"), ("clip", "CLIP"), ("vae", "VAE"), ("audio_vae", "VAE"))
|
|
|
|
|
|
def _loaders(slots, schemas):
|
|
"""{input name: slot id}: the slots that hand the trainer its model, text encoder and VAEs.
|
|
|
|
The pipeline names them (a model's pipeline calls them model, clip, vae and audio_vae); a
|
|
pipeline that names them otherwise falls back to the first slot that makes that kind and
|
|
takes none of that kind in, which is a loader and not a patch on one.
|
|
"""
|
|
live = {s["id"]: s for s in slots if not s.get("bypassed")}
|
|
found = {}
|
|
for name, kind in WANTED:
|
|
if name in live and kind in schemas.outputs(live[name]["node"]):
|
|
found[name] = name
|
|
for name, kind in WANTED[:3]:
|
|
if name in found:
|
|
continue
|
|
for slot in live.values():
|
|
if slot["id"] in found.values() or slot["id"] == "audio_vae":
|
|
continue
|
|
if kind in schemas.outputs(slot["node"]) and kind not in schemas.inputs(slot["node"]).values():
|
|
found[name] = slot["id"]
|
|
break
|
|
return found
|
|
|
|
|
|
def _upstream(slots, roots):
|
|
"""The given slots and every slot they read from."""
|
|
by_id = {s["id"]: s for s in slots}
|
|
keep, todo = set(), list(roots)
|
|
while todo:
|
|
sid = todo.pop()
|
|
if sid in keep or sid not in by_id:
|
|
continue
|
|
keep.add(sid)
|
|
for value in (by_id[sid].get("inputs") or {}).values():
|
|
if graph.is_link(value):
|
|
todo.append(value[0])
|
|
return [s for s in slots if s["id"] in keep]
|
|
|
|
|
|
def _model_file(slot):
|
|
"""The file and dtypes a plain diffusion model loader slot names, or {} when the slot is anything else."""
|
|
given = slot.get("inputs") or {}
|
|
name = given.get("model_name")
|
|
if slot.get("node") != "FunUIDiffusionModelLoader" or not isinstance(name, str) or not name or name == "None":
|
|
return {}
|
|
return {"model_name": name, "weight_dtype": given.get("weight_dtype", "default"), "compute_dtype": given.get("compute_dtype", "default")}
|
|
|
|
|
|
def assemble(slots, set_name: str, values: dict, schemas=None):
|
|
"""(prompt, problems). `values` are the training node's settings; the loaders come from `slots`."""
|
|
schemas = schemas or graph.from_comfyui()
|
|
bad = graph.shape_problems(slots)
|
|
if bad:
|
|
return {}, bad
|
|
found = _loaders(slots, schemas)
|
|
missing = [n for n in ("model", "clip", "vae") if n not in found]
|
|
if missing:
|
|
return {}, [f"the pipeline has no {', '.join(missing)} loader: training needs the model, the text encoder and the VAE"]
|
|
trainer = {"id": "train_lora", "node": NODE, "inputs": {"set_name": set_name, **values}}
|
|
by_id = {s["id"]: s for s in slots}
|
|
path = _model_file(by_id[found["model"]])
|
|
if path: # the trainer loads this file itself, once: no loader node runs in front of it
|
|
trainer["inputs"].update(path)
|
|
found = {n: sid for n, sid in found.items() if n != "model"}
|
|
kept = _upstream([s for s in slots if not s.get("bypassed")], found.values())
|
|
for name, sid in found.items():
|
|
trainer["inputs"][name] = [sid, 0]
|
|
return graph.build(kept + [trainer], schemas)
|