diff --git a/comfy/model_management.py b/comfy/model_management.py index 304931eb0..d59d6b354 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1213,7 +1213,8 @@ 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: - v_tensor = comfy.memory_management.interpret_gathered_like(cast_geometry, weight._v_tensor)[0] + raw_tensor = comfy_aimdo.torch.aimdo_to_tensor(weight._v, device) + v_tensor = comfy.memory_management.interpret_gathered_like(cast_geometry, raw_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 19c9031ea..224e218e3 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -1542,7 +1542,6 @@ 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,7 +1556,6 @@ 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 33803b223..057535c8c 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 = s._v_tensor + xfer_dest = comfy_aimdo.torch.aimdo_to_tensor(s._v, device) resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature) if not resident: