37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
"""
|
|
Thin wrapper pairing a STARFlow-V transformer with a comfy ModelPatcher so
|
|
it participates in ComfyUI's VRAM management (offload to CPU / reload to
|
|
GPU on demand) instead of being fully resident on GPU at import/load time.
|
|
"""
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
|
|
class STARFlowVModelWrapper:
|
|
def __init__(self, model: torch.nn.Module, cfg, load_device: Optional[torch.device] = None,
|
|
offload_device: Optional[torch.device] = None):
|
|
import comfy.model_management as mm
|
|
from comfy.model_patcher import ModelPatcher
|
|
|
|
self.cfg = cfg
|
|
load_device = load_device or mm.get_torch_device()
|
|
offload_device = offload_device or mm.unet_offload_device()
|
|
if not hasattr(model, "device"):
|
|
model.device = offload_device
|
|
size = mm.module_size(model)
|
|
self.patcher = ModelPatcher(model, load_device=load_device, offload_device=offload_device, size=size)
|
|
|
|
def load(self) -> torch.nn.Module:
|
|
import comfy.model_management as mm
|
|
mm.load_models_gpu([self.patcher])
|
|
return self.patcher.model
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
return self.patcher.load_device
|
|
|
|
@property
|
|
def model(self) -> torch.nn.Module:
|
|
return self.patcher.model
|