mmgp 1.2.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

mmgp-2.0.1.dist-info/METADATA

@@ -0,0 +1,137 @@
+Metadata-Version: 2.1
+Name: mmgp
+Version: 2.0.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
+
+
+<p align="center">
+  <H2>Memory Management 2.0 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 12 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 (32 GB of RAM and 12 VRAM) and to run it at a very good speed on a high end consumer config (48 GB of RAM and 24 GB of VRAM).
+
+Each profile may use the following: 
+- Smart preloading of models in RAM to reduce RAM requirements
+- 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 in 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 quantized models
+
+## 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_Fastest: at least 48 GB of RAM and 24 GB of VRAM : the fastest well suited for a RTX 3090 / RTX 4090
+- HighRAM_LowVRAM_Fast (recommended): 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_Medium: at least 32 GB of RAM and 24 GB of VRAM : so so speed but adapted for RTX 3090 / RTX 4090 with limited RAM
+- LowRAM_LowVRAM_Slow: at least 32 GB of RAM and 12 GB of VRAM :  if have little VRAM or generate longer videos 
+- VerylowRAM_LowVRAM_Slowest: 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
+
+By default the 'transformer' will be quantized to 8 bits for all profiles. If you don't want that you may specify the optional parameter *quantizeTransformer = False*.
+
+## Alternatively you may want to create your own profile with specific parameters:
+
+For example:
+```
+  from mmgp import offload
+  offload.all(pipe, pinInRAM=True, modelsToQuantize = ["text_encoder_2"] )
+```  
+- pinInRAM: Boolean (for all models) or List of models ids to pin in RAM. Every model pinned in RAM will load much faster (4 times) but this requires more RAM
+- modelsToQuantize: list of model ids to quantize on the fly. If the corresponding model is already quantized, this option will be ignored.
+- 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.
+- budgets: either a number in mega bytes (for all models, if 0 unlimited budget) or a dictionary that maps model ids to mega bytes : define the budget in VRAM (in fact the real number is 2.5 this number) that is allocated in VRAM for each model. 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. Turning on the PinInRAM accelerates greatly (4x) small budgets but consumes usually 50% more RAM.
+
+
+##  Going further
+
+The module includes several tools to package a light version of your favorite video / image generator:
+- *save_model(model, file_path, do_quantize = False, quantization_type = 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).
+If the model is saved in a quantized format, an extra file that ends with '_map.json' will be created and needed to reload the model again.
+
+- *load_model_data(model, file_path: str)*\
+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.
+
+- *fast_load_transformers_model(model_path: str)*\
+Initialize (build the model hierarchy in memory) and fast load the corresponding tensors of a 'transformers' library model.
+The advantages over the original *from_pretrained* method is that the 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 model are processed in a transparent way.
+Please note that you need to keep the original file transformers 'config.json' in the same directory. 
+
+
+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 need to 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 that there should be always one model whose Id is 'transformer'. It 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).
+
+Becareful, 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-2.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mmgp.py,sha256=UuAUF76QIve6j2qEkYucxaOq4uo9mKptP4AdkQKr8Eg,45152
+mmgp-2.0.1.dist-info/LICENSE.md,sha256=HjzvY2grdtdduZclbZ46B2M-XpT4MDCxFub5ZwTWq2g,93
+mmgp-2.0.1.dist-info/METADATA,sha256=y-6bIJqU6FrX4NMVXheTjs7n2PeoG-kilyyULqgxnt4,8601
+mmgp-2.0.1.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+mmgp-2.0.1.dist-info/top_level.txt,sha256=waGaepj2qVfnS2yAOkaMu4r9mJaVjGbEi6AwOUogU_U,14
+mmgp-2.0.1.dist-info/RECORD,,

mmgp.py

@@ -1,24 +1,28 @@
-# ------------------ Memory Management for the GPU Poor by DeepBeepMeep (mmgp)------------------
+# ------------------ Memory Management 2.0 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
 # times in a pipe (eg VAE).
 #
 # Requirements:
-# - GPU: RTX 3090/ RTX 4090 (24 GB of VRAM)
-# - RAM: minimum 48 GB, recommended 64 GB 
+# - VRAM: minimum 12 GB, recommended 24 GB (RTX 3090/ RTX 4090)
+# - RAM: minimum 24 GB, recommended 48 - 64 GB 
 #
 # 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, you may want to choose between 4 profiles depending on your hardware, for instance:
+#   from mmgp import offload, profile_type
+#   offload.profile(pipe, profile_type.HighRAM_LowVRAM_Fast)
+# Alternatively you may want to your own parameters, for instance:
 #   from mmgp import offload
-#   offload.all(pipe)
+#   offload.all(pipe, pinInRAM=true, modelsToQuantize = ["text_encoder_2"] )
 # The 'transformer' model that contains usually the video or image generator is quantized on the fly by default to 8 bits so that it can fit into 24 GB of VRAM. 
 # If you want to save time on disk and reduce the loading time, you may want to load directly a prequantized model. In that case you need to set the option quantizeTransformer to False to turn off on the fly quantization.
 # You can specify a list of additional models string ids to quantize (for instance the text_encoder) using the optional argument modelsToQuantize. This may be useful if you have less than 48 GB of RAM.
 # Note that there is little advantage on the GPU / VRAM side to quantize text encoders as their inputs are usually quite light. 
-# Conversely if you have more than 64GB RAM you may want to enable RAM pinning with the option pinInRAM = True. You will get in return super fast loading / unloading of models
+# Conversely if you have more than 48GB RAM you may want to enable RAM pinning with the option pinInRAM = True. You will get in return super fast loading / unloading of models
 # (this can save significant time if the same pipeline is run multiple times in a row)
 # 
 # Sometime there isn't an explicit pipe object as each submodel is loaded separately in the main app. If this is the case, you need to create a dictionary that manually maps all the models.
@@ -53,10 +57,15 @@ import torch
 import gc
 import time
 import functools
