mmgp 3.3.4__py3-none-any.whl → 3.4.1__py3-none-any.whl

mmgp/offload.py

@@ -1,4 +1,4 @@
-# ------------------ Memory Management 3.3.4 for the GPU Poor by DeepBeepMeep (mmgp)------------------
+# ------------------ Memory Management 3.4.1 for the GPU Poor by DeepBeepMeep (mmgp)------------------
 #
 # This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  can run smoothly on a 24 GB GPU limited card. 
 # This a replacement for the accelerate library that should in theory manage offloading, but doesn't work properly with models that are loaded / unloaded several
@@ -89,7 +89,7 @@ mmm = safetensors2.mmm
 default_verboseLevel = 1
 
 ONE_MB =  1048576
-sizeofbfloat16 = torch.bfloat16.itemsize
+sizeofhalffloat = torch.bfloat16.itemsize
 sizeofint8 = torch.int8.itemsize
 total_pinned_bytes = 0
 max_pinnable_bytes = 0
@@ -149,10 +149,17 @@ def _compute_verbose_level(level):
     safetensors2.verboseLevel = level
     return level
 
-def _get_perc_reserved_mem_max(perc_reserved_mem_max):
-    if perc_reserved_mem_max<=0:             
+def _get_perc_reserved_mem_max(perc_reserved_mem_max = 0):
+    if perc_reserved_mem_max <=0:
+        perc_reserved_mem_max = os.getenv("perc_reserved_mem_max", 0)
+
+    if perc_reserved_mem_max <= 0:             
         perc_reserved_mem_max = 0.40 if os.name == 'nt' else 0.5        
-    return  perc_reserved_mem_max 
+    return perc_reserved_mem_max
+    
+def _get_max_reservable_memory(perc_reserved_mem_max = 0):
+    max_reservable_memory = perc_reserved_mem_max * physical_memory
+    return  max_reservable_memory
 
 def _detect_main_towers(model, min_floors = 5):
     cur_blocks_prefix = None
@@ -214,7 +221,7 @@ def _get_model(model_path):
     _filename = _path[-1]
     _path = _path[:-1]
     if len(_path)<=1:
-        raise("file not found")
+        raise Exception("file not found")
     else:
         try:
             from huggingface_hub import  hf_hub_download #snapshot_download,    
@@ -290,8 +297,9 @@ def _get_tensor_ref(p):
         return p.data_ptr()
 
 
-# BIG_TENSOR_MAX_SIZE = 2**28 # 256 MB
-BIG_TENSOR_MAX_SIZE = 2**27 # 128 MB
+BIG_TENSOR_MAX_SIZE = 2**28 # 256 MB
+BIG_TENSOR_MIN_SIZE = 2**26 # 64 MB
+RESERVED_RAM_MIN_AVAILABLE = 2**27 # 128 MB
 
 def _extract_tie_weights_from_sd(sd , sd_name, verboseLevel =1):
     tied_weights = {}
@@ -322,7 +330,7 @@ def _extract_tie_weights_from_sd(sd , sd_name, verboseLevel =1):
 
 def _pin_sd_to_memory(sd, sd_name, tied_weights = None, gig_tensor_size = BIG_TENSOR_MAX_SIZE, verboseLevel = 1):
     global max_pinnable_bytes, total_pinned_bytes
-    if max_pinnable_bytes > 0 and  max_pinnable_bytes >= max_pinnable_bytes:
+    if max_pinnable_bytes > 0 and  total_pinned_bytes >= max_pinnable_bytes:
 
         if  verboseLevel>=1 :
             print(f"Unable pin data of '{sd_name}' to reserved RAM as there is no reserved RAM left")
@@ -357,6 +365,12 @@ def _pin_sd_to_memory(sd, sd_name, tied_weights = None, gig_tensor_size = BIG_TE
     last_big_tensor = 0
     total = 0  
 
+    try:
+        dummy_pinned_tensor = torch.empty( RESERVED_RAM_MIN_AVAILABLE, dtype= torch.uint8, pin_memory=True, device="cpu")
+    except:
+        print("There isn't any Reserved RAM left, you may need to choose a profile with a higher number that requires less Reserved RAM or set OS env 'perc_reserved_mem_max' to a value less 0.3")
+        return
+    
     for size in big_tensors_sizes:
         try:
             current_big_tensor = torch.empty( size, dtype= torch.uint8, pin_memory=True, device="cpu")
@@ -367,6 +381,7 @@ def _pin_sd_to_memory(sd, sd_name, tied_weights = None, gig_tensor_size = BIG_TE
 
         last_big_tensor += 1
         total += size
+    del dummy_pinned_tensor
 
         
     tensor_no = 0
@@ -401,10 +416,10 @@ def _pin_sd_to_memory(sd, sd_name, tied_weights = None, gig_tensor_size = BIG_TE
     return 
 
 
-def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = True, big_tensor_size = BIG_TENSOR_MAX_SIZE, verboseLevel = 1):
+def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = True, big_tensor_size = BIG_TENSOR_MAX_SIZE, perc_reserved_mem_max = 0,verboseLevel = 1):
 
     global max_pinnable_bytes, total_pinned_bytes
-    if max_pinnable_bytes > 0 and  max_pinnable_bytes >= max_pinnable_bytes:
+    if max_pinnable_bytes > 0 and  total_pinned_bytes >= max_pinnable_bytes:
 
         if  verboseLevel>=1 :
             print(f"Unable pin data of '{model_id}' to reserved RAM as there is no reserved RAM left")
@@ -414,6 +429,8 @@ def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = Tru
         towers_names, _ = _detect_main_towers(model)
 
 
+    perc_reserved_mem_max = _get_perc_reserved_mem_max(perc_reserved_mem_max)
+    max_reservable_memory = _get_max_reservable_memory(perc_reserved_mem_max) 
 
     current_big_tensor_size = 0
     big_tensor_no  = 0
@@ -502,8 +519,12 @@ def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = Tru
     
 
     failed_planned_allocation = False
-
     gc.collect()
+    try:
+        dummy_pinned_tensor = torch.empty( RESERVED_RAM_MIN_AVAILABLE, dtype= torch.uint8, pin_memory=True, device="cpu")
+    except:
+        print("There isn't any Reserved RAM left, you may need to choose a profile with a higher number that requires less Reserved RAM or set OS env 'perc_reserved_mem_max' to a value less than{perc_reserved_mem_max}")
+        return
 
     last_allocated_big_tensor = -1        
     tensor_no = 0
@@ -530,16 +551,21 @@ def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = Tru
             big_tensor_no, offset, length = tensor_map_indexes[tensor_no]
             if last_allocated_big_tensor <  big_tensor_no:
                 last_allocated_big_tensor += 1
-                size = big_tensors_sizes[last_allocated_big_tensor] 
+                size = max(big_tensors_sizes[last_allocated_big_tensor], BIG_TENSOR_MIN_SIZE) 
                 try:
-                    # if total > 7000 * ONE_MB:
-                    #     raise  Exception ("test no more reserved RAM")
+                    if max_reservable_memory > 0 and ( (total_pinned_bytes + total + size) >= max_reservable_memory):
+                        dummy_pinned_tensor = None
+                        failed_planned_allocation = True
+                        max_pinnable_bytes = total_pinned_bytes + total
+                        break
+
                     current_big_tensor = torch.empty( size, dtype= torch.uint8, pin_memory=True, device="cpu")
                     big_tensors.append(current_big_tensor)
                 except:
-                    print(f"Unable to pin more tensors for this model as the maximum reservable memory has been reached ({total/ONE_MB:.2f})")
-                    max_pinnable_bytes = total + total_pinned_bytes
+                    print(f"Unable to pin more tensors for this model as the maximum reservable memory has been reached ({total/ONE_MB:.2f}).")
+                    dummy_pinned_tensor = None
                     failed_planned_allocation = True
+                    max_pinnable_bytes = total_pinned_bytes + total
                     break
 
                 total += size
@@ -569,6 +595,7 @@ def _pin_to_memory(model, model_id, partialPinning = False, pinnedPEFTLora = Tru
                 p.data = _move_to_pinned_tensor(p.data, current_big_tensor, offset, length)
             tensor_no += 1
         del p
+    del dummy_pinned_tensor
     model._pinned_bytes = total
     total_pinned_bytes += total
     del params_dict
@@ -591,8 +618,23 @@ def _welcome():
     if welcome_displayed:
          return 
     welcome_displayed = True
-    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.3.4) by DeepBeepMeep ************{ENDC}{UNBOLD}")
+    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.4.1) by DeepBeepMeep ************{ENDC}{UNBOLD}")
 
