mmgp 3.0.3__py3-none-any.whl → 3.1.0__py3-none-any.whl

mmgp/offload.py

@@ -1,4 +1,4 @@
-# ------------------ Memory Management 3.0 for the GPU Poor by DeepBeepMeep (mmgp)------------------
+# ------------------ Memory Management 3.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
@@ -61,13 +61,25 @@ import sys
 import os
 import json
 import psutil
+try:    
+    from diffusers.utils.peft_utils import set_weights_and_activate_adapters, get_peft_kwargs
+except:
+    set_weights_and_activate_adapters = None
+    get_peft_kwargs = None
+    pass
+try:    
+    from peft import LoraConfig, inject_adapter_in_model, set_peft_model_state_dict
+except:
+    inject_adapter_in_model = None
+    pass
+
 from mmgp import safetensors2
 from mmgp import profile_type
 
-from optimum.quanto import freeze, qfloat8, qint8, quantize, QModuleMixin, QTensor, WeightQBytesTensor, quantize_module 
-
+from optimum.quanto import freeze,  qfloat8, qint4 , qint8, quantize, QModuleMixin, QTensor,  quantize_module 
 
 
+shared_state = {}
 
 mmm = safetensors2.mmm
 
@@ -127,6 +139,9 @@ def move_tensors(obj, device):
         return _list
     else:
         raise TypeError("Tensor or list / dict of tensors expected")
+def _get_module_name(v):
+    return v.__module__.lower()
+
 
 def _compute_verbose_level(level):
     if level <0:        
@@ -139,33 +154,75 @@ def _get_max_reservable_memory(perc_reserved_mem_max):
         perc_reserved_mem_max = 0.40 if os.name == 'nt' else 0.5        
     return  perc_reserved_mem_max * physical_memory
 
-def _detect_main_towers(model, verboseLevel=1):
+def _detect_main_towers(model, min_floors = 5, verboseLevel=1):
     cur_blocks_prefix = None
     towers_modules= []
     towers_names= []
 
+    floors_modules= []
+    tower_name = None
+
+
     for submodule_name, submodule in model.named_modules():  
+
         if submodule_name=='':
             continue
 
-        if isinstance(submodule, torch.nn.ModuleList):
-            newList =False
-            if cur_blocks_prefix == None:
-                cur_blocks_prefix = submodule_name + "."
-                newList = True
-            else:
-                if not submodule_name.startswith(cur_blocks_prefix):
-                    cur_blocks_prefix = submodule_name + "."
-                    newList = True
+        if cur_blocks_prefix != None:
+            if submodule_name.startswith(cur_blocks_prefix):
+                depth_prefix = cur_blocks_prefix.split(".")
+                depth_name = submodule_name.split(".")
+                level  =  depth_name[len(depth_prefix)-1]                        
+                pre , num = _extract_num_from_str(level)
+
+                if num != cur_blocks_seq: 
+                    floors_modules.append(submodule)
 
-            if newList and len(submodule)>=5:
-                towers_names.append(submodule_name)
-                towers_modules.append(submodule)
+                cur_blocks_seq = num
+            else:
+                if len(floors_modules) >= min_floors:
+                    towers_modules += floors_modules
+                    towers_names.append(tower_name)
+                tower_name = None
+                floors_modules= []
+                cur_blocks_prefix, cur_blocks_seq = None, -1
+
+        if cur_blocks_prefix == None:
+            pre , num = _extract_num_from_str(submodule_name)
+            if isinstance(submodule, (torch.nn.ModuleList)):  
+                cur_blocks_prefix, cur_blocks_seq = pre + ".",  -1
+                tower_name = submodule_name + ".*" 
+            elif num >=0:
+                cur_blocks_prefix, cur_blocks_seq = pre, num
+                tower_name = submodule_name[ :-1] + "*" 
+                floors_modules.append(submodule)
+
+    if len(floors_modules) >= min_floors:
+        towers_modules += floors_modules
+        towers_names.append(tower_name)
+
+    # for submodule_name, submodule in model.named_modules():  
+    #     if submodule_name=='':
+    #         continue
+
+    #     if isinstance(submodule, torch.nn.ModuleList):
+    #         newList =False
+    #         if cur_blocks_prefix == None:
+    #             cur_blocks_prefix = submodule_name + "."
+    #             newList = True
+    #         else:
+    #             if not submodule_name.startswith(cur_blocks_prefix):
+    #                 cur_blocks_prefix = submodule_name + "."
+    #                 newList = True
+
+    #         if newList and len(submodule)>=5:
+    #             towers_names.append(submodule_name)
+    #             towers_modules.append(submodule)
                 
-        else:                
-            if cur_blocks_prefix is not None:
-                if not submodule_name.startswith(cur_blocks_prefix):
-                    cur_blocks_prefix = None 
+    #     else:                
+    #         if cur_blocks_prefix is not None:
+    #             if not submodule_name.startswith(cur_blocks_prefix):
+    #                 cur_blocks_prefix = None 
 
     return towers_names, towers_modules
 
@@ -179,7 +236,7 @@ def _get_model(model_path):
     _path = Path(model_path).parts
     _filename = _path[-1]
     _path = _path[:-1]
-    if len(_path)==1:
+    if len(_path)<=1:
         raise("file not found")
     else:
         from huggingface_hub import  hf_hub_download #snapshot_download,    