+import sys
+import json
+
 from optimum.quanto import freeze, qfloat8, qint8, quantize, QModuleMixin, QTensor
 
 
 
+ONE_MB =  1048576
+
 cotenants_map = { 
                              "text_encoder": ["vae", "text_encoder_2"],
                              "text_encoder_2": ["vae", "text_encoder"],                             
@@ -79,10 +88,107 @@ def move_tensors(obj, device):
     else:
         raise TypeError("Tensor or list / dict of tensors expected")
 
+def _quantize(model_to_quantize, weights=qint8, verboseLevel = 1, threshold = 1000000000, model_id = None):
+    
+    sizeofbfloat16 = torch.bfloat16.itemsize
+
+    def compute_submodule_size(submodule):
+        size = 0
+        for p in submodule.parameters(recurse=False):
+            size  += torch.numel(p.data) * sizeofbfloat16
+
+        for p in submodule.buffers(recurse=False):
+            size  += torch.numel(p.data) * sizeofbfloat16
+
+        return size
+    
+    total_size =0
+    total_excluded = 0
+    exclude_list = []
+    submodule_size = 0
+    submodule_names = []
+    cur_blocks_prefix = None
+    prev_blocks_prefix = None
+
+    print(f"Quantization of model '{model_id}' started")
+
+    for submodule_name, submodule in model_to_quantize.named_modules():  
+        if isinstance(submodule, QModuleMixin):
+            if verboseLevel>=1:
+                print("No quantization to do as model is already quantized")
+            return False
+
+
+        if submodule_name=='':
+            continue
+
+
+        flush = False
+        if isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
+            if cur_blocks_prefix == None:
+                cur_blocks_prefix = submodule_name + "."
+                flush = True                    
+            else:
+                #if cur_blocks_prefix != submodule_name[:len(cur_blocks_prefix)]:
+                if not submodule_name.startswith(cur_blocks_prefix):
+                    cur_blocks_prefix = submodule_name + "."
+                    flush = True                    
+        else:                
+            if cur_blocks_prefix is not None:
+                #if not cur_blocks_prefix == submodule_name[0:len(cur_blocks_prefix)]:
+                if not submodule_name.startswith(cur_blocks_prefix):
+                    cur_blocks_prefix = None 
+                    flush = True                    
+
+        if flush:
+            if submodule_size <= threshold:
+                exclude_list += submodule_names
+                if verboseLevel >=2:
+                    print(f"Excluded size {submodule_size/ONE_MB:.1f} MB: {prev_blocks_prefix} : {submodule_names}")
+                total_excluded += submodule_size
+
+            submodule_size = 0
+            submodule_names = []
+        prev_blocks_prefix = cur_blocks_prefix
+        size = compute_submodule_size(submodule)
+        submodule_size += size
+        total_size += size
+        submodule_names.append(submodule_name)
+
+    if submodule_size > 0 and submodule_size <= threshold:
+        exclude_list += submodule_names
+        if verboseLevel >=2:
+            print(f"Excluded size {submodule_size/ONE_MB:.1f} MB: {prev_blocks_prefix} : {submodule_names}")
+        total_excluded += submodule_size
+
+    perc_excluded =total_excluded/ total_size if total_size >0 else 1
+    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.")
+        exclude_list = None
+
+            # we are obviously loading a model that has been already quantized
+
+    quantize(model_to_quantize,weights, exclude= exclude_list)
+    freeze(model_to_quantize)
+    torch.cuda.empty_cache()
+    gc.collect()         
+    print(f"Quantization of model '{model_id}' done")
+
+    return True
 
 def get_model_name(model):
     return model.name
 
+import enum    
+class profile_type(int, enum.Enum): 
+    HighRAM_HighVRAM_Fastest = 1
+    HighRAM_LowVRAM_Fast = 2
+    LowRAM_HighVRAM_Medium = 3
+    LowRAM_LowVRAM_Slow = 4
+    VerylowRAM_LowVRAM_Slowest = 5
+
 class HfHook:
     def __init__(self):
         self.execution_device = "cuda"
@@ -94,28 +200,57 @@ class offload:
     def __init__(self):
         self.active_models = []
         self.active_models_ids = []
+        self.active_subcaches = {}        
         self.models = {}
-        self.verbose = False
+        self.verboseLevel = 0
         self.models_to_quantize = []
         self.pinned_modules_data = {}
-        self.params_of_modules = {}
-        self.pinTensors = False
+        self.blocks_of_modules = {}
+        self.blocks_of_modules_sizes = {}
+        self.compile = False
         self.device_mem_capacity = torch.cuda.get_device_properties(0).total_memory
         self.last_reserved_mem_check =0
+        self.loaded_blocks = {}
+        self.prev_blocks_names = {}
+        self.next_blocks_names = {}
+        self.default_stream = torch.cuda.default_stream(torch.device("cuda")) # torch.cuda.current_stream()
+        self.transfer_stream = torch.cuda.Stream()
+        self.async_transfers = False
+
 
-    def collect_module_parameters(self, module: torch.nn.Module, module_params):
-        if isinstance(module, (torch.nn.ModuleList, torch.nn.Sequential)):
-            for i in range(len(module)):
-                current_layer = module[i]
-                module_params.extend(current_layer.parameters())
-                module_params.extend(current_layer.buffers())
+    def add_module_to_blocks(self, model_id, blocks_name, submodule, prev_block_name):
+
+        entry_name = model_id if blocks_name is None else model_id + "/" + blocks_name
+        if entry_name in self.blocks_of_modules:
+            blocks_params = self.blocks_of_modules[entry_name]
+            blocks_params_size = self.blocks_of_modules_sizes[entry_name]
         else:
-            for p in module.parameters(recurse=False):
-                module_params.append(p)
-            for p in module.buffers(recurse=False):
-                module_params.append(p)
-            for sub_module in module.children():
-                self.collect_module_parameters(sub_module, module_params)
+            blocks_params = []
+            self.blocks_of_modules[entry_name] = blocks_params
+            blocks_params_size = 0
+            if blocks_name !=None:
+                prev_entry_name = None if prev_block_name == None else  model_id + "/" + prev_block_name
+                self.prev_blocks_names[entry_name] =  prev_entry_name
+                if not prev_block_name == None:
+                    self.next_blocks_names[prev_entry_name] = entry_name        
+
+        for p in submodule.parameters(recurse=False):
+            blocks_params.append(p)
+            if isinstance(p, QTensor):
+                blocks_params_size += p._data.nbytes
+                blocks_params_size += p._scale.nbytes
+            else:
+                blocks_params_size += p.data.nbytes
+
+        for p in submodule.buffers(recurse=False):
+            blocks_params.append(p)      
+            blocks_params_size += p.data.nbytes
+
+
+        self.blocks_of_modules_sizes[entry_name] = blocks_params_size
+
+        return blocks_params_size
+
 
     def can_model_be_cotenant(self, model_id):
         potential_cotenants= cotenants_map.get(model_id, None)
@@ -126,45 +261,113 @@ class offload:
                 return False    
         return True
 
-    def gpu_load(self, model_id):
-        model = self.models[model_id]
-        self.active_models.append(model)
-        self.active_models_ids.append(model_id)
-        if self.verbose:
+    @torch.compiler.disable()    
+    def gpu_load_blocks(self, model_id, blocks_name, async_load = False):
+        if blocks_name != None:
+            self.loaded_blocks[model_id] = blocks_name           
+
+        def cpu_to_gpu(stream_to_use, blocks_params, record_for_stream = None):
+            with torch.cuda.stream(stream_to_use):
+                for p in blocks_params:
+                    if isinstance(p, QTensor):
+                        p._data = p._data.cuda(non_blocking=True)             
+                        p._scale = p._scale.cuda(non_blocking=True)
+                    else:
+                        p.data = p.data.cuda(non_blocking=True) 
+                    
+                    if record_for_stream != None:
+                        if isinstance(p, QTensor):
+                            p._data.record_stream(record_for_stream)
+                            p._scale.record_stream(record_for_stream)
+                        else:
+                            p.data.record_stream(record_for_stream)
+
+
+        entry_name = model_id if blocks_name is None else model_id + "/" + blocks_name
+        if self.verboseLevel >=2:
+            model = self.models[model_id]
             model_name = model._get_name()
-            print(f"Loading model {model_name} ({model_id}) in GPU")
-        if not self.pinInRAM: 
-            model.to("cuda")
+            print(f"Loading model {entry_name} ({model_name}) in GPU")
+ 
+
+        if self.async_transfers and blocks_name != None:
+            first = self.prev_blocks_names[entry_name] == None
+            next_blocks_entry = self.next_blocks_names[entry_name] if entry_name in self.next_blocks_names else None
+            if first:
+                cpu_to_gpu(torch.cuda.current_stream(), self.blocks_of_modules[entry_name])
+                # if next_blocks_entry != None:
+                #     self.transfer_stream.wait_stream(self.default_stream)
+            # else:
+            #     self.transfer_stream.wait_stream(self.default_stream)
+            torch.cuda.synchronize()
+
+            if next_blocks_entry != None:
+                cpu_to_gpu(self.transfer_stream, self.blocks_of_modules[next_blocks_entry]) #, self.default_stream
+
         else:
-            module_params = self.params_of_modules[model_id]            
-            for p in module_params:
+            # if self.async_transfers:
+            #     self.transfer_stream.wait_stream(self.default_stream)
+            cpu_to_gpu(self.default_stream, self.blocks_of_modules[entry_name])
+            torch.cuda.synchronize()
+ 
+
+    @torch.compiler.disable()  
+    def gpu_unload_blocks(self, model_id, blocks_name):
+        if blocks_name != None:
+            self.loaded_blocks[model_id] = None 
+
+        blocks_name = model_id if blocks_name is None else model_id + "/" + blocks_name
+
+        if self.verboseLevel >=2:
+            model = self.models[model_id]
+            model_name = model._get_name()
+            print(f"Unloading model {blocks_name} ({model_name}) from GPU")
+ 
+        blocks_params = self.blocks_of_modules[blocks_name]
+
+        if model_id in self.pinned_modules_data:
+            pinned_parameters_data = self.pinned_modules_data[model_id]
+            for p in blocks_params:
                 if isinstance(p, QTensor):
-                    p._data = p._data.cuda(non_blocking=True)             
-                    p._scale = p._scale.cuda(non_blocking=True)
+                    data = pinned_parameters_data[p]
+                    p._data = data[0]          
+                    p._scale = data[1]             
                 else:
-                    p.data = p.data.cuda(non_blocking=True) #
-        # torch.cuda.current_stream().synchronize()    
+                    p.data = pinned_parameters_data[p]
+        else:
+            for p in blocks_params:
+                if isinstance(p, QTensor):
+                    p._data = p._data.cpu()
+                    p._scale = p._scale.cpu()     
+                else:
+                    p.data = p.data.cpu()           
+
+
+
     @torch.compiler.disable()  
+    def gpu_load(self, model_id):
+        model = self.models[model_id]
+        self.active_models.append(model)
+        self.active_models_ids.append(model_id)
+
+        self.gpu_load_blocks(model_id, None)
+
+        # torch.cuda.current_stream().synchronize()    
+
     def unload_all(self):
-        for model, model_id in zip(self.active_models, self.active_models_ids):
-            if not self.pinInRAM:
-                model.to("cpu")
-            else:
-                module_params = self.params_of_modules[model_id]            
-                pinned_parameters_data = self.pinned_modules_data[model_id]
-                for p in module_params:
-                    if isinstance(p, QTensor):
-                        data = pinned_parameters_data[p]
-                        p._data = data[0]          
-                        p._scale = data[1]             
-                    else:
-                        p.data = pinned_parameters_data[p]            
-               
+        for model_id in self.active_models_ids:
+            self.gpu_unload_blocks(model_id, None)      
+            loaded_block = self.loaded_blocks[model_id]
+            if loaded_block != None:
+                self.gpu_unload_blocks(model_id, loaded_block)      
+                self.loaded_blocks[model_id] = None  
  
         self.active_models = []
         self.active_models_ids = []
+        self.active_subcaches = []
         torch.cuda.empty_cache()
         gc.collect()
+        self.last_reserved_mem_check = time.time()
 
     def move_args_to_gpu(self, *args, **kwargs):
         new_args= []
@@ -188,10 +391,12 @@ class offload:
         
         return new_args, new_kwargs
 
-    def ready_to_check_mem(self, forceMemoryCheck):
+    def ready_to_check_mem(self):
+        if self.compile:
+            return
         cur_clock = time.time()
         # can't check at each call if we can empty the cuda cache as quering the reserved memory value is a time consuming operation
-        if not forceMemoryCheck and (cur_clock - self.last_reserved_mem_check)<0.200:
+        if (cur_clock - self.last_reserved_mem_check)<0.200:
             return False
         self.last_reserved_mem_check = cur_clock
         return True        
@@ -199,20 +404,70 @@ class offload:
 
     def empty_cache_if_needed(self):
         mem_reserved = torch.cuda.memory_reserved()
-        if mem_reserved >= 0.9*self.device_mem_capacity:            
+        mem_threshold = 0.9*self.device_mem_capacity
+        if mem_reserved >= mem_threshold:            
             mem_allocated = torch.cuda.memory_allocated()
             if mem_allocated <= 0.70 * mem_reserved: 
                 # print(f"Cuda empty cache triggered as Allocated Memory ({mem_allocated/1024000:0f} MB) is lot less than Cached Memory ({mem_reserved/1024000:0f} MB)  ")
                 torch.cuda.empty_cache()
+                tm= time.time()
+                if self.verboseLevel >=2:
+                    print(f"Empty Cuda cache at {tm}")
                 # print(f"New cached memory after purge is {torch.cuda.memory_reserved()/1024000:0f} MB)  ")
 
-    def hook_me_light(self, target_module, forceMemoryCheck, previous_method):
-        def check_empty_cache(module, *args, **kwargs):
-            if self.ready_to_check_mem(forceMemoryCheck):
+
+    def any_param_or_buffer(self, target_module: torch.nn.Module):
+        
+        for _ in target_module.parameters(recurse= False):
+            return True
+        
+        for _ in target_module.buffers(recurse= False):
+            return True
+        
+        return False
+
+
+
+    def hook_me_light(self, target_module, model_id,blocks_name, previous_method, context):
+
+        anyParam = self.any_param_or_buffer(target_module)
+
+        def check_empty_cuda_cache(module, *args, **kwargs):
+            if self.ready_to_check_mem():
                 self.empty_cache_if_needed()
             return previous_method(*args, **kwargs) 
-  
-        setattr(target_module, "forward", functools.update_wrapper(functools.partial(check_empty_cache, target_module), previous_method) )
+
+
+        def load_module_blocks(module,  *args, **kwargs):
+            #some_context = context #for debugging
+            if blocks_name == None:
+                if self.ready_to_check_mem():
+                    self.empty_cache_if_needed()
+            else:                
+                loaded_block = self.loaded_blocks[model_id]
+                if (loaded_block == None or loaded_block != blocks_name) :
+                    if loaded_block != None:
+                        self.gpu_unload_blocks(model_id, loaded_block)
+                        if self.ready_to_check_mem():
+                            self.empty_cache_if_needed()
+                    self.loaded_blocks[model_id] = blocks_name
+                    self.gpu_load_blocks(model_id, blocks_name)
+            return previous_method(*args, **kwargs) 
+
+        if hasattr(target_module, "_mm_id"):
+            orig_model_id = getattr(target_module, "_mm_id")
+            if self.verboseLevel >=2:
+                print(f"Model '{model_id}' shares module '{target_module._get_name()}' with module '{orig_model_id}' ")
+            assert not anyParam
+            return
+        setattr(target_module, "_mm_id", model_id)
+
+
+        if  blocks_name != None and anyParam: 
+            setattr(target_module, "forward", functools.update_wrapper(functools.partial(load_module_blocks, target_module), previous_method) )
+            #print(f"new cache:{blocks_name}")
+        else:
+            setattr(target_module, "forward", functools.update_wrapper(functools.partial(check_empty_cuda_cache, target_module), previous_method) )
 
         
     def hook_me(self, target_module, model, model_id, module_id, previous_method):
@@ -236,13 +491,9 @@ class offload:
             return
         setattr(target_module, "_mm_id", model_id)
 
-        # 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)  
-        if not hasattr(target_module, "_hf_hook"):
-            setattr(target_module, "_hf_hook", HfHook())
         setattr(target_module, "forward", functools.update_wrapper(functools.partial(check_change_module, target_module), previous_method) )
 
-        if not self.verbose:
+        if not self.verboseLevel >=1:
             return
 
         if module_id == None or module_id =='':
@@ -262,22 +513,185 @@ class offload:
     #         self.unhook_module(module)
 
 
+    @staticmethod
+    def fast_load_transformers_model(model_path: str):
+        """
+        quick version of .LoadfromPretrained of  the transformers library
+        used to build a model and load the corresponding weights (quantized or not)
+        """       
+
+        from transformers import AutoConfig
+
+        if model_path.endswith(".sft") or model_path.endswith(".safetensors"):
+            config_path = model_path[ : model_path.rfind("/")]
+        else:
+            raise("full model path expected")
+        config_fullpath = config_path +"/config.json"
+
+        import os.path
+        if not os.path.isfile(config_fullpath):
+            raise("a 'config.json' that describes the model is required in the directory of the model")
+
+        with open(config_fullpath, "r", encoding="utf-8") as reader:
+            text = reader.read()
+        transformer_config= json.loads(text)
+        architectures = transformer_config["architectures"]
+        class_name = architectures[0] 
+
+        module = __import__("transformers")
+        transfomer_class = getattr(module, class_name)
+ 
+        config = AutoConfig.from_pretrained(config_path)     
+
+        from accelerate import init_empty_weights
+        #needed to keep inits of non persistent buffers
+        with init_empty_weights():
+            model = transfomer_class(config)
+        
+        model = model.base_model
+        torch.set_default_device('cpu')
+        model.apply(model._initialize_weights)
+
+        #missing_keys, unexpected_keys =  
+        offload.load_model_data(model,model_path, strict = True )
+
+        return model
+        # #            text_encoder.final_layer_norm = text_encoder.norm
+    #     model = model.base_model
+    #     model.final_layer_norm = model.norm
+    #     self.model = model
+
+
+
+    @staticmethod
+    def load_model_data(model, file_path: str, device=torch.device('cpu'), strict = True):
+        """
+        Load a model, detect if it has been previously quantized using quanto and do the extra setup if necessary
+        """
+        from optimum.quanto import requantize
+        import safetensors.torch 
+
+        if "quanto" in file_path.lower():
+            pos = str.rfind(file_path, ".")
+            if pos > 0:
+                quantization_map_path = file_path[:pos]
+            quantization_map_path += "_map.json"
+
+
+            with open(quantization_map_path, 'r') as f:
+                quantization_map = json.load(f)
+
+            state_dict = safetensors.torch.load_file(file_path)
+
+            # change dtype of current meta model parameters because 'requantize' won't update the dtype on non quantized parameters
+            for k, p in model.named_parameters():
+                if not k in quantization_map and k in state_dict:
+                    p_in_sd = state_dict[k] 
+                    if p.data.dtype != p_in_sd.data.dtype:
+                         p.data = p.data.to(p_in_sd.data.dtype)
+
+            requantize(model, state_dict, quantization_map, device)    
+
+            # for k, p in model.named_parameters():
+            #     if p.data.dtype == torch.float32:
+            #         pass        
+            
+
+            # del state_dict
+            return
+
+        else:
+            if ".safetensors" in file_path or ".sft" in file_path: 
+                state_dict = safetensors.torch.load_file(file_path)
+
+            else:
+
+                state_dict = torch.load(file_path, weights_only=True)
+                if "module" in state_dict:
+                    state_dict = state_dict["module"]
+
+
+        model.load_state_dict(state_dict, strict = strict,  assign = True ) #strict=True,
+
+
+        return
+
+    @staticmethod
+    def save_model(model, file_path, do_quantize = False, quantization_type = qint8 ):
+        """save the weights of a model and quantize them if requested
+        These weights can be loaded again using 'load_model_data'
+        """       
+        import safetensors.torch
+        pos = str.rfind(file_path, ".")
+        if pos > 0:
+            file_path = file_path[:pos]
+        
+        if do_quantize:
+            _quantize(model, weights=quantization_type)
+        
+# #        state_dict = {k: v.clone().contiguous() for k, v in model.state_dict().items()}
+#         state_dict = {k: v  for k, v in model.state_dict().items()}
+
+
+
+        safetensors.torch.save_file(model.state_dict(), file_path + '.safetensors')
+
+        if do_quantize:
+            from optimum.quanto import quantization_map
+
+            with open(file_path + '_map.json', 'w') as f:
+                json.dump(quantization_map(model), f)
+
   
 
     @classmethod
-    def all(cls, pipe_or_dict_of_modules, quantizeTransformer = True, pinInRAM = False,  verbose = True, modelsToQuantize = None ):
+    def all(cls, pipe_or_dict_of_modules, quantizeTransformer = True, pinInRAM = False,  verboseLevel = 1, modelsToQuantize = None, budgets= 0, info = None):
+        """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
+        pinInRAM: move models in reserved memor. This allows very fast performance but requires 50% extra RAM (usually >=64 GB)
+        modelsToQuantize: a list of models to be also quantized on the fly (e.g the text_encoder), useful to reduce bith RAM and VRAM consumption
+        budgets: 0 by default (unlimited). If non 0, it corresponds to the maximum size in MB that every model will occupy at any moment
+         (in fact the real usage is twice this number). It is very efficient to reduce VRAM consumption but this feature may be very slow
+         if pinInRAM is not enabled
+        """        
+        
         self = cls()
-        self.verbose = verbose
+        self.verboseLevel = verboseLevel
         self.pinned_modules_data = {}
+        model_budgets = {}
 
+#        model_budgets = {"text_encoder_2": 3400 }
+        HEADER = '\033[95m'
+        ENDC = '\033[0m'
+        BOLD ='\033[1m'
+        UNBOLD ='\033[0m'
+        
+        print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 2.0) by DeepBeepMeep ************{ENDC}{UNBOLD}")
+        if info != None:
+            print(info)
+        budget = 0
+        if not budgets is None:
+            if isinstance(budgets , dict):
+                model_budgets = budgets
+            else:
+                budget = int(budgets) * ONE_MB
+
+        if (budgets!= None or budget >0) :
+            self.async_transfers = True
+
+        pinInRAM = True
         # compile not working yet or slower
-        compile = False
-        self.pinInRAM = pinInRAM 
+        compile = False # True
+        #quantizeTransformer = False
+        #self.async_transfers = False
+        self.compile = compile
+
         pipe = None
-        preloadInRAM = True
         torch.set_default_device('cuda')
         if hasattr(pipe_or_dict_of_modules, "components"):
-            pipe_or_dict_of_modules.to("cpu") #XXXX
+            # commented as it not very useful and generates warnings
+            #pipe_or_dict_of_modules.to("cpu") #XXXX
             # create a fake Accelerate parameter so that lora loading doesn't change the device
             pipe_or_dict_of_modules.hf_device_map = torch.device("cuda")
             pipe = pipe_or_dict_of_modules
@@ -291,114 +705,178 @@ class offload:
             modelsToQuantize = [modelsToQuantize]
         if quantizeTransformer:
             modelsToQuantize.append("transformer")
+            
         self.models_to_quantize = modelsToQuantize
+        models_already_loaded = []
+
+        modelsToPin = None
+        pinAllModels = False
+        if isinstance(pinInRAM, bool):
+            pinAllModels = pinInRAM
+        elif isinstance(pinInRAM, list):            
+            modelsToPin = pinInRAM
+        else:
+            modelsToPin = [pinInRAM]
+
  #       del  models["transformer"] # to test everything but the transformer that has a much longer loading