+def change_dtype(model, new_dtype, exclude_buffers = False):
+    for submodule_name, submodule in model.named_modules():  
+        if hasattr(submodule, "_lock_dtype"):
+            continue
+        for n, p in submodule.named_parameters(recurse = False):
+            if p.data.dtype != new_dtype:
+                p.data = p.data.to(new_dtype)
+
+        if not exclude_buffers:
+            for p in submodule.buffers(recurse=False):
+                if p.data.dtype != new_dtype:
+                    p.data = p.data.to(new_dtype)
+
+    return model
+            
 def _extract_num_from_str(num_in_str):
     size = len(num_in_str)
     for i in range(size):
@@ -728,12 +770,16 @@ def _quantize(model_to_quantize, weights=qint8, verboseLevel = 1, threshold = 2*
                 tied_weights[submodule_name]=  (n, ) + match 
             else:
                 cache_ref[ref] = (submodule_name, n)
-                size  += torch.numel(p.data) * sizeofbfloat16
+                size  += torch.numel(p.data) * sizeofhalffloat
 
         for p in submodule.buffers(recurse=False):
-            size  += torch.numel(p.data) * sizeofbfloat16
-
+            size  += torch.numel(p.data) * sizeofhalffloat
 
+        already_added = False
+        if hasattr(submodule, "_lock_dtype"):
+            submodule_size += size
+            submodule_names.append(submodule_name)
+            already_added = True
 
         if not any(submodule_name.startswith(pre) for pre in tower_names):
             flush = False
@@ -751,8 +797,9 @@ def _quantize(model_to_quantize, weights=qint8, verboseLevel = 1, threshold = 2*
                 submodule_size = 0
                 submodule_names = []
             prev_blocks_prefix = cur_blocks_prefix
-            submodule_size += size
-            submodule_names.append(submodule_name)
+            if not already_added:
+                submodule_size += size
+                submodule_names.append(submodule_name)
         total_size += size
 
     if submodule_size >0  : 
@@ -853,7 +900,7 @@ def split_linear_modules(model, map ):
                     sub_bias = torch.split(bias, split_sizes, dim=0)
                     for sub_name, _subdata, _subbias, _subscale in zip(mapped_modules, sub_data, sub_bias, sub_scale):
                         with init_empty_weights():
-                            sub_module = QLinear(_subdata.shape[1], _subdata.shape[0], bias=bias != None, device ="cpu", dtype=torch.bfloat16)
+                            sub_module = QLinear(_subdata.shape[1], _subdata.shape[0], bias=bias != None, device ="cpu", dtype=weight.dtype)
                         sub_module.weight = torch.nn.Parameter(WeightQBytesTensor.create(weight.qtype, weight.axis, _subdata.size(), weight.stride(), _subdata, _subscale, activation_qtype=weight.activation_qtype, requires_grad=weight.requires_grad ))
                         if bias != None:                        
                             sub_module.bias = torch.nn.Parameter(_subbias)
@@ -866,7 +913,7 @@ def split_linear_modules(model, map ):
                     sub_bias = torch.split(bias, split_sizes, dim=0)
                     for sub_name, subdata, subbias in zip(mapped_modules, sub_data, sub_bias):
                         with init_empty_weights():
-                            sub_module = torch.nn.Linear( subdata.shape[1], subdata.shape[0], bias=bias != None, device ="cpu", dtype=torch.bfloat16)
+                            sub_module = torch.nn.Linear( subdata.shape[1], subdata.shape[0], bias=bias != None, device ="cpu", dtype=weight.dtype)
                         sub_module.weight = torch.nn.Parameter(subdata , requires_grad=False)
                         if bias != None:
                             sub_module.bias = torch.nn.Parameter(subbias)
@@ -1045,8 +1092,8 @@ def load_loras_into_model(model, lora_path, lora_multi = None, activate_all_lora
                     #     loras_module_data = dict()
                     #     loras_model_data[module] = loras_module_data
                     loras_adapter_data =  loras_module_data.get(adapter_name, None)
-                    lora_A = None if lora_A == None else lora_A.to(torch.bfloat16) 
-                    lora_B = None if lora_B == None else lora_B.to(torch.bfloat16) 
+                    lora_A = None if lora_A == None else lora_A.to(module.weight.dtype) 
+                    lora_B = None if lora_B == None else lora_B.to(module.weight.dtype) 
                     if loras_adapter_data == None:
                         alpha = lora_alphas.get(k[:-len("lora_X.weight")] + "alpha", 1.) 
                         loras_adapter_data = [lora_A, lora_B, alpha]
@@ -1320,9 +1367,13 @@ def load_model_data(model, file_path: str, do_quantize = False, quantizationType
             if k.endswith(missing_keys[0]):
                 base_model_prefix = k[:-len(missing_keys[0])]
                 break
-        state_dict = filter_state_dict(state_dict,base_model_prefix)
+        if base_model_prefix == None:
+            raise Exception("Missing keys: {missing_keys}")
+        state_dict = filter_state_dict(state_dict, base_model_prefix)
         missing_keys , unexpected_keys = model.load_state_dict(state_dict, False,  assign = True )
     del state_dict
+    if len(unexpected_keys) > 0 and verboseLevel >=2:
+        print(f"Unexpected keys while loading '{file_path}': {unexpected_keys}")
 
     for k,p in model.named_parameters():
         if p.is_meta:
@@ -1483,6 +1534,7 @@ class offload:
         self.transfer_stream = torch.cuda.Stream()
         self.async_transfers = False
         self.parameters_ref  = {} 
+        self.max_reservable_memory = 0
 
         global last_offload_obj
         last_offload_obj = self
@@ -1729,22 +1781,22 @@ class offload:
         gc.collect()
         self.last_reserved_mem_check = time.time()
 
-    def move_args_to_gpu(self, *args, **kwargs):
+    def move_args_to_gpu(self, dtype, *args, **kwargs):
         new_args= []
         new_kwargs={}
+
         for arg in args:
             if torch.is_tensor(arg):    
                 if arg.dtype == torch.float32:
-                    arg = arg.to(torch.bfloat16).cuda(non_blocking=True)             
+                    arg = arg.to(dtype).cuda(non_blocking=True)
                 elif not arg.is_cuda:
                     arg = arg.cuda(non_blocking=True)
             new_args.append(arg)
-
         for k in kwargs:
             arg = kwargs[k]
             if torch.is_tensor(arg):
                 if arg.dtype == torch.float32:
-                    arg = arg.to(torch.bfloat16).cuda(non_blocking=True)             
+                    arg = arg.to(dtype).cuda(non_blocking=True)             
                 elif not arg.is_cuda:
                     arg = arg.cuda(non_blocking=True)             
             new_kwargs[k]= arg
@@ -1805,7 +1857,7 @@ class offload:
         loras_scaling = model._loras_scaling
         training = False
 
-
+        dtype = weight.dtype 
         if weight.shape[-1] < x.shape[-2] : # sum base weight and lora matrices instead of applying input on each sub lora matrice if input is too large. This will save a lot VRAM and compute
             if len(active_adapters) > 0:
                 if isinstance(submodule, QModuleMixin): 
@@ -1833,7 +1885,7 @@ class offload:
             result = torch.nn.functional.linear(x, weight, bias=submodule.bias)
 
             if len(active_adapters) > 0:
-                x = x.to(torch.bfloat16)
+                x = x.to(dtype)
 
                 for active_adapter in active_adapters:
                     data = loras_data.get(active_adapter + '_GPU', None)
@@ -1873,7 +1925,7 @@ class offload:
         return functools.update_wrapper(functools.partial(lora_linear_forward, submodule), old_forward)
 
     def ensure_model_loaded(self, model_id):
-        if  model_id in self.active_models_ids:
+        if model_id in self.active_models_ids:
             return
         # new_model_id = getattr(module, "_mm_id") 
         # do not always unload existing models if it is more efficient to keep in them in the GPU 
@@ -1895,8 +1947,9 @@ class offload:
         target_module.register_forward_pre_hook(preload_blocks_for_compile)        
 
 
-    def hook_check_empty_cache_needed(self, target_module, model_id, blocks_name, previous_method,  context):
+    def hook_check_empty_cache_needed(self, target_module, model, model_id, blocks_name, previous_method,  context):
 
+        dtype = model._dtype
         qint4quantization =  isinstance(target_module, QModuleMixin) and  target_module.weight!= None and  target_module.weight.qtype == qint4 
         if qint4quantization:
             pass
@@ -1921,8 +1974,8 @@ class offload:
                     self.empty_cache_if_needed()
             elif blocks_name != self.loaded_blocks[model_id] and blocks_name not in self.preloaded_blocks_per_model[model_id]:
                 self.gpu_load_blocks(model_id, blocks_name)
-            if qint4quantization:
-                args, kwargs = self.move_args_to_gpu(*args, **kwargs)
+            if qint4quantization and dtype !=None:
+                args, kwargs = self.move_args_to_gpu(dtype, *args, **kwargs)
 
             return previous_method(*args, **kwargs) 
 
@@ -1933,11 +1986,16 @@ class offload:
 
         
     def hook_change_module(self, target_module, model, model_id, module_id, previous_method):
+        if hasattr(target_module, "_lock_dtype"):
+            dtype = target_module._lock_dtype 
+        else:
+            dtype = model._dtype
 
         def check_change_module(module, *args, **kwargs):            
             self.ensure_model_loaded(model_id)
             # transfer leftovers inputs that were incorrectly created in the RAM (mostly due to some .device tests that returned incorrectly "cpu")
-            args, kwargs = self.move_args_to_gpu(*args, **kwargs)
+            if dtype != None:
+                args, kwargs = self.move_args_to_gpu(dtype, *args, **kwargs)
             return previous_method(*args, **kwargs) 
   
         if hasattr(target_module, "_mm_id"):
@@ -2069,7 +2127,7 @@ class offload:
 
 
 
-def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, partialPinning = False, loras = None, quantizeTransformer = True,  extraModelsToQuantize = None, quantizationType = qint8, budgets= 0, workingVRAM = None, asyncTransfers = True, compile = False, convertFloatToBfloat16 = True, perc_reserved_mem_max = 0, coTenantsMap = None, verboseLevel = -1):
+def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, partialPinning = False, loras = None, quantizeTransformer = True,  extraModelsToQuantize = None, quantizationType = qint8, budgets= 0, workingVRAM = None, asyncTransfers = True, compile = False, convertWeightsFloatTo = torch.bfloat16, perc_reserved_mem_max = 0, coTenantsMap = None, verboseLevel = -1):
     """Hook to a pipeline or a group of modules in order to reduce their VRAM requirements:
     pipe_or_dict_of_modules : the pipeline object or a dictionary of modules of the model
     quantizeTransformer: set True by default will quantize on the fly the video / image model
@@ -2083,6 +2141,7 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
     self.verboseLevel = verboseLevel
     safetensors2.verboseLevel = verboseLevel
     self.modules_data = {}
+    
     model_budgets = {}
 
     windows_os =  os.name == 'nt'
@@ -2159,9 +2218,8 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
       #  torch._logging.set_logs(recompiles=True)
       #  torch._inductor.config.realize_opcount_threshold = 100 # workaround bug "AssertionError: increase TRITON_MAX_BLOCK['X'] to 4096."
 
-     
-    perc_reserved_mem_max = _get_perc_reserved_mem_max(perc_reserved_mem_max) 
-    max_reservable_memory = perc_reserved_mem_max * physical_memory
+    perc_reserved_mem_max = _get_perc_reserved_mem_max(perc_reserved_mem_max)
+    max_reservable_memory = _get_max_reservable_memory(perc_reserved_mem_max) 
 
     estimatesBytesToPin = 0
     for model_id in models: 
@@ -2176,30 +2234,39 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
         modelPinned = (pinAllModels or model_id in modelsToPin) and not hasattr(current_model,"_already_pinned")
 
         current_model_size = 0
-        for n, p in current_model.named_parameters():
-            p.requires_grad = False
-            if isinstance(p, QTensor):
-                # # fix quanto bug (seems to have been fixed)   
-                # if not modelPinned and p._scale.dtype == torch.float32:
-                #     p._scale = p._scale.to(torch.bfloat16) 
-                if p._qtype == qint4:
-                    if hasattr(p,"_scale_shift"):
-                        current_model_size +=  torch.numel(p._scale_shift) * p._scale_shift.element_size()
-                    else:
-                        current_model_size +=  torch.numel(p._scale) * p._shift.element_size() + torch.numel(p._scale) * p._shift.element_size()
+        model_dtype = None
+        
+        for _ , m in current_model.named_modules():
+            ignore_dtype = hasattr(m, "_lock_dtype")
+            for n, p in m.named_parameters(recurse = False):
+                p.requires_grad = False
+                if isinstance(p, QTensor):
+                    if p._qtype == qint4:
+                        if hasattr(p,"_scale_shift"):
+                            current_model_size +=  torch.numel(p._scale_shift) * p._scale_shift.element_size()
+                        else:
+                            current_model_size +=  torch.numel(p._scale) * p._shift.element_size() + torch.numel(p._scale) * p._shift.element_size()
 
-                    current_model_size +=  torch.numel(p._data._data) * p._data._data.element_size()
+                        current_model_size +=  torch.numel(p._data._data) * p._data._data.element_size()
 
-                else:
-                    current_model_size +=  torch.numel(p._scale) * p._scale.element_size()
-                    current_model_size +=  torch.numel(p._data) * p._data.element_size()
+                    else:
+                        current_model_size +=  torch.numel(p._scale) * p._scale.element_size()
+                        current_model_size +=  torch.numel(p._data) * p._data.element_size()
+                    dtype = p._scale.dtype
 
-            else:
-                if convertFloatToBfloat16 and p.data.dtype == torch.float32:
-                    # convert any left overs float32 weight to bloat16 to divide by 2 the model memory footprint
-                    p.data = p.data.to(torch.bfloat16)
-                current_model_size +=  torch.numel(p.data) * p.data.element_size()
-                            
+                else:
+                    if not ignore_dtype:
+                        dtype = p.data.dtype
+                        if convertWeightsFloatTo != None and  dtype == torch.float32 :
+                            # convert any left overs float32 weight to bfloat16 / float16 to divide by 2 the model memory footprint
+                            dtype = convertWeightsFloatTo if model_dtype == None else model_dtype
+                            p.data = p.data.to(dtype)
+                        if model_dtype== None:
+                            model_dtype = dtype
+                        else:
+                            assert model_dtype == dtype
+                    current_model_size +=  torch.numel(p.data) * p.data.element_size()
+                current_model._dtype = model_dtype
         for b in current_model.buffers():
             # do not convert 32 bits float to 16 bits since buffers are few (and potential gain low) and usually they are needed for precision calculation (for instance Rope)
             current_model_size +=  torch.numel(b.data) * b.data.element_size()
@@ -2267,7 +2334,7 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
                 if self.verboseLevel >=1:
                     print(f"Model '{model_id}' already pinned to reserved memory")
             else:
-                _pin_to_memory(current_model, model_id, partialPinning= partialPinning, pinnedPEFTLora = pinnedPEFTLora, verboseLevel=verboseLevel)            
+                _pin_to_memory(current_model, model_id, partialPinning= partialPinning, pinnedPEFTLora = pinnedPEFTLora, perc_reserved_mem_max = perc_reserved_mem_max, verboseLevel=verboseLevel)            
     
         current_budget = model_budgets[model_id]
         cur_blocks_prefix, prev_blocks_name, cur_blocks_name,cur_blocks_seq, is_mod_seq = None, None, None, -1, False
@@ -2317,7 +2384,7 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
                     elif compilationInThisOne and submodule in towers_modules: 
                         self.hook_preload_blocks_for_compilation(submodule, model_id, cur_blocks_name, context = submodule_name )
                     else:
-                        self.hook_check_empty_cache_needed(submodule, model_id, cur_blocks_name, submodule_method, context = submodule_name )
+                        self.hook_check_empty_cache_needed(submodule, current_model, model_id, cur_blocks_name, submodule_method, context = submodule_name )
 
                 self.add_module_to_blocks(model_id, cur_blocks_name, submodule, prev_blocks_name, submodule_name)
 

{mmgp-3.3.4.dist-info → mmgp-3.4.1.dist-info}/METADATA

@@ -1,197 +1,197 @@
-Metadata-Version: 2.4
-Name: mmgp
-Version: 3.3.4
-Summary: Memory Management for the GPU Poor
-Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
-License:                     GNU GENERAL PUBLIC LICENSE
-                               Version 3, 29 June 2007
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
-License-File: LICENSE.md
-Requires-Dist: torch>=2.1.0
-Requires-Dist: optimum-quanto
-Requires-Dist: accelerate
-Requires-Dist: safetensors
-Requires-Dist: psutil
-Dynamic: license-file
-
-
-<p align="center">
-  <H2>Memory Management 3.3.4  for the GPU Poor by DeepBeepMeep</H2>	
-</p>
-
-
-This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  can run smoothly on a 12 to 24 GB GPU limited card. 
-This a replacement for the accelerate library that should in theory manage offloading, but doesn't work properly with models that are loaded / unloaded several
-times in a pipe (eg VAE).
-
-Requirements:
-- VRAM: minimum 6 GB, recommended 24 GB (RTX 3090/ RTX 4090) 
-- RAM: minimum 24 GB, recommended 48 GB 
-
-This module features 5 profiles in order to able to run the model at a decent speed on a low end consumer config (24 GB of RAM and 6 VRAM) and to run it at a very good speed (if not the best) on a high end consumer config (48 GB of RAM and 24 GB of VRAM).\
-These RAM requirements are for Linux systems. Due to different memory management Windows will require an extra 16 GB of RAM to run the corresponding profile.
-
-Each profile may use a combination of the following: 
-- Low RAM consumption (thanks to a rewritten safetensors library) that allows low RAM on the fly quantization
-- Smart automated loading / unloading of models in the GPU to avoid unloading models that may be needed again soon
-- Smart slicing of models to reduce memory occupied by models in the VRAM
-- Ability to pin models to reserved RAM to accelerate transfers to VRAM
-- Async transfers to VRAM to avoid a pause when loading a new slice of a model
-- Automated on the fly quantization or ability to load pre quantized models
-- Pretrained Lora support with low RAM requirements
-- Support for pytorch compilation on Linux and WSL (supported on pure Windows but requires a complex Triton Installation).
-
-## Sample applications that use mmgp
-It is recommended to have a look at these applications to see how mmgp was implemented in each of them:
-- Wan2GP: https://github.com/deepbeepmeep/Wan2GP :\
-An excellent text to video and image to video generator by Alibaba
-
-- Hunyuan3D-2GP: https://github.com/deepbeepmeep/Hunyuan3D-2GP :\
-A great image to 3D and text to 3D tool by the Tencent team. Thanks to mmgp it can run with less than 6 GB of VRAM
-
-- HuanyuanVideoGP: https://github.com/deepbeepmeep/HunyuanVideoGP :\
-One of the best open source Text to Video generator
-
-- FluxFillGP: https://github.com/deepbeepmeep/FluxFillGP :\
-One of the best inpainting / outpainting tools based on Flux that can run with less than 12 GB of VRAM.
-
-- Cosmos1GP: https://github.com/deepbeepmeep/Cosmos1GP :\
-This application include two models: a text to world generator and a image / video to world (probably the best open source image to video generator).
-
-- OminiControlGP: https://github.com/deepbeepmeep/OminiControlGP :\
-A Flux derived application very powerful that can be used to transfer an object of your choice in a prompted scene. With mmgp you can run it with only 6 GB of VRAM.
-
-- YuE GP: https://github.com/deepbeepmeep/YuEGP :\
-A great song generator (instruments + singer's voice) based on prompted Lyrics and a genre description. Thanks to mmgp you can run it with less than 10 GB of VRAM without waiting forever.
-
-## Installation
-First you need to install the module in your current project with:
-```shell
-pip install mmgp
-```
-
-
-## Usage 
-
-It is almost plug and play and just needs to be invoked from the main app just after the model pipeline has been created.
-1) First make sure that the pipeline explictly loads the models in the CPU device, for instance:
-```
-  pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to("cpu")
-```
-
-2) Once every potential Lora has been loaded and merged, add the following lines for a quick setup:
-```
-  from mmgp import offload, profile_type
-  offload.profile(pipe, profile_type.HighRAM_LowVRAM_Fast)
-```
-
-You can choose between 5 profiles depending on your hardware:
-- HighRAM_HighVRAM  (1): at least 48 GB of RAM and 24 GB of VRAM : the fastest well suited for a RTX 3090 / RTX 4090 but consumes much more VRAM, adapted for fast shorter video or small batches of pictures
-- HighRAM_LowVRAM  (2): at least 48 GB of RAM and 12 GB of VRAM : a bit slower, better suited for RTX 3070/3080/4070/4080 or for RTX 3090 / RTX 4090 with large pictures batches or long videos
-- LowRAM_HighVRAM  (3): at least 32 GB of RAM and 24 GB of VRAM : adapted for RTX 3090 / RTX 4090 with limited RAM  but at the cost of VRAM (shorter videos / fewer images)
-- LowRAM_LowVRAM  (4): at least 32 GB of RAM and 12 GB of VRAM :  if you have little VRAM or want to generate longer videos / more images
-- VerylowRAM_LowVRAM  (5): at least 24 GB of RAM and 10 GB of VRAM : if you don't have much it won't be fast but maybe it will work
-
-Profile 2 (High RAM) and 4 (Low RAM) are the most recommended profiles since they are versatile (support for long videos for a slight performance cost).\
-If you use Flux derived applciation profile 1 and 3 will offer much faster generation times.
-In any case, a safe approach is to start from profile 5 (default profile) and then go down progressively to profile 4 and then to profile 2 as long as the app remains responsive or doesn't trigger any out of memory error.
-
-By default the model named 'transformer' will be quantized to 8 bits for all profiles. If you don't want that you may specify the optional parameter *quantizeTransformer = False*.
-
-Every parameter set automatically by a profile can be overridden with one or multiple parameters accepted by *offload.all* (see below):
-```
-  from mmgp import offload, profile_type
-  offload.profile(pipe, profile_type.HighRAM_LowVRAM, budgets = 1000)
-```
-If you want to know which parameter are set by one specific profile you can use the parameter *verboseLevel=2*
-
-**It is highly recommended to put the *from mmgp import offload, profile_type* at the top of your main python file (that is as the first import) so that all the existing safetensors calls are redirected to mmpg.**
- 
-
-## Alternatively you may want to create your own profile with specific parameters:
-
-For example:
-```
-  from mmgp import offload
-  offload.all(pipe, pinnedMemory=True, ExtraModelsToQuantize = ["text_encoder_2"] )
-```  
-- pinnedMemory: Boolean (for all models) or List of models ids to pin to RAM. Every model pinned to RAM will load much faster (up to 2 times) but this requires more RAM
-- quantizeTransformer: boolean by default True. The 'transformer' model in the pipe contains usually the video or image generator is by defaut; quantized on the fly by default to 8 bits. If you want to save time on disk and reduce the loading time, you may want to load directly a prequantized model. If you don't want to quantize the image generator, you need to set the option *quantizeTransformer* to *False* to turn off on the fly quantization.
-- extraModelsToQuantize: list of additional modelids of models to quantize on the fly. If the corresponding model is already quantized, this option will be ignored.
-- budgets: either a number in mega bytes,  (for all models, if 0 unlimited budget) a string that is perecentage of the total VRAM or a dictionary that maps model ids to mega bytes : define the approximate budget in mega bytes that is allocated in  VRAM for a model. Try not to allocate all the available VRAM so that the rest can be used to process the data. To define the default value in the dictionary, you may add entry named "*".
-The smaller this number, the more VRAM left for image data / longer video but also the slower because there will be lots of loading / unloading between the RAM and the VRAM. If model is too big to fit in a budget, it will be broken down in multiples parts that will be unloaded / loaded consequently. The speed of low budget can be  increased (up to 2 times) by turning on the options pinnedMemory and asyncTransfers.
-- workingVRAM: either a number in mega bytes, a string that is perecentage of the total VRAM or a dictionary that maps a model ids to a number in mega bytes that corresponds to a minimum amount of VRAM that should be left for the data processed by the model. This number will prevail if it is in conflict with a too high budget defined for the same model.
-- asyncTransfers: boolean, load to the GPU the next model part while the current part is being processed. This requires twice the budget if any is defined. This may increase speed by 20% (mostly visible on fast modern GPUs).
-- verboseLevel: number between 0 and 2 (1 by default), provides various level of feedback of the different processes
-- compile: list of model ids to compile, may accelerate up x2 depending on the type of GPU. It makes sense to compile only the model that is frequently used such as the "transformer" model in the case of video or image generation. Compilation requires Triton to be installed. Triton is available out of the box on Linux or WSL but requires to be installed with Windows: https://github.com/woct0rdho/triton-windows
-- coTenantsMap: a dictionary that maps a model id to a list of other models with which it accepts to share the VRAM at the same time. This is useful to avoid unefficient loading / unloading when two models processes are interleaved. For instance *coTenantsMap = { "text_encoder_2": ["text_encoder"] }* , here when *text_encoder_2* is loaded it won't unload *text_encoder*. Please note that the reverse is not true as these maps by design are not symetrical to allow tailored workflows. If you need to have as well *text_encoder* that won't unload *text_encoder_2* if it is already loaded *coTenantsMap = { "text_encoder_2": ["text_encoder"], "text_encoder": ["text_encoder_2"] }*
-
-If you are short on RAM and plan to work with quantized models, it is recommended to load pre-quantized models direclty rather than using on the fly quantization, it will be faster and consume slightly less RAM.
- 
-##  Going further
-
-The module includes several tools to package a light version of your favorite video / image generator:
-- *extract_models(string prefix,  obj to explore)*\
-This tool will try to detect for you models that are embedded in a pipeline or in some custom class. It will save you time by building a pipe dictionary required by *offload.all* or "offload.profile*. The prefix correponds to the text that will appear before the name of each model in the dictionary.
-
-- *load_loras_into_model(model, lora_path, lora_multi, activate_all_loras = True)*\
-Load in a model a list of Lora described by a list of path *lora_path* and a list of *weights coefficients*.
-The Lora file must be in the *diffusers* format. This function works also on non diffusers models. However if there is already an official Lora support for a model it is recommended to use the official diffusers functions. By default all the load loras will be activated or they can be activated later using *activate_loras*.
-
--*activate_loras(model, lora_nos, lora_multi = None )*\
-Activate the loras whose nos are in the list of nos. Every lora that is not this list and that was activated previously will be disactivated.
-
-- *save_model(model, file_path, do_quantize = False, quantizationType = qint8 )*\
-Save tensors of a model already loaded in memory in a safetensor format (much faster to reload). You can save it in a quantized format (default qint8 quantization recommended).
-The resulting safetensor file will contain extra fields in its metadata such as the quantization map and its configuration, so you will be able to move the file around without files such as *config.json* or *file_map.json*.
-You will need *load_model_data* or *fast_load_transformers_model* to read the file again . You may also load it using the default *safetensor* librar however you will need to provide in the same directory any complementary file that are usually requested (for instance *config.json*)
-
-- *load_model_data(model, file_path: str, do_quantize = False, quantizationType = qint8, pinToRAM = False, partialPin = False)*\
-Load the tensors data of a model in RAM of a model already initialized with no data. Detect and handle quantized models saved previously with *save_model*.A model can also be quantized on the fly while being loaded. The model which is loaded can be pinned to RAM while it is loaded, this is more RAM efficient than pinning tensors later using *offline.all* or *offline.profile*
-
-- *fast_load_transformers_model(model_path: str, do_quantize = False, quantizationType = qint8, pinToRAM = False, partialPin = False)*\
-Initialize (build the model hierarchy in memory) and fast load the corresponding tensors of a 'transformers' or 'diffusers' library model.
-The advantages over the original *from_pretrained* method is that a full model can fit into a single file with a filename of your choosing (thefore you can have multiple 'transformers' versions of the same model in the same directory) and prequantized models are processed in a transparent way. 
-Last but not least, you can also on the fly pin to RAM the whole model or the most important part of it (partialPin = True) in a more efficient way (faster and requires less RAM) than if you did through *offload.all* or *offload.profile*.
-
-
-The typical workflow wil be:
-1) temporarly insert the *save_model* function just after a model has been fully loaded to save a copy of the model / quantized model.
-2) replace the full initalizing / loading logic with *fast_load_transformers_model* (if there is a *from_pretrained* call to a transformers object) or only the tensor loading functions (*torch.load_model_file* and *torch.load_state_dict*) with *load_model_data after* the initializing logic.
-
-## Special cases
-Sometime there isn't an explicit pipe object as each submodel is loaded separately in the main app. If this is the case, you may try to use *extract_models* or create a dictionary that manually maps all the models.\
-For instance :
-
-
-- for flux derived models: 
-```
-pipe = { "text_encoder": clip, "text_encoder_2": t5, "transformer": model, "vae":ae }
-```
-- for mochi: 
-```
-pipe = { "text_encoder": self.text_encoder, "transformer": self.dit, "vae":self.decoder }
-```
-
-
-Please note it is recommended to have always one model whose Id is 'transformer' so that you can leverage predefined profiles. The 'transformer' corresponds to the main image / video model which usually needs to be quantized (this is done on the fly by default when loading the model).
-
-Be careful, lots of models use the T5 XXL as a text encoder. However, quite often their corresponding pipeline configurations point at the official Google T5 XXL repository 
-where there is a huge 40GB model to download and load. It is cumbersorme as it is a 32 bits model and contains the decoder part of T5 that is not used. 
-I suggest you use instead one of the 16 bits encoder only version available around, for instance:
-```
-text_encoder_2 = T5EncoderModel.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="text_encoder_2", torch_dtype=torch.float16)
-```
-
-Sometime just providing the pipe won't be sufficient as you will need to change the content of the core model: 
-- For instance you may need to disable an existing CPU offload logic that already exists (such as manual calls to move tensors between cuda and the cpu)
-- mmpg to tries to fake the device as being "cuda" but sometimes some code won't be fooled and it will create tensors in the cpu device and this may cause some issues.
-
-You are free to use my module for non commercial use as long you give me proper credits. You may contact me on twitter @deepbeepmeep
-
-Thanks to
----------
-- Huggingface / accelerate for the hooking examples
-- Huggingface / quanto for their very useful quantizer
-- gau-nernst for his Pinnig RAM samples
+Metadata-Version: 2.4
+Name: mmgp
+Version: 3.4.1
+Summary: Memory Management for the GPU Poor
+Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
+License:                     GNU GENERAL PUBLIC LICENSE
+                               Version 3, 29 June 2007
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: torch>=2.1.0
+Requires-Dist: optimum-quanto
+Requires-Dist: accelerate
+Requires-Dist: safetensors
+Requires-Dist: psutil
+Dynamic: license-file
+
+
+<p align="center">
+  <H2>Memory Management 3.4.1 for the GPU Poor by DeepBeepMeep</H2>	
+</p>
+
+
+This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  can run smoothly on a 12 to 24 GB GPU limited card. 
+This a replacement for the accelerate library that should in theory manage offloading, but doesn't work properly with models that are loaded / unloaded several
+times in a pipe (eg VAE).
+
+Requirements:
+- VRAM: minimum 6 GB, recommended 24 GB (RTX 3090/ RTX 4090) 
+- RAM: minimum 24 GB, recommended 48 GB 
+
+This module features 5 profiles in order to able to run the model at a decent speed on a low end consumer config (24 GB of RAM and 6 VRAM) and to run it at a very good speed (if not the best) on a high end consumer config (48 GB of RAM and 24 GB of VRAM).\
+These RAM requirements are for Linux systems. Due to different memory management Windows will require an extra 16 GB of RAM to run the corresponding profile.
+
+Each profile may use a combination of the following: 
+- Low RAM consumption (thanks to a rewritten safetensors library) that allows low RAM on the fly quantization
+- Smart automated loading / unloading of models in the GPU to avoid unloading models that may be needed again soon
+- Smart slicing of models to reduce memory occupied by models in the VRAM
+- Ability to pin models to reserved RAM to accelerate transfers to VRAM
+- Async transfers to VRAM to avoid a pause when loading a new slice of a model
+- Automated on the fly quantization or ability to load pre quantized models
+- Pretrained Lora support with low RAM requirements
+- Support for pytorch compilation on Linux and WSL (supported on pure Windows but requires a complex Triton Installation).
+
+## Sample applications that use mmgp
+It is recommended to have a look at these applications to see how mmgp was implemented in each of them:
+- Wan2GP: https://github.com/deepbeepmeep/Wan2GP :\
+An excellent text to video and image to video generator by Alibaba
+
+- Hunyuan3D-2GP: https://github.com/deepbeepmeep/Hunyuan3D-2GP :\
+A great image to 3D and text to 3D tool by the Tencent team. Thanks to mmgp it can run with less than 6 GB of VRAM
+
+- HuanyuanVideoGP: https://github.com/deepbeepmeep/HunyuanVideoGP :\
+One of the best open source Text to Video generator
+
+- FluxFillGP: https://github.com/deepbeepmeep/FluxFillGP :\
+One of the best inpainting / outpainting tools based on Flux that can run with less than 12 GB of VRAM.
+
+- Cosmos1GP: https://github.com/deepbeepmeep/Cosmos1GP :\
+This application include two models: a text to world generator and a image / video to world (probably the best open source image to video generator).
+
+- OminiControlGP: https://github.com/deepbeepmeep/OminiControlGP :\
+A Flux derived application very powerful that can be used to transfer an object of your choice in a prompted scene. With mmgp you can run it with only 6 GB of VRAM.
+
+- YuE GP: https://github.com/deepbeepmeep/YuEGP :\
+A great song generator (instruments + singer's voice) based on prompted Lyrics and a genre description. Thanks to mmgp you can run it with less than 10 GB of VRAM without waiting forever.
+
+## Installation
+First you need to install the module in your current project with:
+```shell
+pip install mmgp
+```
+
+
+## Usage 
+
+It is almost plug and play and just needs to be invoked from the main app just after the model pipeline has been created.
+1) First make sure that the pipeline explictly loads the models in the CPU device, for instance:
+```
+  pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to("cpu")
+```
+
+2) Once every potential Lora has been loaded and merged, add the following lines for a quick setup:
+```
+  from mmgp import offload, profile_type
+  offload.profile(pipe, profile_type.HighRAM_LowVRAM_Fast)
+```
+
+You can choose between 5 profiles depending on your hardware:
+- HighRAM_HighVRAM  (1): at least 48 GB of RAM and 24 GB of VRAM : the fastest well suited for a RTX 3090 / RTX 4090 but consumes much more VRAM, adapted for fast shorter video or small batches of pictures
+- HighRAM_LowVRAM  (2): at least 48 GB of RAM and 12 GB of VRAM : a bit slower, better suited for RTX 3070/3080/4070/4080 or for RTX 3090 / RTX 4090 with large pictures batches or long videos
+- LowRAM_HighVRAM  (3): at least 32 GB of RAM and 24 GB of VRAM : adapted for RTX 3090 / RTX 4090 with limited RAM  but at the cost of VRAM (shorter videos / fewer images)
+- LowRAM_LowVRAM  (4): at least 32 GB of RAM and 12 GB of VRAM :  if you have little VRAM or want to generate longer videos / more images
+- VerylowRAM_LowVRAM  (5): at least 24 GB of RAM and 10 GB of VRAM : if you don't have much it won't be fast but maybe it will work
+
+Profile 2 (High RAM) and 4 (Low RAM) are the most recommended profiles since they are versatile (support for long videos for a slight performance cost).\
+If you use Flux derived applciation profile 1 and 3 will offer much faster generation times.
+In any case, a safe approach is to start from profile 5 (default profile) and then go down progressively to profile 4 and then to profile 2 as long as the app remains responsive or doesn't trigger any out of memory error.
+
+By default the model named 'transformer' will be quantized to 8 bits for all profiles. If you don't want that you may specify the optional parameter *quantizeTransformer = False*.
+
+Every parameter set automatically by a profile can be overridden with one or multiple parameters accepted by *offload.all* (see below):
+```
+  from mmgp import offload, profile_type
+  offload.profile(pipe, profile_type.HighRAM_LowVRAM, budgets = 1000)
+```
+If you want to know which parameter are set by one specific profile you can use the parameter *verboseLevel=2*
+
+**It is highly recommended to put the *from mmgp import offload, profile_type* at the top of your main python file (that is as the first import) so that all the existing safetensors calls are redirected to mmpg.**
+ 
+
+## Alternatively you may want to create your own profile with specific parameters:
+
+For example:
+```
+  from mmgp import offload
+  offload.all(pipe, pinnedMemory=True, ExtraModelsToQuantize = ["text_encoder_2"] )
+```  
+- pinnedMemory: Boolean (for all models) or List of models ids to pin to RAM. Every model pinned to RAM will load much faster (up to 2 times) but this requires more RAM
+- quantizeTransformer: boolean by default True. The 'transformer' model in the pipe contains usually the video or image generator is by defaut; quantized on the fly by default to 8 bits. If you want to save time on disk and reduce the loading time, you may want to load directly a prequantized model. If you don't want to quantize the image generator, you need to set the option *quantizeTransformer* to *False* to turn off on the fly quantization.
+- extraModelsToQuantize: list of additional modelids of models to quantize on the fly. If the corresponding model is already quantized, this option will be ignored.
+- budgets: either a number in mega bytes,  (for all models, if 0 unlimited budget) a string that is perecentage of the total VRAM or a dictionary that maps model ids to mega bytes : define the approximate budget in mega bytes that is allocated in  VRAM for a model. Try not to allocate all the available VRAM so that the rest can be used to process the data. To define the default value in the dictionary, you may add entry named "*".
+The smaller this number, the more VRAM left for image data / longer video but also the slower because there will be lots of loading / unloading between the RAM and the VRAM. If model is too big to fit in a budget, it will be broken down in multiples parts that will be unloaded / loaded consequently. The speed of low budget can be  increased (up to 2 times) by turning on the options pinnedMemory and asyncTransfers.
+- workingVRAM: either a number in mega bytes, a string that is perecentage of the total VRAM or a dictionary that maps a model ids to a number in mega bytes that corresponds to a minimum amount of VRAM that should be left for the data processed by the model. This number will prevail if it is in conflict with a too high budget defined for the same model.
+- asyncTransfers: boolean, load to the GPU the next model part while the current part is being processed. This requires twice the budget if any is defined. This may increase speed by 20% (mostly visible on fast modern GPUs).
+- verboseLevel: number between 0 and 2 (1 by default), provides various level of feedback of the different processes
+- compile: list of model ids to compile, may accelerate up x2 depending on the type of GPU. It makes sense to compile only the model that is frequently used such as the "transformer" model in the case of video or image generation. Compilation requires Triton to be installed. Triton is available out of the box on Linux or WSL but requires to be installed with Windows: https://github.com/woct0rdho/triton-windows
+- coTenantsMap: a dictionary that maps a model id to a list of other models with which it accepts to share the VRAM at the same time. This is useful to avoid unefficient loading / unloading when two models processes are interleaved. For instance *coTenantsMap = { "text_encoder_2": ["text_encoder"] }* , here when *text_encoder_2* is loaded it won't unload *text_encoder*. Please note that the reverse is not true as these maps by design are not symetrical to allow tailored workflows. If you need to have as well *text_encoder* that won't unload *text_encoder_2* if it is already loaded *coTenantsMap = { "text_encoder_2": ["text_encoder"], "text_encoder": ["text_encoder_2"] }*
+
+If you are short on RAM and plan to work with quantized models, it is recommended to load pre-quantized models direclty rather than using on the fly quantization, it will be faster and consume slightly less RAM.
+ 
+##  Going further
+
+The module includes several tools to package a light version of your favorite video / image generator:
+- *extract_models(string prefix,  obj to explore)*\
+This tool will try to detect for you models that are embedded in a pipeline or in some custom class. It will save you time by building a pipe dictionary required by *offload.all* or "offload.profile*. The prefix correponds to the text that will appear before the name of each model in the dictionary.
+
+- *load_loras_into_model(model, lora_path, lora_multi, activate_all_loras = True)*\
+Load in a model a list of Lora described by a list of path *lora_path* and a list of *weights coefficients*.
+The Lora file must be in the *diffusers* format. This function works also on non diffusers models. However if there is already an official Lora support for a model it is recommended to use the official diffusers functions. By default all the load loras will be activated or they can be activated later using *activate_loras*.
+
+-*activate_loras(model, lora_nos, lora_multi = None )*\
+Activate the loras whose nos are in the list of nos. Every lora that is not this list and that was activated previously will be disactivated.
+
+- *save_model(model, file_path, do_quantize = False, quantizationType = qint8 )*\
+Save tensors of a model already loaded in memory in a safetensor format (much faster to reload). You can save it in a quantized format (default qint8 quantization recommended).
+The resulting safetensor file will contain extra fields in its metadata such as the quantization map and its configuration, so you will be able to move the file around without files such as *config.json* or *file_map.json*.
+You will need *load_model_data* or *fast_load_transformers_model* to read the file again . You may also load it using the default *safetensor* librar however you will need to provide in the same directory any complementary file that are usually requested (for instance *config.json*)
+
+- *load_model_data(model, file_path: str, do_quantize = False, quantizationType = qint8, pinToRAM = False, partialPin = False)*\
+Load the tensors data of a model in RAM of a model already initialized with no data. Detect and handle quantized models saved previously with *save_model*.A model can also be quantized on the fly while being loaded. The model which is loaded can be pinned to RAM while it is loaded, this is more RAM efficient than pinning tensors later using *offline.all* or *offline.profile*
+
+- *fast_load_transformers_model(model_path: str, do_quantize = False, quantizationType = qint8, pinToRAM = False, partialPin = False)*\
+Initialize (build the model hierarchy in memory) and fast load the corresponding tensors of a 'transformers' or 'diffusers' library model.
+The advantages over the original *from_pretrained* method is that a full model can fit into a single file with a filename of your choosing (thefore you can have multiple 'transformers' versions of the same model in the same directory) and prequantized models are processed in a transparent way. 
+Last but not least, you can also on the fly pin to RAM the whole model or the most important part of it (partialPin = True) in a more efficient way (faster and requires less RAM) than if you did through *offload.all* or *offload.profile*.
+
+
+The typical workflow wil be:
+1) temporarly insert the *save_model* function just after a model has been fully loaded to save a copy of the model / quantized model.
+2) replace the full initalizing / loading logic with *fast_load_transformers_model* (if there is a *from_pretrained* call to a transformers object) or only the tensor loading functions (*torch.load_model_file* and *torch.load_state_dict*) with *load_model_data after* the initializing logic.
+
+## Special cases
+Sometime there isn't an explicit pipe object as each submodel is loaded separately in the main app. If this is the case, you may try to use *extract_models* or create a dictionary that manually maps all the models.\
+For instance :
+
+
+- for flux derived models: 
+```
+pipe = { "text_encoder": clip, "text_encoder_2": t5, "transformer": model, "vae":ae }
+```
+- for mochi: 
+```
+pipe = { "text_encoder": self.text_encoder, "transformer": self.dit, "vae":self.decoder }
+```
+
+
+Please note it is recommended to have always one model whose Id is 'transformer' so that you can leverage predefined profiles. The 'transformer' corresponds to the main image / video model which usually needs to be quantized (this is done on the fly by default when loading the model).
+
+Be careful, lots of models use the T5 XXL as a text encoder. However, quite often their corresponding pipeline configurations point at the official Google T5 XXL repository 
+where there is a huge 40GB model to download and load. It is cumbersorme as it is a 32 bits model and contains the decoder part of T5 that is not used. 
+I suggest you use instead one of the 16 bits encoder only version available around, for instance:
+```
+text_encoder_2 = T5EncoderModel.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="text_encoder_2", torch_dtype=torch.float16)
+```
+
+Sometime just providing the pipe won't be sufficient as you will need to change the content of the core model: 
+- For instance you may need to disable an existing CPU offload logic that already exists (such as manual calls to move tensors between cuda and the cpu)
+- mmpg to tries to fake the device as being "cuda" but sometimes some code won't be fooled and it will create tensors in the cpu device and this may cause some issues.
+
+You are free to use my module for non commercial use as long you give me proper credits. You may contact me on twitter @deepbeepmeep
+
+Thanks to
+---------
+- Huggingface / accelerate for the hooking examples
+- Huggingface / quanto for their very useful quantizer
+- gau-nernst for his Pinnig RAM samples

mmgp-3.4.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mmgp/__init__.py,sha256=A9qBwyQMd1M7vshSTOBnFGP1MQvS2hXmTcTCMUcmyzE,509
+mmgp/offload.py,sha256=O2gCDhTM6qtYH10QFkYcd1eI2CBaALhGrhYejA5_OII,110019
+mmgp/safetensors2.py,sha256=rmUbBmK3Dra5prUTTRSVi6-XUFAa9Mj6B5CNPgzt9To,17333
+mmgp-3.4.1.dist-info/licenses/LICENSE.md,sha256=DD-WIS0BkPoWJ_8hQO3J8hMP9K_1-dyrYv1YCbkxcDU,94
+mmgp-3.4.1.dist-info/METADATA,sha256=JGsaE1Xlkb-a68ivHJPaf_wAW2RexkCAru3N_TjUNYc,16350
+mmgp-3.4.1.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+mmgp-3.4.1.dist-info/top_level.txt,sha256=waGaepj2qVfnS2yAOkaMu4r9mJaVjGbEi6AwOUogU_U,14
+mmgp-3.4.1.dist-info/RECORD,,

{mmgp-3.3.4.dist-info → mmgp-3.4.1.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.0.2)
+Generator: setuptools (80.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

{mmgp-3.3.4.dist-info → mmgp-3.4.1.dist-info}/licenses/LICENSE.md

@@ -1,2 +1,2 @@
-                    GNU GENERAL PUBLIC LICENSE
+                    GNU GENERAL PUBLIC LICENSE
                        Version 3, 29 June 2007

mmgp-3.3.4.dist-info/RECORD

@@ -1,9 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mmgp/__init__.py,sha256=A9qBwyQMd1M7vshSTOBnFGP1MQvS2hXmTcTCMUcmyzE,509
-mmgp/offload.py,sha256=WpQK1af2g0qcAm32EguTX8oBHZGKumPX2EqYS-df69Y,106583
-mmgp/safetensors2.py,sha256=rmUbBmK3Dra5prUTTRSVi6-XUFAa9Mj6B5CNPgzt9To,17333
-mmgp-3.3.4.dist-info/licenses/LICENSE.md,sha256=HjzvY2grdtdduZclbZ46B2M-XpT4MDCxFub5ZwTWq2g,93
-mmgp-3.3.4.dist-info/METADATA,sha256=Yk2eSpNITRDHK0lclsP6VXhW0_5hkUNVvXSfk25f7Ds,16154
-mmgp-3.3.4.dist-info/WHEEL,sha256=DK49LOLCYiurdXXOXwGJm6U4DkHkg4lcxjhqwRa0CP4,91
-mmgp-3.3.4.dist-info/top_level.txt,sha256=waGaepj2qVfnS2yAOkaMu4r9mJaVjGbEi6AwOUogU_U,14
-mmgp-3.3.4.dist-info/RECORD,,