@@ -263,7 +320,13 @@ def _pin_to_memory(model, model_id, partialPinning = False, perc_reserved_mem_ma
     # print(f"num params to pin {model_id}: {len(params_list)}")
     for p in params_list:
         if isinstance(p, QTensor):
-            length = torch.numel(p._data) * p._data.element_size() + torch.numel(p._scale) * p._scale.element_size() 
+            if p._qtype == qint4:
+                if hasattr(p,"_scale_shift"):
+                    length = torch.numel(p._data._data) * p._data._data.element_size() + torch.numel(p._scale_shift) * p._scale_shift.element_size() 
+                else:
+                    length = torch.numel(p._data._data) * p._data._data.element_size() + torch.numel(p._scale) * p._scale.element_size() + torch.numel(p._shift) * p._shift.element_size()                     
+            else:
+                length = torch.numel(p._data) * p._data.element_size() + torch.numel(p._scale) * p._scale.element_size() 
         else:
             length = torch.numel(p.data) * p.data.element_size() 
 
@@ -306,10 +369,22 @@ def _pin_to_memory(model, model_id, partialPinning = False, perc_reserved_mem_ma
         if big_tensor_no>=0 and big_tensor_no < last_big_tensor:
             current_big_tensor = big_tensors[big_tensor_no]
             if isinstance(p, QTensor):
-                length1 = torch.numel(p._data) * p._data.element_size() 
-                p._data = _move_to_pinned_tensor(p._data, current_big_tensor, offset, length1)
-                length2 = torch.numel(p._scale) * p._scale.element_size() 
-                p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
+                if p._qtype == qint4:
+                    length1 = torch.numel(p._data._data) * p._data._data.element_size()
+                    p._data._data =  _move_to_pinned_tensor(p._data._data, current_big_tensor, offset, length1)
+                    if hasattr(p,"_scale_shift"):
+                        length2 = torch.numel(p._scale_shift) * p._scale_shift.element_size() 
+                        p._scale_shift = _move_to_pinned_tensor(p._scale_shift, current_big_tensor, offset + length1, length2)
+                    else:
+                        length2 = torch.numel(p._scale) * p._scale.element_size() 
+                        p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
+                        length3 = torch.numel(p._shift) * p._shift.element_size() 
+                        p._shift = _move_to_pinned_tensor(p._shift, current_big_tensor, offset + length1 + length2, length3)
+                else:
+                    length1 = torch.numel(p._data) * p._data.element_size() 
+                    p._data = _move_to_pinned_tensor(p._data, current_big_tensor, offset, length1)
+                    length2 = torch.numel(p._scale) * p._scale.element_size() 
+                    p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
             else:
                 length = torch.numel(p.data) * p.data.element_size() 
                 p.data = _move_to_pinned_tensor(p.data, current_big_tensor, offset, length)
@@ -336,100 +411,16 @@ def _welcome():
     if welcome_displayed:
          return 
     welcome_displayed = True
-    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.0) by DeepBeepMeep ************{ENDC}{UNBOLD}")
-
-
-# def _pin_to_memory_sd(model, sd, model_id, partialPinning = False, perc_reserved_mem_max = 0, verboseLevel = 1):
-#     if  verboseLevel>=1 :
-#         if partialPinning:
-#             print(f"Partial pinning to reserved RAM of data of file '{model_id}' while loading it")
-#         else:            
-#             print(f"Pinning data to reserved RAM of file '{model_id}' while loading it")
-
-#     max_reservable_memory = _get_max_reservable_memory(perc_reserved_mem_max)
-#     if partialPinning:
-#         towers_names, _ = _detect_main_towers(model)
-#         towers_names = [n +"." for n in towers_names]
-
-#     BIG_TENSOR_MAX_SIZE = 2**28 # 256 MB
-#     current_big_tensor_size = 0
-#     big_tensor_no  = 0
-#     big_tensors_sizes = []
-#     tensor_map_indexes = []
-#     total_tensor_bytes = 0
-
-#     for k,t in  sd.items():
-#         include = True
-#         # if isinstance(p, QTensor):
-#         #     length = torch.numel(p._data) * p._data.element_size() + torch.numel(p._scale) * p._scale.element_size() 
-#         # else:
-#         #     length = torch.numel(p.data) * p.data.element_size() 
-#         length = torch.numel(t) * t.data.element_size() 
-
-#         if partialPinning:
-#             include = any(k.startswith(pre) for pre in towers_names) if partialPinning else True
-
-#         if include:
-#             if current_big_tensor_size + length > BIG_TENSOR_MAX_SIZE:
-#                 big_tensors_sizes.append(current_big_tensor_size)
-#                 current_big_tensor_size = 0
-#                 big_tensor_no += 1
-#             tensor_map_indexes.append((big_tensor_no, current_big_tensor_size, length  ))
-#             current_big_tensor_size += length
-#         else:
-#             tensor_map_indexes.append((-1, 0, 0  ))
-#         total_tensor_bytes += length
-
-#     big_tensors_sizes.append(current_big_tensor_size)
-
-#     big_tensors = []
-#     last_big_tensor = 0
-#     total = 0
-    
+    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.1) by DeepBeepMeep ************{ENDC}{UNBOLD}")
 
-#     for size in big_tensors_sizes:
-#         try:
-#             currrent_big_tensor = torch.empty( size, dtype= torch.uint8, pin_memory=True)
-#             big_tensors.append(currrent_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})")
-#             break
-
-#         last_big_tensor += 1
-#         total += size
-
-
-#     tensor_no = 0
-#     for k,t in sd.items():
-#         big_tensor_no, offset, length = tensor_map_indexes[tensor_no]
-#         if big_tensor_no>=0 and big_tensor_no < last_big_tensor:
-#             current_big_tensor = big_tensors[big_tensor_no]
-#             # if isinstance(p, QTensor):
-#             #     length1 = torch.numel(p._data) * p._data.element_size() 
-#             #     p._data = _move_to_pinned_tensor(p._data, current_big_tensor, offset, length1)
-#             #     length2 = torch.numel(p._scale) * p._scale.element_size() 
-#             #     p._scale = _move_to_pinned_tensor(p._scale, current_big_tensor, offset + length1, length2)
-#             # else:
-#             #     length = torch.numel(p.data) * p.data.element_size() 
-#             #     p.data = _move_to_pinned_tensor(p.data, current_big_tensor, offset, length)
-#             length = torch.numel(t) * t.data.element_size() 
-#             t = _move_to_pinned_tensor(t, current_big_tensor, offset, length)
-#             sd[k] = t
-#         tensor_no += 1
-
-#     global total_pinned_bytes
-#     total_pinned_bytes += total
-
-#     if verboseLevel >=1:
-#         if total_tensor_bytes == total:        
-#             print(f"The whole model was pinned to reserved RAM: {last_big_tensor} large blocks spread across {total/ONE_MB:.2f} MB")
-#         else:
-#             print(f"{total/ONE_MB:.2f} MB were pinned to reserved RAM out of {total_tensor_bytes/ONE_MB:.2f} MB")
-
-#     model._already_pinned = True
-
-
-#     return 
+def _extract_num_from_str(num_in_str):
+    for i in range(len(num_in_str)):
+        if not num_in_str[-i-1:].isnumeric():
+            if i == 0:
+                return num_in_str, -1
+            else:             
+                return num_in_str[: -i],  int(num_in_str[-i:])                    
+    return  "", int(num_in_str)
 
 def  _quantize_dirty_hack(model):
     # dirty hack: add a hook on state_dict() to return a fake non quantized state_dict if called by Lora Diffusers initialization functions
@@ -536,10 +527,14 @@ def _quantize(model_to_quantize, weights=qint8, verboseLevel = 1, threshold = 10
     prev_blocks_prefix = None
 
     if hasattr(model_to_quantize, "_quanto_map"):
+        for k, entry in model_to_quantize._quanto_map.items():
+            weights  =  entry["weights"]
+            print(f"Model '{model_id}' is already quantized to format '{weights}'")
+            return False
         print(f"Model '{model_id}' is already quantized")
         return False
-    
-    print(f"Quantization of model '{model_id}' started")
+
+    print(f"Quantization of model '{model_id}' started to format '{weights}'")
 
     for submodule_name, submodule in model_to_quantize.named_modules():  
         if isinstance(submodule, QModuleMixin):
@@ -594,18 +589,18 @@ def _quantize(model_to_quantize, weights=qint8, verboseLevel = 1, threshold = 10
     if verboseLevel >=2:
         print(f"Total Excluded {total_excluded/ONE_MB:.1f} MB oF {total_size/ONE_MB:.1f} that is {perc_excluded*100:.2f}%")
     if perc_excluded >= 0.10:
-        print(f"Too many many modules are excluded, there is something wrong with the selection, switch back to full quantization.")
+        print(f"Too many modules are excluded, there is something wrong with the selection, switch back to full quantization.")
         exclude_list = None
 
 
     #quantize(model_to_quantize,weights, exclude= exclude_list)
-    pass
+
     for name, m in model_to_quantize.named_modules():
         if exclude_list is None or not any( name == module_name for module_name in exclude_list):
             _quantize_submodule(model_to_quantize, name, m, weights=weights, activations=None, optimizer=None)
 
-    # force read non quantized parameters so that their lazy tensors and corresponding mmap are released
-    # otherwise we may end up to keep in memory both the quantized and the non quantize model
+    # force to read non quantized parameters so that their lazy tensors and corresponding mmap are released
+    # otherwise we may end up keeping in memory both the quantized and the non quantize model
     for m in model_to_quantize.modules():
         # do not read quantized weights (detected them directly or behind an adapter)
         if isinstance(m, QModuleMixin) or hasattr(m, "base_layer") and  isinstance(m.base_layer, QModuleMixin): 
@@ -620,18 +615,271 @@ def _quantize(model_to_quantize, weights=qint8, verboseLevel = 1, threshold = 10
             b.data = b.data + 0 
 
     
+
     freeze(model_to_quantize)
     torch.cuda.empty_cache()
     gc.collect()         
     quantization_map = _quantization_map(model_to_quantize)
     model_to_quantize._quanto_map = quantization_map
 
+    if hasattr(model_to_quantize, "_already_pinned"):
+        delattr(model_to_quantize, "_already_pinned")
+
     _quantize_dirty_hack(model_to_quantize)
 
     print(f"Quantization of model '{model_id}' done")
 
     return True
 
+def load_loras_into_model(model, lora_path, lora_multi = None, verboseLevel = -1):
+    verboseLevel = _compute_verbose_level(verboseLevel)
+
+    if inject_adapter_in_model == None or set_weights_and_activate_adapters == None or  get_peft_kwargs == None:
+        raise Exception("Unable to load Lora, missing 'peft' and / or 'diffusers' modules")
+    
+    if not isinstance(lora_path, list):
+        lora_path = [lora_path]
+    
+    if lora_multi is None:
+        lora_multi = [1. for _ in lora_path]
+
+    for i, path in enumerate(lora_path):
+        adapter_name = str(i)
+
+        state_dict = safetensors2.torch_load_file(path)
+        
+        keys = list(state_dict.keys())
+        if len(keys) == 0:
+            raise Exception(f"Empty Lora '{path}'")
+
+
+        network_alphas = {}
+        for k in keys:
+            if "alpha" in k:
+                alpha_value = state_dict.pop(k)
+                if not ( (torch.is_tensor(alpha_value) and torch.is_floating_point(alpha_value)) or isinstance(
+                    alpha_value, float
+                )):
+                    network_alphas[k] =  torch.tensor( float(alpha_value.item() ) )
+
+        pos = keys[0].find(".")
+        prefix = keys[0][0:pos]
+        if not any( prefix.startswith(some_prefix) for some_prefix in ["diffusion_model", "transformer"]): 
+            msg = f"No compatible weight was found in Lora file '{path}'. Please check that it is compatible with the Diffusers format."
+            raise Exception(msg)
+
+        transformer = model
+
+        transformer_keys = [k for k in keys if k.startswith(prefix)]
+        state_dict = {
+            k.replace(f"{prefix}.", ""): v for k, v in state_dict.items() if k in transformer_keys
+        }
+
+        sd_keys = state_dict.keys()
+        if len(sd_keys) == 0:
+            print(f"No compatible weight was found in Lora file '{path}'. Please check that it is compatible with the Diffusers format.")
+            return
+
+        # is_correct_format = all("lora" in key for key in state_dict.keys())
+
+
+
+
+        # check with first key if is not in peft format
+        # first_key = next(iter(state_dict.keys()))
+        # if "lora_A" not in first_key:
+        #     state_dict = convert_unet_state_dict_to_peft(state_dict)
+
+        if adapter_name in getattr(transformer, "peft_config", {}):
+            raise ValueError(
+                f"Adapter name {adapter_name} already in use in the transformer - please select a new adapter name."
+            )
+
+        rank = {}
+        for key, val in state_dict.items():
+            if "lora_B" in key:
+                rank[key] = val.shape[1]
+
+        if network_alphas is not None and len(network_alphas) >= 1:
+            alpha_keys = [k for k in network_alphas.keys() if k.startswith(prefix) and k.split(".")[0] == prefix]
+            network_alphas = {k.replace(f"{prefix}.", ""): v for k, v in network_alphas.items() if k in alpha_keys}
+
+        lora_config_kwargs = get_peft_kwargs(rank, network_alpha_dict=network_alphas, peft_state_dict=state_dict)
+        
+        lora_config = LoraConfig(**lora_config_kwargs)
+        peft_kwargs = {}        
+        peft_kwargs["low_cpu_mem_usage"] = True
+        inject_adapter_in_model(lora_config, model, adapter_name=adapter_name, **peft_kwargs)
+
+        incompatible_keys = set_peft_model_state_dict(transformer, state_dict, adapter_name, **peft_kwargs)
+
+        warn_msg = ""
+        if incompatible_keys is not None:
+            # Check only for unexpected keys.
+            unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
+            if unexpected_keys:
+                pass
+        if verboseLevel >=1:
+            print(f"Lora '{path}' was loaded in model '{_get_module_name(model)}'")
+    set_weights_and_activate_adapters(model,[ str(i) for i in range(len(lora_multi))], lora_multi)
+
+
+def fast_load_transformers_model(model_path: str, do_quantize = False, quantizationType =  qint8, pinToMemory = False, partialPinning = False, verboseLevel = -1):
+    """
+    quick version of .LoadfromPretrained of  the transformers library
+    used to build a model and load the corresponding weights (quantized or not)
+    """       
+
+    
+    import os.path
+    from accelerate import init_empty_weights
+ 
+    if not (model_path.endswith(".sft") or model_path.endswith(".safetensors")):
+        raise Exception("full model path to file expected")
+
+    model_path = _get_model(model_path)
+    verboseLevel = _compute_verbose_level(verboseLevel)
+
+    with safetensors2.safe_open(model_path) as f:
+        metadata = f.metadata() 
+
+    if metadata is None:
+        transformer_config = None
+    else:
+        transformer_config = metadata.get("config", None)
+
+    if transformer_config == None:
+        config_fullpath =  os.path.join(os.path.dirname(model_path), "config.json")
+
+        if not os.path.isfile(config_fullpath):
+            raise Exception("a 'config.json' that describes the model is required in the directory of the model or inside the safetensor file")
+
+        with open(config_fullpath, "r", encoding="utf-8") as reader:
+            text = reader.read()
+        transformer_config= json.loads(text)
+
+
+    if "architectures" in transformer_config: 
+        architectures = transformer_config["architectures"]
+        class_name = architectures[0] 
+
+        module = __import__("transformers")
+        map = {  "T5WithLMHeadModel" : "T5EncoderModel"}
+        class_name = map.get(class_name, class_name)
+        transfomer_class = getattr(module, class_name)
+        from transformers import AutoConfig
+
+        import tempfile
+        with tempfile.NamedTemporaryFile("w", delete = False,  encoding ="utf-8") as fp: 
+            fp.write(json.dumps(transformer_config))
+            fp.close()
+            config_obj = AutoConfig.from_pretrained(fp.name)     
+        os.remove(fp.name)
+
+        #needed to keep inits of non persistent buffers
+        with init_empty_weights():
+            model = transfomer_class(config_obj)
+                
+        model = model.base_model
+
+    elif "_class_name" in transformer_config:
+        class_name = transformer_config["_class_name"]
+
+        module = __import__("diffusers")
+        transfomer_class = getattr(module, class_name)
+
+        with init_empty_weights():
+            model = transfomer_class.from_config(transformer_config)
+
+
+    torch.set_default_device('cpu')
+
+    model._config = transformer_config
+            
+    load_model_data(model,model_path, do_quantize = do_quantize, quantizationType = quantizationType, pinToMemory= pinToMemory, partialPinning= partialPinning, verboseLevel=verboseLevel )
+
+    return model
+
+
+
+def load_model_data(model, file_path: str, do_quantize = False, quantizationType = qint8, pinToMemory = False, partialPinning = False, verboseLevel = -1):
+    """
+    Load a model, detect if it has been previously quantized using quanto and do the extra setup if necessary
+    """
+
+    file_path = _get_model(file_path)
+    verboseLevel = _compute_verbose_level(verboseLevel)
+
+    model = _remove_model_wrapper(model)
+
+    # if pinToMemory and do_quantize:
+    #     raise Exception("Pinning and Quantization can not be used at the same time")
+
+    if not (".safetensors" in file_path or ".sft" in file_path): 
+        if pinToMemory:
+            raise Exception("Pinning to memory while loading only supported for safe tensors files")
+        state_dict = torch.load(file_path, weights_only=True)
+        if "module" in state_dict:
+            state_dict = state_dict["module"]
+    else:
+        state_dict, metadata = _safetensors_load_file(file_path)
+            
+        if metadata is None:
+            quantization_map = None
+        else:
+            quantization_map = metadata.get("quantization_map", None)
+            config = metadata.get("config", None)
+            if config is not None:
+                model._config = config
+
+
+
+        if quantization_map is None:
+            pos = str.rfind(file_path, ".")
+            if pos > 0:
+                quantization_map_path = file_path[:pos]
+            quantization_map_path += "_map.json"
+
+            if os.path.isfile(quantization_map_path):
+                with open(quantization_map_path, 'r') as f:
+                    quantization_map = json.load(f)
+
+
+
+        if quantization_map is None :
+            if "quanto" in file_path and not do_quantize:
+                print("Model seems to be quantized by quanto but no quantization map was found whether inside the model or in a separate '{file_path[:json]}_map.json' file")
+        else:
+            _requantize(model, state_dict, quantization_map)    
+
+    missing_keys , unexpected_keys = model.load_state_dict(state_dict, False,  assign = True )
+    # if len(missing_keys) > 0:
+    #     sd_crap = { k : None for k in missing_keys}
+    #     missing_keys , unexpected_keys = model.load_state_dict(sd_crap, strict =False,  assign = True )
+    del state_dict
+
+    for k,p in model.named_parameters():
+        if p.is_meta:
+            txt  = f"Incompatible State Dictionary or 'Init_Empty_Weights' not set since parameter '{k}' has no data"
+            raise Exception(txt)
+    for k,b in model.named_buffers():
+        if b.is_meta:
+            txt  = f"Incompatible State Dictionary or 'Init_Empty_Weights' not set since buffer '{k}' has no data"
+            raise Exception(txt)
+
+    if do_quantize:
+        if quantization_map is None:
+            if _quantize(model, quantizationType, verboseLevel=verboseLevel, model_id=file_path):
+                quantization_map = model._quanto_map  
+        else:
+            if verboseLevel >=1:
+                print("Model already quantized")
+
+    if pinToMemory:
+        _pin_to_memory(model, file_path, partialPinning = partialPinning, verboseLevel = verboseLevel)
+
+    return
+
 def get_model_name(model):
     return model.name
 
@@ -663,6 +911,7 @@ class offload:
         self.async_transfers = False
         global last_offload_obj
         last_offload_obj = self
+
         
     def add_module_to_blocks(self, model_id, blocks_name, submodule, prev_block_name):
 
@@ -684,15 +933,25 @@ class offload:
 
         for k,p in submodule.named_parameters(recurse=False):
             if isinstance(p, QTensor):
-                blocks_params.append( (submodule, k, p._data, p._scale) )
-                blocks_params_size += p._data.nbytes
-                blocks_params_size += p._scale.nbytes
+                blocks_params.append( (submodule, k, p ) )
+
+                if p._qtype == qint4:
+                    if hasattr(p,"_scale_shift"):
+                        blocks_params_size += torch.numel(p._scale_shift) * p._scale_shift.element_size()
+                        blocks_params_size += torch.numel(p._data._data) * p._data._data.element_size()
+                    else:
+                        blocks_params_size += torch.numel(p._scale) * p._scale.element_size()
+                        blocks_params_size += torch.numel(p._shift) * p._shift.element_size()
+                        blocks_params_size += torch.numel(p._data._data) * p._data._data.element_size()
+                else:
+                    blocks_params_size += torch.numel(p._scale) * p._scale.element_size()
+                    blocks_params_size += torch.numel(p._data) * p._data.element_size()
             else:
-                blocks_params.append( (submodule, k, p.data, None) )
-                blocks_params_size += p.data.nbytes
+                blocks_params.append( (submodule, k, p  ) )
+                blocks_params_size += torch.numel(p.data) * p.data.element_size()
 
         for k, p in submodule.named_buffers(recurse=False):
-            blocks_params.append( (submodule, k, p.data, None) )
+            blocks_params.append( (submodule, k, p) )
             blocks_params_size += p.data.nbytes
 
 
@@ -710,34 +969,28 @@ class offload:
                 return False    
         return True
 
-    def gpu_load_blocks(self, model_id, blocks_name, async_load = False):
+    @torch.compiler.disable()
+    def gpu_load_blocks(self, model_id, blocks_name):
         # cl = clock.start()
 
         if blocks_name != None:
             self.loaded_blocks[model_id] = blocks_name           
 
         entry_name = model_id if blocks_name is None else model_id + "/" + blocks_name
-
-        def cpu_to_gpu(stream_to_use, blocks_params, record_for_stream = None):
+        
+        def cpu_to_gpu(stream_to_use, blocks_params): #, record_for_stream = None
             with torch.cuda.stream(stream_to_use):
                 for param in blocks_params:
-                    parent_module, n, data, scale  = param
-                    p = getattr(parent_module, n)
-                    if isinstance(p, QTensor):
-                        q = WeightQBytesTensor.create(p.qtype, p.axis, p.size(), p.stride(), data.cuda(non_blocking=True), scale.cuda(non_blocking=True), activation_qtype=p.activation_qtype, requires_grad=p.requires_grad )
-                        #q = p.to("cuda", non_blocking=True)
-                        q = torch.nn.Parameter(q , requires_grad=False)
-                        setattr(parent_module, n , q)
-                        del p
-                    else:
-                        p.data = p.data.cuda(non_blocking=True) 
-                    
-                    if record_for_stream != None:
-                        if isinstance(p, QTensor):
-                            q._data.record_stream(record_for_stream)
-                            q._scale.record_stream(record_for_stream)
-                        else:
-                            p.data.record_stream(record_for_stream)
+                    parent_module, n, p  = param
+                    q = p.to("cuda", non_blocking=True)
+                    q = torch.nn.Parameter(q , requires_grad=False)
+                    setattr(parent_module, n , q)
+                    # if record_for_stream != None:
+                    #     if isinstance(p, QTensor):
+                    #         q._data.record_stream(record_for_stream)
+                    #         q._scale.record_stream(record_for_stream)
+                    #     else:
+                    #         p.data.record_stream(record_for_stream)
 
 
         if self.verboseLevel >=2:
@@ -762,7 +1015,7 @@ class offload:
         # cl.stop()
         # print(f"load time: {cl.format_time_gap()}")
 
-
+    @torch.compiler.disable()
     def gpu_unload_blocks(self, model_id, blocks_name):
         # cl = clock.start()
         if blocks_name != None:
@@ -776,23 +1029,14 @@ class offload:
             print(f"Unloading model {blocks_name} ({model_name}) from GPU")
  
         blocks_params = self.blocks_of_modules[blocks_name]
-
         for param in blocks_params:
-            parent_module, n, data, scale = param
-            p = getattr(parent_module, n) 
-            if isinstance(p, QTensor):
-                # need to change the parameter directly from the module as it can't be swapped in place due to a memory leak in the pytorch compiler
-                q = WeightQBytesTensor.create(p.qtype, p.axis, p.size(), p.stride(), data, scale, activation_qtype=p.activation_qtype, requires_grad=p.requires_grad )
-                q = torch.nn.Parameter(q , requires_grad=False)
-                setattr(parent_module, n , q)
-                del p
-            else:
-                p.data = data
-
+            parent_module, n, p  = param
+            q = torch.nn.Parameter(p , requires_grad=False)
+            setattr(parent_module, n , q)
         # cl.stop()
         # print(f"unload time: {cl.format_time_gap()}")
 
-
+    # @torch.compiler.disable()
     def gpu_load(self, model_id):
         model = self.models[model_id]
         self.active_models.append(model)
@@ -824,8 +1068,8 @@ class offload:
             if torch.is_tensor(arg):    
                 if arg.dtype == torch.float32:
                     arg = arg.to(torch.bfloat16).cuda(non_blocking=True)             
-                else:
-                    arg = arg.cuda(non_blocking=True)             
+                elif not arg.is_cuda:
+                    arg = arg.cuda(non_blocking=True)
             new_args.append(arg)
 
         for k in kwargs:
@@ -833,7 +1077,7 @@ class offload:
             if torch.is_tensor(arg):
                 if arg.dtype == torch.float32:
                     arg = arg.to(torch.bfloat16).cuda(non_blocking=True)             
-                else:
+                elif not arg.is_cuda:
                     arg = arg.cuda(non_blocking=True)             
             new_kwargs[k]= arg
         
@@ -874,10 +1118,10 @@ class offload:
         
         return False
 
-    def hook_load_data_if_needed(self, target_module, model_id,blocks_name, context):
+    def hook_preload_blocks_for_compilation(self, target_module, model_id,blocks_name, context):
 
-        @torch.compiler.disable()
-        def load_data_if_needed(module,  *args, **kwargs):
+        # @torch.compiler.disable()
+        def preload_blocks_for_compile(module,  *args, **kwargs):
             some_context = context #for debugging
             if blocks_name == None:
                 if self.ready_to_check_mem():
@@ -891,12 +1135,17 @@ class offload:
                             self.empty_cache_if_needed()
                     self.loaded_blocks[model_id] = blocks_name
                     self.gpu_load_blocks(model_id, blocks_name)
-
-        target_module.register_forward_pre_hook(load_data_if_needed)        
+        # need to be registered before the forward not to be break the efficiency of the compilation chain
+        # it should be at the top of the compilation as this type of hook in the middle of a chain seems to break memory performance
+        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):
 
+        qint4quantization =  isinstance(target_module, QModuleMixin) and  target_module.weight!= None and  target_module.weight.qtype == qint4 
+        if qint4quantization:
+            pass
+
         def check_empty_cuda_cache(module, *args, **kwargs):
             # if self.ready_to_check_mem():
             #     self.empty_cache_if_needed()
@@ -912,6 +1161,8 @@ class offload:
                             self.empty_cache_if_needed()
                     self.loaded_blocks[model_id] = blocks_name
                     self.gpu_load_blocks(model_id, blocks_name)
+            if qint4quantization:
+                args, kwargs = self.move_args_to_gpu(*args, **kwargs)
 
             return previous_method(*args, **kwargs) 
 
@@ -959,177 +1210,18 @@ class offload:
             print(f"Hooked in model '{model_id}' ({model_name})")
 
 
-    # Not implemented yet, but why would one want to get rid of these features ?
-    # def unhook_module(module: torch.nn.Module):
-    #     if not hasattr(module,"_mm_id"):
-    #         return
-        
-    #     delattr(module, "_mm_id")
-                 
-    # def unhook_all(parent_module: torch.nn.Module):
-    #     for module in parent_module.components.items():
-    #         self.unhook_module(module)
-
-def fast_load_transformers_model(model_path: str, do_quantize = False, quantization_type = qint8, pinToMemory = False, partialPinning = False,  verboseLevel = -1):
-    """
-    quick version of .LoadfromPretrained of  the transformers library
-    used to build a model and load the corresponding weights (quantized or not)
-    """       
-
-    
-    import os.path
-    from accelerate import init_empty_weights
- 
-    if not (model_path.endswith(".sft") or model_path.endswith(".safetensors")):
-        raise Exception("full model path to file expected")
-
-    model_path = _get_model(model_path)
-    verboseLevel = _compute_verbose_level(verboseLevel)
-
-    with safetensors2.safe_open(model_path) as f:
-        metadata = f.metadata() 
-
-    if metadata is None:
-        transformer_config = None
-    else:
-        transformer_config = metadata.get("config", None)
-
-    if transformer_config == None:
-        config_fullpath =  os.path.join(os.path.dirname(model_path), "config.json")
-
-        if not os.path.isfile(config_fullpath):
-            raise Exception("a 'config.json' that describes the model is required in the directory of the model or inside the safetensor file")
-
-        with open(config_fullpath, "r", encoding="utf-8") as reader:
-            text = reader.read()
-        transformer_config= json.loads(text)
-
-
-    if "architectures" in transformer_config: 
-        architectures = transformer_config["architectures"]
-        class_name = architectures[0] 
-
-        module = __import__("transformers")
-        transfomer_class = getattr(module, class_name)
-        from transformers import AutoConfig
-
-        import tempfile
-        with tempfile.NamedTemporaryFile("w", delete = False,  encoding ="utf-8") as fp: 
-            fp.write(json.dumps(transformer_config))
-            fp.close()
-            config_obj = AutoConfig.from_pretrained(fp.name)     
-        os.remove(fp.name)
-
-        #needed to keep inits of non persistent buffers
-        with init_empty_weights():
-            model = transfomer_class(config_obj)
-                
-        model = model.base_model
-
-    elif "_class_name" in transformer_config:
-        class_name = transformer_config["_class_name"]
-
-        module = __import__("diffusers")
-        transfomer_class = getattr(module, class_name)
-
-        with init_empty_weights():
-            model = transfomer_class.from_config(transformer_config)
-
-
-    torch.set_default_device('cpu')
-
-    model._config = transformer_config
-            
-    load_model_data(model,model_path, do_quantize = do_quantize, quantization_type = quantization_type, pinToMemory= pinToMemory, partialPinning= partialPinning, verboseLevel=verboseLevel )
-
-    return model
-
-
-
-def load_model_data(model, file_path: str, do_quantize = False, quantization_type = qint8, pinToMemory = False, partialPinning = False, verboseLevel = -1):
-    """
-    Load a model, detect if it has been previously quantized using quanto and do the extra setup if necessary
-    """
-
-    file_path = _get_model(file_path)
-    verboseLevel = _compute_verbose_level(verboseLevel)
-
-    model = _remove_model_wrapper(model)
-
-    # if pinToMemory and do_quantize:
-    #     raise Exception("Pinning and Quantization can not be used at the same time")
-
-    if not (".safetensors" in file_path or ".sft" in file_path): 
-        if pinToMemory:
-            raise Exception("Pinning to memory while loading only supported for safe tensors files")
-        state_dict = torch.load(file_path, weights_only=True)
-        if "module" in state_dict:
-            state_dict = state_dict["module"]
-    else:
-        state_dict, metadata = _safetensors_load_file(file_path)
-
-
-        # if pinToMemory:
-        #     _pin_to_memory_sd(model,state_dict, file_path, partialPinning = partialPinning, perc_reserved_mem_max = perc_reserved_mem_max, verboseLevel = verboseLevel)
-
-        # with safetensors2.safe_open(file_path) as f:
-        #     metadata = f.metadata() 
-
-            
-        if metadata is None:
-            quantization_map = None
-        else:
-            quantization_map = metadata.get("quantization_map", None)
-            config = metadata.get("config", None)
-            if config is not None:
-                model._config = config
-
-
-
-        if quantization_map is None:
-            pos = str.rfind(file_path, ".")
-            if pos > 0:
-                quantization_map_path = file_path[:pos]
-            quantization_map_path += "_map.json"
-
-            if os.path.isfile(quantization_map_path):
-                with open(quantization_map_path, 'r') as f:
-                    quantization_map = json.load(f)
-
-
-
-        if quantization_map is None :
-            if "quanto" in file_path and not do_quantize:
-                print("Model seems to be quantized by quanto but no quantization map was found whether inside the model or in a separate '{file_path[:json]}_map.json' file")
-        else:
-            _requantize(model, state_dict, quantization_map)    
-
-    missing_keys , unexpected_keys = model.load_state_dict(state_dict, strict = quantization_map is None,  assign = True )
-    del state_dict
-
-    if do_quantize:
-        if quantization_map is None:
-            if _quantize(model, quantization_type, verboseLevel=verboseLevel, model_id=file_path):
-                quantization_map = model._quanto_map  
-        else:
-            if verboseLevel >=1:
-                print("Model already quantized")
-
-    if pinToMemory:
-        _pin_to_memory(model, file_path, partialPinning = partialPinning, verboseLevel = verboseLevel)
-
-    return
-
-def save_model(model, file_path, do_quantize = False, quantization_type = qint8, verboseLevel = -1 ):
+def save_model(model, file_path, do_quantize = False, quantizationType = qint8, verboseLevel = -1, config_file_path = None ):
     """save the weights of a model and quantize them if requested
     These weights can be loaded again using 'load_model_data'
     """       
     
     config = None
-
     verboseLevel = _compute_verbose_level(verboseLevel)
-
-    if hasattr(model, "_config"):
+    if config_file_path !=None:
+        with open(config_file_path, "r", encoding="utf-8") as reader:
+            text = reader.read()
+            config= json.loads(text)
+    elif hasattr(model, "_config"):
         config = model._config
     elif hasattr(model, "config"):
         config_fullpath = None
@@ -1147,7 +1239,7 @@ def save_model(model, file_path, do_quantize = False, quantization_type = qint8,
                 config= json.loads(text)
 
     if do_quantize:
-        _quantize(model, weights=quantization_type, model_id=file_path)
+        _quantize(model, weights=quantizationType, model_id=file_path)
     
     quantization_map = getattr(model, "_quanto_map", None)
 
@@ -1155,12 +1247,12 @@ def save_model(model, file_path, do_quantize = False, quantization_type = qint8,
         print(f"Saving file '{file_path}")
     safetensors2.torch_write_file(model.state_dict(),  file_path , quantization_map = quantization_map, config = config)
     if verboseLevel >=1:
-        print(f"File '{file_path} saved")
+        print(f"File '{file_path}' saved")
 
 
 
 
-def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = True,  extraModelsToQuantize = None, budgets= 0, asyncTransfers = True, compile = False, perc_reserved_mem_max = 0, verboseLevel = -1):
+def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = True,  extraModelsToQuantize = None, quantizationType = qint8, budgets= 0, asyncTransfers = True, compile = False, perc_reserved_mem_max = 0, 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
@@ -1238,13 +1330,14 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = Tru
     self.anyCompiledModule = compileAllModels or len(modelsToCompile)>0
     if self.anyCompiledModule:
         torch._dynamo.config.cache_size_limit = 10000
+        torch.compiler.reset()
+
       #  torch._logging.set_logs(recompiles=True)
       #  torch._inductor.config.realize_opcount_threshold = 100 # workaround bug "AssertionError: increase TRITON_MAX_BLOCK['X'] to 4096."
 
     max_reservable_memory = _get_max_reservable_memory(perc_reserved_mem_max)
 
     estimatesBytesToPin = 0
-
     for model_id in models: 
         current_model: torch.nn.Module = models[model_id] 
         # make sure that no RAM or GPU memory is not allocated for gradiant / training
@@ -1252,19 +1345,30 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = Tru
         
         # if the model has just been quantized so there is no need to quantize it again
         if model_id in models_to_quantize:
-            _quantize(current_model, weights=qint8, verboseLevel = self.verboseLevel, model_id=model_id)
+            _quantize(current_model, weights=quantizationType, verboseLevel = self.verboseLevel, model_id=model_id)
 
         modelPinned = (pinAllModels or model_id in modelsToPin) and not hasattr(current_model,"_already_pinned")
 
-        current_model_size = 0    
-        # load all the remaining unread lazy safetensors in RAM to free open cache files 
-        for p in current_model.parameters():
+        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) 
-                current_model_size +=  torch.numel(p._scale) * p._scale.element_size()
-                current_model_size +=  torch.numel(p._data) * p._data.element_size()
+                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()
+
+                else:
+                    current_model_size +=  torch.numel(p._scale) * p._scale.element_size()
+                    current_model_size +=  torch.numel(p._data) * p._data.element_size()
+
             else:
                 if p.data.dtype == torch.float32:
                     # convert any left overs float32 weight to bloat16 to divide by 2 the model memory footprint
@@ -1272,7 +1376,7 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = Tru
                 current_model_size +=  torch.numel(p.data) * p.data.element_size()
                             
         for b in current_model.buffers():
-            if b.data.dtype == torch.float32:
+            if b.data.dtype == torch.float32: 
                 # convert any left overs float32 weight to bloat16 to divide by 2 the model memory footprint
                 b.data = b.data.to(torch.bfloat16)
             current_model_size +=  torch.numel(b.data) * b.data.element_size()
@@ -1298,22 +1402,21 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = Tru
     #  Hook forward methods of modules 
     for model_id in models: 
         current_model: torch.nn.Module = models[model_id] 
-        current_budget = model_budgets[model_id]
-        current_size = 0
-        cur_blocks_prefix, prev_blocks_name, cur_blocks_name,cur_blocks_seq = None, None, None, -1
-        self.loaded_blocks[model_id] = None
         towers_names, towers_modules = _detect_main_towers(current_model)
-        towers_names = [n +"." for n in towers_names]
         if self.verboseLevel>=2 and len(towers_names)>0:
             print(f"Potential iterative blocks found in model '{model_id}':{towers_names}")
         # compile main iterative modules stacks ("towers")
-        if compileAllModels or model_id in modelsToCompile :
-            #torch.compiler.reset()
+        compilationInThisOne = compileAllModels or model_id in modelsToCompile 
+        if compilationInThisOne:
             if self.verboseLevel>=1:
-                print(f"Pytorch compilation of model '{model_id}' is scheduled.")
-            for tower in towers_modules:
-                for submodel in tower:
-                    submodel.forward= torch.compile(submodel.forward,  backend= "inductor", mode="default" ) # , fullgraph= True, mode= "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs",  
+                if len(towers_modules)>0:
+                    print(f"Pytorch compilation of model '{model_id}' is scheduled.")
+                else:
+                    print(f"Pytorch compilation of model '{model_id}' is not yet supported.")
+
+            for submodel in towers_modules:
+                # for submodel in tower:
+                submodel.forward= torch.compile(submodel.forward,  backend= "inductor", mode="default" ) # , fullgraph= True, mode= "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs",  
                     #dynamic=True,
                 
         if pinAllModels or model_id in modelsToPin:
@@ -1323,6 +1426,11 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = Tru
             else:
                 _pin_to_memory(current_model, model_id, partialPinning= partialPinning, perc_reserved_mem_max=perc_reserved_mem_max, verboseLevel=verboseLevel)            
 
+        current_budget = model_budgets[model_id]
+        current_size = 0
+        cur_blocks_prefix, prev_blocks_name, cur_blocks_name,cur_blocks_seq = None, None, None, -1
+        self.loaded_blocks[model_id] = None
+
         for submodule_name, submodule in current_model.named_modules():  
             # create a fake 'accelerate' parameter so that the _execution_device property returns always "cuda" 
             # (it is queried in many pipelines even if offloading is not properly implemented)  
@@ -1331,44 +1439,43 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, quantizeTransformer = Tru
 
             if submodule_name=='':
                 continue
-            newListItem = False
+            
             if current_budget > 0:
-                if isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)): #
-                    if cur_blocks_prefix == None:
-                        cur_blocks_prefix = submodule_name + "."
+                if cur_blocks_prefix != None:
+                    if submodule_name.startswith(cur_blocks_prefix):
+                        depth_prefix = cur_blocks_prefix.split(".")
+                        depth_name = submodule_name.split(".")
+                        level  =  depth_name[len(depth_prefix)-1]                        
+                        pre , num = _extract_num_from_str(level)
+                        if num != cur_blocks_seq and (cur_blocks_seq == -1 or current_size > current_budget): 
+                            prev_blocks_name = cur_blocks_name
+                            cur_blocks_name =  cur_blocks_prefix + str(num)
+                            # print(f"new block: {model_id}/{cur_blocks_name} - {submodule_name}")
+                        cur_blocks_seq = num
                     else:
-                        #if cur_blocks_prefix != submodule_name[:len(cur_blocks_prefix)]:
-                        if not submodule_name.startswith(cur_blocks_prefix):
-                            cur_blocks_prefix = submodule_name + "."
-                            cur_blocks_name,cur_blocks_seq = None, -1
-                else:
-                    
-                    if cur_blocks_prefix is not None:
-                        if submodule_name.startswith(cur_blocks_prefix):
-                            num = int(submodule_name[len(cur_blocks_prefix):].split(".")[0])
-                            newListItem= num != cur_blocks_seq
-                            if num != cur_blocks_seq and (cur_blocks_name == None or current_size > current_budget): 
-                                prev_blocks_name = cur_blocks_name
-                                cur_blocks_name = cur_blocks_prefix + str(num)
-                                # print(f"new block: {model_id}/{cur_blocks_name} - {submodule_name}")
-                            cur_blocks_seq = num
-                        else:
-                            cur_blocks_prefix, prev_blocks_name, cur_blocks_name,cur_blocks_seq = None, None, None, -1
-
+                        cur_blocks_prefix, prev_blocks_name, cur_blocks_name,cur_blocks_seq = None, None, None, -1
+
+                if cur_blocks_prefix == None:
+                    pre , num = _extract_num_from_str(submodule_name)
+                    if isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):  
+                        cur_blocks_prefix, prev_blocks_name, cur_blocks_seq = pre + ".", None, -1
+                    elif num >=0:
+                        cur_blocks_prefix, prev_blocks_name, cur_blocks_seq = pre, None, num
+                        cur_blocks_name = submodule_name
+                        # print(f"new block: {model_id}/{cur_blocks_name} - {submodule_name}")
+                          
+ 
             if hasattr(submodule, "forward"):
                 submodule_method = getattr(submodule, "forward")
                 if callable(submodule_method):   
                     if len(submodule_name.split("."))==1:
                         self.hook_change_module(submodule, current_model, model_id, submodule_name, submodule_method)