- #       models = { 'transformer': pipe_or_dict_of_modules["transformer"]} # to test only the transformer
+        sizeofbfloat16 = torch.bfloat16.itemsize
+   #        
+   #     models = { 'transformer': pipe_or_dict_of_modules["transformer"]} # to test only the transformer
+
+
         for model_id in models: 
             current_model: torch.nn.Module = models[model_id] 
+            modelPinned = pinAllModels or (modelsToPin != None and model_id in modelsToPin)
             # make sure that no RAM or GPU memory is not allocated for gradiant / training
-            current_model.to("cpu").eval() #XXXXX
-
+            current_model.to("cpu").eval()
+            already_loaded = False
             # Quantize model just before transferring it to the RAM to keep OS cache file
             # open as short as possible. Indeed it seems that as long as the lazy safetensors 
             # are not fully fully loaded, the OS won't be able to release the corresponding cache file in RAM.
             if model_id in self.models_to_quantize:
-                print(f"Quantization of model '{model_id}' started")
-                quantize(current_model, weights=qint8)
-                freeze(current_model)
-                print(f"Quantization of model '{model_id}' done")
-                torch.cuda.empty_cache()
-                gc.collect()         
 
+                already_quantized = _quantize(current_model, weights=qint8, verboseLevel = self.verboseLevel, model_id=model_id)
+                if not already_quantized:
+                    already_loaded = True
+                    models_already_loaded.append(model_id) 
 
