diff --git a/comfy/model_management.py b/comfy/model_management.py index b8d0801ec..6018c1ab6 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1208,8 +1208,7 @@ def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False, str signature = comfy_aimdo.model_vbar.vbar_fault(weight._v) if signature is not None: - raw_tensor = comfy_aimdo.torch.aimdo_to_tensor(weight._v, device) - v_tensor = comfy.memory_management.interpret_gathered_like(cast_geometry, raw_tensor)[0] + v_tensor = comfy.memory_management.interpret_gathered_like(cast_geometry, weight._v_tensor)[0] if not comfy_aimdo.model_vbar.vbar_signature_compare(signature, weight._v_signature): weight._v_signature = signature #Send it over diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index f03b4c3be..b9a117a7c 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -1543,6 +1543,7 @@ class ModelPatcherDynamic(ModelPatcher): if vbar is not None and not hasattr(m, "_v"): m._v = vbar.alloc(v_weight_size) + m._v_tensor = comfy_aimdo.torch.aimdo_to_tensor(m._v, device_to) allocated_size += v_weight_size else: @@ -1557,6 +1558,7 @@ class ModelPatcherDynamic(ModelPatcher): weight_size = geometry.numel() * geometry.element_size() if vbar is not None and not hasattr(weight, "_v"): weight._v = vbar.alloc(weight_size) + weight._v_tensor = comfy_aimdo.torch.aimdo_to_tensor(weight._v, device_to) weight._model_dtype = model_dtype allocated_size += weight_size vbar.set_watermark_limit(allocated_size) diff --git a/comfy/ops.py b/comfy/ops.py index 0f4eca7c7..ea0d70702 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -87,7 +87,7 @@ def cast_bias_weight_with_vbar(s, dtype, device, bias_dtype, non_blocking, compu signature = comfy_aimdo.model_vbar.vbar_fault(s._v) if signature is not None: - xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device) + xfer_dest = s._v_tensor resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature) if not resident: