file_ref: object

lock: object

offset: int

if isinstance(tensor, QuantizedTensor):

if not read_tensor_file_slice_into(tensor._qdata,

destination._qdata if destination is not None else None, stream=stream,

destination2=(destination2._qdata if destination2 is not None else None)):

return False

if destination is not None:

dst_orig_dtype = destination._params.orig_dtype

destination._params.copy_from(tensor._params, non_blocking=False)

destination._params = dataclasses.replace(destination._params, orig_dtype=dst_orig_dtype)

if destination2 is not None:

dst_orig_dtype = destination2._params.orig_dtype

destination2._params.copy_from(destination._params if destination is not None else tensor._params, non_blocking=True)

destination2._params = dataclasses.replace(destination2._params, orig_dtype=dst_orig_dtype)

return True

info = getattr(tensor.untyped_storage(), "_comfy_tensor_file_slice", None)

if info is None:

return False

if destination is not None and destination.device.type != "cpu" and destination2 is None:

destination2 = destination

destination = None

file_obj = info.file_ref

if (file_obj is None

or (destination is None and destination2 is None)

or (destination is not None and (destination.device.type != "cpu" or destination.numel() * destination.element_size() < info.size))

or (destination2 is not None and (destination2.device.type == "cpu" or destination2.numel() * destination2.element_size() < info.size))

or tensor.numel() * tensor.element_size() != info.size

or tensor.storage_offset() != 0

or not tensor.is_contiguous()):

return False

if info.size == 0:

return True

if destination is None:

stream_ptr = getattr(stream, "cuda_stream", 0) if stream is not None else 0

comfy_aimdo.host_buffer.read_file_to_device(file_obj, info.offset, info.size,

stream_ptr, destination2.data_ptr(),

destination2.device.index,

mark_cold=False)

return True

hostbuf = getattr(destination.untyped_storage(), "_comfy_hostbuf", None)

if hostbuf is not None:

stream_ptr = getattr(stream, "cuda_stream", 0) if stream is not None else 0

device_ptr = destination2.data_ptr() if destination2 is not None else 0

with info.lock:

hostbuf.read_file_slice(file_obj, info.offset, info.size,

offset=destination.data_ptr() - hostbuf.get_raw_address(),

stream=stream_ptr,

device_ptr=device_ptr,

device=None if destination2 is None else destination2.device.index)

return True

if not hasattr(file_obj, "seek") or not hasattr(file_obj, "readinto"):

return False

buf_type = ctypes.c_ubyte * info.size

view = memoryview(buf_type.from_address(destination.data_ptr()))

try:

with info.lock:

file_obj.seek(info.offset)

done = 0

while done < info.size:

try:

n = file_obj.readinto(view[done:])

except OSError:

return False

if n <= 0:

return False

done += n

return True

finally:

shape: any

dtype: torch.dtype

def element_size(self):

info = torch.finfo(self.dtype) if self.dtype.is_floating_point else torch.iinfo(self.dtype)

return info.bits // 8

def numel(self):

geometries = []

for t in tensors:

if t is None or isinstance(t, QuantizedTensor):

geometries.append(t)

continue

tdtype = t.dtype

if hasattr(t, "_model_dtype"):

tdtype = t._model_dtype

if dtype is not None:

tdtype = dtype

geometries.append(TensorGeometry(shape=t.shape, dtype=tdtype))

if isinstance(tensor, list):

return sum([vram_aligned_size(t) for t in tensor])

if isinstance(tensor, QuantizedTensor):

inner_tensors, _ = tensor.__tensor_flatten__()

return vram_aligned_size([ getattr(tensor, attr) for attr in inner_tensors ])

if tensor is None:

return 0

size = tensor.numel() * tensor.element_size()

aligment_req = 1024

offset = 0

dest_views = []

if gathered.dim() != 1 or gathered.element_size() != 1:

raise ValueError(f"Buffer must be 1D and single-byte (got {gathered.dim()}D {gathered.dtype})")

for tensor in tensors:

if tensor is None:

dest_views.append(None)

continue

if isinstance(tensor, QuantizedTensor):

inner_tensors, qt_ctx = tensor.__tensor_flatten__()

templates = { attr: getattr(tensor, attr) for attr in inner_tensors }

else:

templates = { "data": tensor }

actuals = {}

for attr, template in templates.items():

size = template.numel() * template.element_size()

if offset + size > gathered.numel():

raise ValueError(f"Buffer too small: needs {offset + size} bytes, but only has {gathered.numel()}. ")

actuals[attr] = gathered[offset:offset+size].view(dtype=template.dtype).view(template.shape)

offset += vram_aligned_size(template)

if isinstance(tensor, QuantizedTensor):

dest_views.append(QuantizedTensor.__tensor_unflatten__(actuals, qt_ctx, 0, 0))

else:

dest_views.append(actuals["data"])

global extra_ram_release_callback

global RAM_CACHE_HEADROOM

extra_ram_release_callback = callback