-    
-            if preloadInRAM: # 
-                # load all the remaining unread lazy safetensors in RAM to free open cache files 
-                for p in current_model.parameters():
-                    # Preread every tensor in RAM except tensors that have just been quantified
-                    # and are no longer needed
-                    if isinstance(p, QTensor):
-                        # fix quanto bug (see below) now as he won't have any opportunity to do it during RAM pinning 
-                        if not pinInRAM and p._scale.dtype == torch.float32:
-                            p._scale = p._scale.to(torch.bfloat16) 
 
+            current_model_size = 0    
+            # load all the remaining unread lazy safetensors in RAM to free open cache files 
+            for p in current_model.parameters():
+                # Preread every tensor in RAM except tensors that have just been quantified
+                # and are no longer needed
+                if isinstance(p, QTensor):
+                    # fix quanto bug (see below) now as he won't have any opportunity to do it during RAM pinning 
+                    if not modelPinned and p._scale.dtype == torch.float32:
+                        p._scale = p._scale.to(torch.bfloat16) 
+                    current_model_size +=  torch.numel(p._scale) * sizeofbfloat16
+                    current_model_size +=  torch.numel(p._data) * sizeofbfloat16 / 2
+                    if pinInRAM and not already_loaded:
+                        # Force flushing the lazy load so that reserved memory can be freed when we are ready to pin
+                        p._scale = p._scale + 0
+                        p._data = p._data + 0
+                else:
+                    if 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)
                     else:
