import math import ctypes import dataclasses import torch from typing import NamedTuple import comfy_aimdo.host_buffer from comfy.quant_ops import QuantizedTensor class TensorFileSlice(NamedTuple): file_ref: object lock: object offset: int size: int def read_tensor_file_slice_into(tensor, destination, stream=None, destination2=None): 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: view.release() class TensorGeometry(NamedTuple): 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): return math.prod(self.shape) def tensors_to_geometries(tensors, dtype=None): 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)) return geometries def vram_aligned_size(tensor): 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 return (size + aligment_req - 1) // aligment_req * aligment_req def interpret_gathered_like(tensors, gathered): 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"]) return dest_views aimdo_enabled = False extra_ram_release_callback = None RAM_CACHE_HEADROOM = 0 def set_ram_cache_release_state(callback, headroom): global extra_ram_release_callback global RAM_CACHE_HEADROOM extra_ram_release_callback = callback RAM_CACHE_HEADROOM = max(0, int(headroom)) def extra_ram_release(target, free_active=False): if extra_ram_release_callback is None: return 0 return extra_ram_release_callback(target, free_active=free_active)