366 lines
11 KiB
Python
366 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
import comfy.model_management as model_management
|
|
|
|
from .kj_loader import CheckpointLoaderResident, DiffusionModelLoaderResident, DiffusionModelSelectorResident
|
|
from .residency import REGISTRY
|
|
|
|
|
|
def _entry_report_json(obj: Any) -> str:
|
|
entry = REGISTRY.entry_for_object(obj)
|
|
if entry is None:
|
|
return json.dumps(
|
|
{
|
|
"tracked": False,
|
|
"policy": REGISTRY.get_policy(),
|
|
"message": "Object is not currently bound in the GPU Resident Loader registry.",
|
|
},
|
|
indent=2,
|
|
sort_keys=True,
|
|
)
|
|
REGISTRY.refresh_runtime_state()
|
|
payload = {
|
|
"tracked": True,
|
|
"policy": REGISTRY.get_policy(),
|
|
"entry": entry.as_dict(),
|
|
}
|
|
return json.dumps(payload, indent=2, sort_keys=True)
|
|
|
|
|
|
def _patcher_for_clip(clip):
|
|
return getattr(clip, "patcher", None)
|
|
|
|
|
|
def _patcher_for_vae(vae):
|
|
return getattr(vae, "patcher", None)
|
|
|
|
|
|
def _preload_patcher(patcher, *, sticky: bool, priority: int) -> None:
|
|
if patcher is None:
|
|
raise RuntimeError("Expected a patcher-capable object, but no patcher was found.")
|
|
model_management.load_models_gpu([patcher], force_full_load=True)
|
|
REGISTRY.set_sticky(patcher, sticky=sticky, priority=priority)
|
|
REGISTRY.touch(patcher)
|
|
REGISTRY.refresh_runtime_state()
|
|
|
|
|
|
def _evict_patcher(patcher, *, unpatch_weights: bool) -> bool:
|
|
if patcher is None:
|
|
raise RuntimeError("Expected a patcher-capable object, but no patcher was found.")
|
|
unloaded = False
|
|
for loaded in list(model_management.current_loaded_models):
|
|
if loaded.model is patcher or loaded.model.is_clone(patcher):
|
|
loaded.model_unload(unpatch_weights=unpatch_weights)
|
|
unloaded = True
|
|
if unloaded and hasattr(model_management, "soft_empty_cache"):
|
|
model_management.soft_empty_cache()
|
|
REGISTRY.refresh_runtime_state()
|
|
return unloaded
|
|
|
|
|
|
class SetGlobalResidencyPolicy:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"policy": (
|
|
["legacy", "balanced", "prefer_gpu", "sticky_gpu"],
|
|
{"default": "sticky_gpu", "tooltip": "Select the global ingest/offload policy used by the startup patcher."},
|
|
)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("active_policy",)
|
|
FUNCTION = "set_policy"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def set_policy(self, policy: str):
|
|
return (REGISTRY.set_policy(policy),)
|
|
|
|
|
|
class RegistrySnapshot:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("registry_json",)
|
|
FUNCTION = "snapshot"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def snapshot(self):
|
|
return (REGISTRY.snapshot_json(),)
|
|
|
|
|
|
class PinModelResidency:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"sticky": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "pin"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def pin(self, model, sticky: bool, priority: int):
|
|
REGISTRY.set_sticky(model, sticky=sticky, priority=priority)
|
|
return (model,)
|
|
|
|
|
|
class PinClipResidency:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"sticky": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP",)
|
|
FUNCTION = "pin"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def pin(self, clip, sticky: bool, priority: int):
|
|
patcher = _patcher_for_clip(clip)
|
|
if patcher is None:
|
|
raise RuntimeError("Expected a CLIP object with a patcher, but no patcher was found.")
|
|
REGISTRY.set_sticky(patcher, sticky=sticky, priority=priority)
|
|
return (clip,)
|
|
|
|
|
|
class PinVAEResidency:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
"sticky": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE",)
|
|
FUNCTION = "pin"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def pin(self, vae, sticky: bool, priority: int):
|
|
patcher = _patcher_for_vae(vae)
|
|
if patcher is None:
|
|
raise RuntimeError("Expected a VAE object with a patcher, but no patcher was found.")
|
|
REGISTRY.set_sticky(patcher, sticky=sticky, priority=priority)
|
|
return (vae,)
|
|
|
|
|
|
class PreloadModelToGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"sticky": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "preload"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def preload(self, model, sticky: bool, priority: int):
|
|
_preload_patcher(model, sticky=sticky, priority=priority)
|
|
return (model,)
|
|
|
|
|
|
class PreloadClipToGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"sticky": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP",)
|
|
FUNCTION = "preload"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def preload(self, clip, sticky: bool, priority: int):
|
|
_preload_patcher(_patcher_for_clip(clip), sticky=sticky, priority=priority)
|
|
return (clip,)
|
|
|
|
|
|
class PreloadVAEToGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
"sticky": ("BOOLEAN", {"default": True}),
|
|
"priority": ("INT", {"default": 0, "min": -100, "max": 100, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE",)
|
|
FUNCTION = "preload"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def preload(self, vae, sticky: bool, priority: int):
|
|
_preload_patcher(_patcher_for_vae(vae), sticky=sticky, priority=priority)
|
|
return (vae,)
|
|
|
|
|
|
class EvictModelFromGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"unpatch_weights": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL", "STRING")
|
|
RETURN_NAMES = ("model", "eviction_status")
|
|
FUNCTION = "evict"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def evict(self, model, unpatch_weights: bool):
|
|
unloaded = _evict_patcher(model, unpatch_weights=unpatch_weights)
|
|
return model, ("evicted" if unloaded else "not_loaded")
|
|
|
|
|
|
class EvictClipFromGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"clip": ("CLIP",),
|
|
"unpatch_weights": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CLIP", "STRING")
|
|
RETURN_NAMES = ("clip", "eviction_status")
|
|
FUNCTION = "evict"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def evict(self, clip, unpatch_weights: bool):
|
|
unloaded = _evict_patcher(_patcher_for_clip(clip), unpatch_weights=unpatch_weights)
|
|
return clip, ("evicted" if unloaded else "not_loaded")
|
|
|
|
|
|
class EvictVAEFromGPU:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
"unpatch_weights": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE", "STRING")
|
|
RETURN_NAMES = ("vae", "eviction_status")
|
|
FUNCTION = "evict"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def evict(self, vae, unpatch_weights: bool):
|
|
unloaded = _evict_patcher(_patcher_for_vae(vae), unpatch_weights=unpatch_weights)
|
|
return vae, ("evicted" if unloaded else "not_loaded")
|
|
|
|
|
|
class ReportModelResidency:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"model": ("MODEL",)}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("report_json",)
|
|
FUNCTION = "report"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def report(self, model):
|
|
return (_entry_report_json(model),)
|
|
|
|
|
|
class ReportClipResidency:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"clip": ("CLIP",)}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("report_json",)
|
|
FUNCTION = "report"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def report(self, clip):
|
|
return (_entry_report_json(_patcher_for_clip(clip)),)
|
|
|
|
|
|
class ReportVAEResidency:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"vae": ("VAE",)}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("report_json",)
|
|
FUNCTION = "report"
|
|
CATEGORY = "GPU Resident Loader/residency"
|
|
|
|
def report(self, vae):
|
|
return (_entry_report_json(_patcher_for_vae(vae)),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DiffusionModelSelectorResident": DiffusionModelSelectorResident,
|
|
"DiffusionModelLoaderResident": DiffusionModelLoaderResident,
|
|
"CheckpointLoaderResident": CheckpointLoaderResident,
|
|
"SetGlobalResidencyPolicy": SetGlobalResidencyPolicy,
|
|
"RegistrySnapshot": RegistrySnapshot,
|
|
"PinModelResidency": PinModelResidency,
|
|
"PinClipResidency": PinClipResidency,
|
|
"PinVAEResidency": PinVAEResidency,
|
|
"PreloadModelToGPU": PreloadModelToGPU,
|
|
"PreloadClipToGPU": PreloadClipToGPU,
|
|
"PreloadVAEToGPU": PreloadVAEToGPU,
|
|
"EvictModelFromGPU": EvictModelFromGPU,
|
|
"EvictClipFromGPU": EvictClipFromGPU,
|
|
"EvictVAEFromGPU": EvictVAEFromGPU,
|
|
"ReportModelResidency": ReportModelResidency,
|
|
"ReportClipResidency": ReportClipResidency,
|
|
"ReportVAEResidency": ReportVAEResidency,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DiffusionModelSelectorResident": "Diffusion Model Selector Resident",
|
|
"DiffusionModelLoaderResident": "Diffusion Model Loader Resident",
|
|
"CheckpointLoaderResident": "Checkpoint Loader Resident",
|
|
"SetGlobalResidencyPolicy": "Set Global Residency Policy",
|
|
"RegistrySnapshot": "Registry Snapshot",
|
|
"PinModelResidency": "Pin Model Residency",
|
|
"PinClipResidency": "Pin CLIP Residency",
|
|
"PinVAEResidency": "Pin VAE Residency",
|
|
"PreloadModelToGPU": "Preload Model To GPU",
|
|
"PreloadClipToGPU": "Preload CLIP To GPU",
|
|
"PreloadVAEToGPU": "Preload VAE To GPU",
|
|
"EvictModelFromGPU": "Evict Model From GPU",
|
|
"EvictClipFromGPU": "Evict CLIP From GPU",
|
|
"EvictVAEFromGPU": "Evict VAE From GPU",
|
|
"ReportModelResidency": "Report Model Residency",
|
|
"ReportClipResidency": "Report CLIP Residency",
|
|
"ReportVAEResidency": "Report VAE Residency",
|
|
}
|