-                        if 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)
-                        else:
-                            # force reading the tensors from the disk by pretending to modify them
-                            p.data = p.data + 0
-                     
+                        # force reading the tensors from the disk by pretending to modify them
+                        p.data = p.data + 0
+
+                    current_model_size +=  torch.numel(p.data) * p.data.element_size()
+                             
+            for b in current_model.buffers():
+                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)
+                else:
+                    # force reading the tensors from the disk by pretending to modify them
+                    b.data = b.data + 0
+
+                current_model_size +=  torch.numel(p.data) * p.data.element_size()
+
+            if model_id not in self.models:
+                self.models[model_id] = current_model
+
+            
+            model_budget = model_budgets[model_id] * ONE_MB if model_id in model_budgets else budget
+
+            if  model_budget > 0 and model_budget > current_model_size:
+                model_budget = 0
+            
+            model_budgets[model_id] = model_budget
+
+        # Pin in RAM models only once they have been fully loaded otherwise there will be some contention (at least on Linux OS) in the non pageable memory
+        # between partially loaded lazy safetensors and pinned tensors   
+        for model_id in models: 
+            current_model: torch.nn.Module = models[model_id] 
+            if not (pinAllModels or modelsToPin != None and model_id in modelsToPin):
+                continue
+            if verboseLevel>=1:
+                print(f"Pinning tensors of '{model_id}' in RAM")
+            gc.collect()                     
+            pinned_parameters_data = {}
+            for p in current_model.parameters():
+                if isinstance(p, QTensor):
+                    # pin in memory both quantized data and scales of quantized parameters
+                    # but don't pin .data as it corresponds to the original tensor that we don't want to reload
+                    p._data = p._data.pin_memory()
+                    # fix quanto bug (that seems to have been fixed since&) that allows _scale to be float32 if the original weight was float32 
+                    # (this may cause type mismatch between dequantified bfloat16 weights and float32 scales)
+                    p._scale = p._scale.to(torch.bfloat16).pin_memory() if p._scale.dtype == torch.float32 else p._scale.pin_memory()
+                    pinned_parameters_data[p]=[p._data, p._scale]
+                else:
+                    p.data = p.data.pin_memory() 
+                    pinned_parameters_data[p]=p.data 
+            for b in current_model.buffers():
+                b.data = b.data.pin_memory()
+
+            pinned_buffers_data = {b: b.data for b in current_model.buffers()}
+            pinned_parameters_data.update(pinned_buffers_data)
+            self.pinned_modules_data[model_id]=pinned_parameters_data
 
-            addModelFlag  = False
 
