45 lines
1.4 KiB
Python
45 lines
1.4 KiB
Python
from collections.abc import Callable
|
|
from typing import TypeVar
|
|
|
|
import torch
|
|
|
|
import comfy.model_management
|
|
import comfy.model_patcher
|
|
|
|
|
|
T = TypeVar("T", bound=torch.Tensor)
|
|
|
|
|
|
class ManagedAuxiliaryModel:
|
|
"""Keep a small auxiliary model under ComfyUI's load/offload lifecycle."""
|
|
|
|
def __init__(self, factory: Callable[[], torch.nn.Module]):
|
|
model = factory().eval().requires_grad_(False)
|
|
load_device = comfy.model_management.vae_device()
|
|
offload_device = comfy.model_management.vae_offload_device()
|
|
model.to(offload_device)
|
|
|
|
self.model = model
|
|
self.patcher = comfy.model_patcher.CoreModelPatcher(
|
|
model,
|
|
load_device=load_device,
|
|
offload_device=offload_device,
|
|
)
|
|
|
|
def _dtype(self) -> torch.dtype:
|
|
parameter = next(self.model.parameters(), None)
|
|
return parameter.dtype if parameter is not None else torch.float32
|
|
|
|
def run(
|
|
self,
|
|
value: torch.Tensor,
|
|
postprocess: Callable[[torch.Tensor], T] | None = None,
|
|
) -> T:
|
|
comfy.model_management.load_models_gpu([self.patcher])
|
|
device = self.patcher.load_device
|
|
with comfy.model_management.cuda_device_context(device):
|
|
output = self.model(value.to(device=device, dtype=self._dtype()))
|
|
if postprocess is not None:
|
|
output = postprocess(output)
|
|
return output.to(comfy.model_management.intermediate_device())
|