-                    elif newListItem: 
-                        self.hook_load_data_if_needed(submodule, model_id, cur_blocks_name, context = submodule_name )
+                    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 )
 
-
-                    current_size = self.add_module_to_blocks(model_id, cur_blocks_name, submodule, prev_blocks_name)
-
-
+                current_size = self.add_module_to_blocks(model_id, cur_blocks_name, submodule, prev_blocks_name)
 
 
     if self.verboseLevel >=2:
@@ -1406,7 +1513,7 @@ def profile(pipe_or_dict_of_modules, profile_no: profile_type =  profile_type.Ve
         modules= modules.components 
 
     modules = {k: _remove_model_wrapper(v) for k, v in modules.items() if isinstance(v, torch.nn.Module)}
-    module_names = {k: v.__module__.lower() for k, v in modules.items() }
+    module_names = {k: _get_module_name(v) for k, v in modules.items() }
 
     default_extraModelsToQuantize = []
     quantizeTransformer = True
@@ -1414,11 +1521,12 @@ def profile(pipe_or_dict_of_modules, profile_no: profile_type =  profile_type.Ve
     models_to_scan = ("text_encoder", "text_encoder_2")
     candidates_to_quantize = ("t5", "llama", "llm")
     for model_id  in models_to_scan:
-        name = module_names[model_id]
-        for candidate in candidates_to_quantize:
-            if candidate in name:
-                default_extraModelsToQuantize.append(model_id)
-                break
+        if model_id in module_names: 
+            name = module_names[model_id]
+            for candidate in candidates_to_quantize:
+                if candidate in name:
+                    default_extraModelsToQuantize.append(model_id)
+                    break
 
 
     # transformer (video or image generator) should be as small as possible not to occupy space that could be used by actual image data
@@ -1427,6 +1535,7 @@ def profile(pipe_or_dict_of_modules, profile_no: profile_type =  profile_type.Ve
     default_budgets = { "transformer" : 600 , "text_encoder": 3000, "text_encoder_2": 3000 }
     extraModelsToQuantize = None
     asyncTransfers = True
+    budgets = None
 
     if profile_no == profile_type.HighRAM_HighVRAM:
         pinnedMemory= True