-            current_block_sequence = None
+        #  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
+             
             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)  
+                if  not hasattr(submodule, "_hf_hook"):
+                    setattr(submodule, "_hf_hook", HfHook())
+
+                if submodule_name=='':
+                    continue
+
+                if current_budget > 0:
+                    if isinstance(submodule, (torch.nn.ModuleList, torch.nn.Sequential)):
+                        if cur_blocks_prefix == None:
+                            cur_blocks_prefix = submodule_name + "."
+                        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 cur_blocks_prefix == submodule_name[0:len(cur_blocks_prefix)]:
+                            if submodule_name.startswith(cur_blocks_prefix):
+                                num = int(submodule_name[len(cur_blocks_prefix):].split(".")[0])
+                                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
+
                 if hasattr(submodule, "forward"):
                     submodule_method = getattr(submodule, "forward")
                     if callable(submodule_method):   
-                        addModelFlag  = True 
-                        if submodule_name=='' or len(submodule_name.split("."))==1:
-                            # hook only the first two levels of  modules with the full suite of processing                             
+                        if len(submodule_name.split("."))==1:
+                            # hook only the first level of modules with the full suite of processing                             
                             self.hook_me(submodule, current_model, model_id, submodule_name, submodule_method)
-                        else:
-                            forceMemoryCheck = False
-                            pos = submodule_name.find(".0.") 
-                            if pos > 0:
-                                if current_block_sequence ==  None:
-                                    new_candidate = submodule_name[0:pos+3]
-                                    if len(new_candidate.split("."))<=4:
-                                        current_block_sequence = new_candidate  
-                                        # force a memory check when initiating a new sequence of blocks as the shapes of tensor will certainly change
-                                        # and memory reusability is less likely
-                                        # we limit this check to the first level of blocks as quering the cuda cache is time consuming
-                                        forceMemoryCheck = True
-                                else:
-                                    if current_block_sequence != submodule_name[0:len(current_block_sequence)]:
-                                         current_block_sequence = None
-                            self.hook_me_light(submodule, forceMemoryCheck, submodule_method)
-
-                            
-            if addModelFlag:
-                if model_id not in self.models:
-                    self.models[model_id] = current_model
-
-        # Pin in RAM models only once they have been fully loaded otherwise there may be some contention in the non pageable memory
-        # between partially loaded lazy safetensors and pinned tensors   
-        if pinInRAM: 
-            if verbose:
-                print("Pinning model tensors in RAM")
-            torch.cuda.empty_cache()
-            gc.collect()                     
-            for model_id in models: 
-                pinned_parameters_data = {}
-                current_model: torch.nn.Module = models[model_id] 
-                for p in current_model.parameters():
-                    if isinstance(p, QTensor):
-                        # pin in memory both quantized data and scales of quantized parameters
-                        # but don't pin .data as it corresponds to the original tensor that we don't want to reload
-                        p._data = p._data.pin_memory()
-                        # fix quanto bug that allows _scale to be float32 if the original weight was float32 
-                        # (this may cause type mismatch between dequantified bfloat16 weights and float32 scales)
-                        p._scale = p._scale.to(torch.bfloat16).pin_memory() if p._scale.dtype == torch.float32 else p._scale.pin_memory()
-                        pinned_parameters_data[p]=[p._data, p._scale]
-                    else:
-                        p.data = p.data.pin_memory() 
-                        pinned_parameters_data[p]=p.data 
-                for b in current_model.buffers():
-                    b.data = b.data.pin_memory()
+                        else: 
+                            # force a memory check when initiating a new sequence of blocks as the shapes of tensor will certainly change
+                            # and memory reusability is less likely
+                            # we limit this check to the first level of blocks as quering the cuda cache is time consuming
+                            self.hook_me_light(submodule, model_id, cur_blocks_name, submodule_method, context = submodule_name)
 
-                pinned_buffers_data = {b: b.data for b in current_model.buffers()}
-                pinned_parameters_data.update(pinned_buffers_data)
-                self.pinned_modules_data[model_id]=pinned_parameters_data
+                        if compile and cur_blocks_name != None and model_id == "transformer" and "_blocks" in submodule_name:
+                             submodule.compile(mode="reduce-overhead" ) #mode= "max-autotune"     
 
-                module_params = []      
-                self.params_of_modules[model_id] = module_params           
-                self.collect_module_parameters(current_model,module_params)
+                        current_size = self.add_module_to_blocks(model_id, cur_blocks_name, submodule, prev_blocks_name)
 
-        if compile:
-            if verbose:
+
+        if compile and False:
+            if verboseLevel>=1:
                 print("Torch compilation started")
             torch._dynamo.config.cache_size_limit = 10000
             # if pipe != None and hasattr(pipe, "__call__"):
@@ -409,13 +887,65 @@ class offload:
                     current_model.compile(mode= "max-autotune")                                 
             #models["transformer"].compile()
                 
-            if verbose:
+            if verboseLevel>=1:
                 print("Torch compilation done")
 
+        if verboseLevel >=2:
+            for n,b in self.blocks_of_modules_sizes.items():
+                print(f"Size of submodel '{n}': {b/ONE_MB:.1f} MB")
+
         torch.cuda.empty_cache()
         gc.collect()         
 
-
         return self
 
-            
+
+
+    @staticmethod
+    def profile(pipe_or_dict_of_modules,profile_no: profile_type, quantizeTransformer = True):
+        """Apply a configuration profile that depends on your hardware:
+        pipe_or_dict_of_modules : the pipeline object or a dictionary of modules of the model
+        profile_name : num of the profile:
+            HighRAM_HighVRAM_Fastest (=1): at least 48 GB of RAM and 24 GB of VRAM : the fastest well suited for a RTX 3090 / RTX 4090
+            HighRAM_LowVRAM_Fast (=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_Medium (=3): at least 32 GB of RAM and 24 GB of VRAM : so so speed but adapted for RTX 3090 / RTX 4090 with limited RAM
+            LowRAM_LowVRAM_Slow (=4): at least 32 GB of RAM and 12 GB of VRAM : if have little VRAM or generate longer videos 
+            VerylowRAM_LowVRAM_Slowest (=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
+        quantizeTransformer: bool = True, the main model is quantized by default for all the profiles, you may want to disable that to get the best image quality
+        """      
+
+
+        modules = pipe_or_dict_of_modules
+        if hasattr(modules, "components"):
+            modules= modules.components 
+        any_T5 = False
+        if "text_encoder_2" in modules:
+            text_encoder_2 =  modules["text_encoder_2"]
+            any_T5 = "t5" in text_encoder_2.__module__.lower()
+        extra_mod_to_quantize = ("text_encoder_2" if any_T5 else "text_encoder")  
+
+        # transformer (video or image generator) should be as small as possible to not occupy space that could be used by actual image data
+        # on the other hand the text encoder should be quite large (as long as it fits in 10 GB of VRAM) to reduce sequence offloading
+
+        budgets = { "transformer" : 600 , "text_encoder": 3000, "text_encoder_2": 3000 }
+
+        if profile_no == profile_type.HighRAM_HighVRAM_Fastest:
+            info = "You have chosen a Very Fast profile that requires at least 48 GB of RAM and 24 GB of VRAM."
+            return offload.all(pipe_or_dict_of_modules, pinInRAM= True, info = info, quantizeTransformer= quantizeTransformer)
+        elif profile_no == profile_type.HighRAM_LowVRAM_Fast:
+            info = "You have chosen a Fast profile that requires at least 48 GB of RAM and 12 GB of VRAM."
+            return offload.all(pipe_or_dict_of_modules, pinInRAM= True, budgets=budgets, info = info, quantizeTransformer= quantizeTransformer )
+        elif profile_no == profile_type.LowRAM_HighVRAM_Medium:
+            info = "You have chosen a Medium speed profile that requires at least 32 GB of RAM and 24 GB of VRAM."
+            return offload.all(pipe_or_dict_of_modules, pinInRAM= "transformer", modelsToQuantize= extra_mod_to_quantize ,  info = info, quantizeTransformer= quantizeTransformer)
+        elif profile_no == profile_type.LowRAM_LowVRAM_Slow:
+            info = "You have chosen the Slow profile that requires at least 32 GB of RAM and 12 GB of VRAM."
+            return offload.all(pipe_or_dict_of_modules, pinInRAM= "transformer", modelsToQuantize= extra_mod_to_quantize ,  budgets=budgets, info = info, quantizeTransformer= quantizeTransformer)
+        elif profile_no == profile_type.VerylowRAM_LowVRAM_Slowest:
+            budgets["transformer"] = 400
+            info = "You have chosen the Slowest profile that requires at least 24 GB of RAM and 10 GB of VRAM."
+            return offload.all(pipe_or_dict_of_modules, pinInRAM= False, modelsToQuantize= extra_mod_to_quantize ,  budgets=budgets, info = info, quantizeTransformer= quantizeTransformer)
+        else:
+            raise("Unknown profile")
+

mmgp-1.2.0.dist-info/METADATA

@@ -1,109 +0,0 @@
-Metadata-Version: 2.1
-Name: mmgp
-Version: 1.2.0
-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
-
-
-<p align="center">
-  <H2>Memory Management 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 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:
-- GPU: RTX 3090/ RTX 4090 (24 GB of VRAM)
-- RAM: minimum 48 GB, recommended 64 GB 
-
-## Usage 
-First you need to install the module in your current project with:
-```shell
-pip install mmgp
-```
-
-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:
-
-```
-  from mmgp import offload
-  offload.all(pipe)
-```  
-
-## Options
-The 'transformer' model in the pipe contains usually the video or image generator is 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. In that case you need to set the option *quantizeTransformer* to *False* to turn off on the fly quantization.
-
-You can specify a list of additional models string ids to quantize (for instance the text_encoder) using the optional argument *modelsToQuantize* for instance *modelsToQuantize = ["text_encoder_2"]*.This may be useful if you have less than 48 GB of RAM.
-
-Note that there is little advantage on the GPU / VRAM side to quantize text encoders as their inputs are usually quite light. 
-
-Conversely if you have more than 64GB of RAM you may want to enable RAM pinning with the option *pinInRAM = True*. You will get in return super fast loading / unloading of models
-(this can save significant time if the same pipeline is run multiple times in a row)
-
-In Summary, if you have:
-- Between 32 GB and 48 GB of RAM
-```
-  offload.all(pipe, modelsToQuantize = ["text_encoder_2"]) # for Flux models
-  #OR
-  offload.all(pipe, modelsToQuantize = ["text_encoder"]) # for HunyuanVideo models
-
-```  
-
-- Between 48 GB and 64 GB of RAM
-```
-  offload.all(pipe)
-```  
-- More than 64 GB of RAM
-```
-  offload.all(pipe), pinInRAM = True
-```
-
-## Special
-Sometime there isn't an explicit pipe object as each submodel is loaded separately in the main app. If this is the case, you need to 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 that there should be always one model whose Id is 'transformer'. It 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).
-
-Becareful, 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-1.2.0.dist-info/RECORD

@@ -1,7 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mmgp.py,sha256=IijgE22bUPl98VvXwoC1qngmkdWU11YXjkiksp8o1hY,21418
-mmgp-1.2.0.dist-info/LICENSE.md,sha256=HjzvY2grdtdduZclbZ46B2M-XpT4MDCxFub5ZwTWq2g,93
-mmgp-1.2.0.dist-info/METADATA,sha256=jRXi-iNZ_3zNNVxMC1qmVDd7ylq8kAr5Y5FgYyBvVh4,4897
-mmgp-1.2.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-mmgp-1.2.0.dist-info/top_level.txt,sha256=waGaepj2qVfnS2yAOkaMu4r9mJaVjGbEi6AwOUogU_U,14
-mmgp-1.2.0.dist-info/RECORD,,