mmgp 1.0.0__py3-none-any.whl

mmgp/__init__.py

@@ -0,0 +1,13 @@
+try:
+    from ._version import (
+        version as __version__,  # type: ignore
+        version_tuple,
+    )
+except ImportError:
+    __version__ = "unknown (no version information available)"
+    version_tuple = (0, 0, "unknown", "noinfo")
+
+from pathlib import Path
+
+PACKAGE = __package__.replace("_", "-")
+PACKAGE_ROOT = Path(__file__).parent

mmgp/_version.py

@@ -0,0 +1,16 @@
+# file generated by setuptools_scm
+# don't change, don't track in version control
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple, Union
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '1.0.0'
+__version_tuple__ = version_tuple = (1, 0, 0)

mmgp/mmgp.py

@@ -0,0 +1,405 @@
+# ------------------ Memory Management for the GPU Poor by DeepBeepMeep ------------------
+#
+# This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  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
+# - RAM: minimum 48 GB, recommended 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:
+#   from mmgp import offload
+#   offload.me(pipe)
+# If you don't have enough RAM you may disable RAM pinning but model switching option pinInRAM= False
+# 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 is corresponds to the main image / video model which usually needs to be quantized (this is done by default)
+# 
+# Becareful, lots of models uses the T5 XXL as a text encoder. However, quite often their corresponding pipeline configuratons points 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)
+#
+# You are free to use my code as long you give me proper credits. You may contact me on twitter @deepbeepmeep
+#
+# Credits 
+# -------
+# Huggingface / accelerate for the hooking examples
+# Huggingface / quanto for their very useful quantizer
+# gau-nernst for his Pinnig RAM examples
+
+
+#
+import torch
+#
+import gc
+import time
+import functools
+from optimum.quanto import freeze, qfloat8, qint8, quantize, QModuleMixin, QTensor
+
+
+# config Dimension X (CogVideo derived ) : Quantization False: because Lora applied later
+
+
+cotenants_map = { 
+                             "text_encoder": ["vae", "text_encoder_2"],
+                             "text_encoder_2": ["vae", "text_encoder"],                             
+                             }
+
+# useful functions to move a group of tensors (to design custom offload patches)
+def move_tensors(obj, device):
+    if torch.is_tensor(obj):
+        return obj.to(device)
+    elif isinstance(obj, dict):
+        _dict = {}
+        for k, v in obj.items():
+            _dict[k] = move_tensors(v, device)
+        return _dict
+    elif isinstance(obj, list):
+        _list = []
+        for v in obj:
+            _list.append(move_tensors(v, device))
+        return _list
+    else:
+        raise TypeError("Tensor or list / dict of tensors expected")
+
+
+def get_model_name(model):
+    return model.name
+
+class HfHook:
+    def __init__(self):
+        self.execution_device = "cuda"
+
+    def detach_hook(self, module):
+        pass
+
+class offload:
+    def __init__(self):
+        self.active_models = []
+        self.active_models_ids = []
+        self.models = {}
+        self.verbose = False
+        self.models_to_quantize = []
+        self.pinned_modules_data = {}
+        self.params_of_modules = {}
+        self.pinTensors = False
+        self.device_mem_capacity = torch.cuda.get_device_properties(0).total_memory
+        self.last_reserved_mem_check =0
+
+    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())
+        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)
+
+    def can_model_be_cotenant(self, model_id):
+        potential_cotenants= cotenants_map.get(model_id, None)
+        if potential_cotenants is None: 
+            return False
+        for existing_cotenant in self.active_models_ids:
+            if existing_cotenant not in potential_cotenants: 
+                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:
+            model_name = model._get_name()
+            print(f"Loading model {model_name} ({model_id}) in GPU")
+        if not self.pinInRAM: 
+            model.to("cuda")
+        else:
+            module_params = self.params_of_modules[model_id]            
+            for p in module_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) #
+        # 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]            
+               
+ 
+        self.active_models = []
+        self.active_models_ids = []
+        torch.cuda.empty_cache()
+        gc.collect()
+
+    def move_args_to_gpu(self, *args, **kwargs):
+        new_args= []
+        new_kwargs={}
+        for arg in args:
+            if torch.is_tensor(arg):    
+                if arg.dtype == torch.float32:
+                    arg = arg.to(torch.bfloat16).cuda(non_blocking=True)             
+                else:
+                    arg = arg.cuda(non_blocking=True)             
+            new_args.append(arg)
+
+        for k in kwargs:
+            arg = kwargs[k]
+            if torch.is_tensor(arg):
+                if arg.dtype == torch.float32:
+                    arg = arg.to(torch.bfloat16).cuda(non_blocking=True)             
+                else:
+                    arg = arg.cuda(non_blocking=True)             
+            new_kwargs[k]= arg
+        
+        return new_args, new_kwargs
+
+    def ready_to_check_mem(self, forceMemoryCheck):
+        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:
+            return False
+        self.last_reserved_mem_check = cur_clock
+        return True        
+
+
+    def empty_cache_if_needed(self):
+        mem_reserved = torch.cuda.memory_reserved()
+        if mem_reserved >= 0.9*self.device_mem_capacity:            
+            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()
+                # 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):
+   #     @torch.compiler.disable()        
+        def check_empty_cache(module, *args, **kwargs):
+            if self.ready_to_check_mem(forceMemoryCheck):
+                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 hook_me(self, target_module, model, model_id, module_id, previous_method):
+        @torch.compiler.disable()        
+        def check_change_module(module, *args, **kwargs):
+            performEmptyCacheTest = False
+            if not model_id in self.active_models_ids:
+                new_model_id = getattr(module, "_mm_id") 
+                # do not always unload existing models if it is more efficient to keep in them in the GPU 
+                # (e.g: small modules whose calls are text encoders) 
+                if not self.can_model_be_cotenant(new_model_id) :
+                    self.unload_all()
+                    performEmptyCacheTest = False
+                self.gpu_load(new_model_id)
+            # transfer leftovers inputs that were incorrectly created in the RAM (mostly due to some .device tests that returned incorrectly "cpu")
+            args, kwargs = self.move_args_to_gpu(*args, **kwargs)
+            if performEmptyCacheTest:
+                self.empty_cache_if_needed()
+            return previous_method(*args, **kwargs) 
+  
+        if hasattr(target_module, "_mm_id"):
+            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:
+            return
+
+        if module_id == None or module_id =='':
+            model_name = model._get_name()
+            print(f"Hooked in model {model_name} ({model_id})")
+
+
+    # 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)
+
+
+  
+
+    @classmethod
+    def all(cls, pipe_or_dict_of_modules, quantizeTransformer = True, pinInRAM = True, compile= True, verbose = True):
+        self = cls()
+        self.verbose = verbose
+        self.pinned_modules_data = {}
+
+        self.pinInRAM = pinInRAM 
+
+        preloadInRAM = True
+        torch.set_default_device('cuda')
+        if hasattr(pipe_or_dict_of_modules, "components"):
+            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_or_dict_of_modules= pipe_or_dict_of_modules.components 
+        
+        
+        
+        models = {k: v for k, v in pipe_or_dict_of_modules.items() if isinstance(v, torch.nn.Module)}
+
+
+        if quantizeTransformer:
+            self.models_to_quantize = ["transformer"]
+ #       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
+        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
+            current_model.to("cpu").eval() #XXXXX
+
+            # 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()         
+
+
+    
+            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) 
+
+                    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
+                     
+
+            addModelFlag  = False
+
+            current_block_sequence = None
+            for submodule_name, submodule in current_model.named_modules():  
+                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                             
+                            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()
+
+                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
+
+                module_params = []      
+                self.params_of_modules[model_id] = module_params           
+                self.collect_module_parameters(current_model,module_params)
+
+        if compile:
+            if verbose:
+                print("Torch compilation started")
+            torch._dynamo.config.cache_size_limit = 10000
+            for model_id in models: 
+                    current_model: torch.nn.Module = models[model_id]
+                    current_model.compile()                                 
+            #models["transformer"].compile()        
+            if verbose:
+                print("Torch compilation done")
+
+        torch.cuda.empty_cache()
+        gc.collect()         
+
+
+        return self
+
+            

mmgp-1.0.0.dist-info/LICENSE.md

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

mmgp-1.0.0.dist-info/METADATA

@@ -0,0 +1,249 @@
+Metadata-Version: 2.1
+Name: mmgp
+Version: 1.0.0
+Summary: Memory Management for the GPU Poor
+Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
+License:                                  Apache License
+                                   Version 2.0, January 2004
+                                http://www.apache.org/licenses/
+        
+           TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+        
+           1. Definitions.
+        
+              "License" shall mean the terms and conditions for use, reproduction,
+              and distribution as defined by Sections 1 through 9 of this document.
+        
+              "Licensor" shall mean the copyright owner or entity authorized by
+              the copyright owner that is granting the License.
+        
+              "Legal Entity" shall mean the union of the acting entity and all
+              other entities that control, are controlled by, or are under common
+              control with that entity. For the purposes of this definition,
+              "control" means (i) the power, direct or indirect, to cause the
+              direction or management of such entity, whether by contract or
+              otherwise, or (ii) ownership of fifty percent (50%) or more of the
+              outstanding shares, or (iii) beneficial ownership of such entity.
+        
+              "You" (or "Your") shall mean an individual or Legal Entity
+              exercising permissions granted by this License.
+        
+              "Source" form shall mean the preferred form for making modifications,
+              including but not limited to software source code, documentation
+              source, and configuration files.
+        
+              "Object" form shall mean any form resulting from mechanical
+              transformation or translation of a Source form, including but
+              not limited to compiled object code, generated documentation,
+              and conversions to other media types.
+        
+              "Work" shall mean the work of authorship, whether in Source or
+              Object form, made available under the License, as indicated by a
+              copyright notice that is included in or attached to the work
+              (an example is provided in the Appendix below).
+        
+              "Derivative Works" shall mean any work, whether in Source or Object
+              form, that is based on (or derived from) the Work and for which the
+              editorial revisions, annotations, elaborations, or other modifications
+              represent, as a whole, an original work of authorship. For the purposes
+              of this License, Derivative Works shall not include works that remain
+              separable from, or merely link (or bind by name) to the interfaces of,
+              the Work and Derivative Works thereof.
+        
+              "Contribution" shall mean any work of authorship, including
+              the original version of the Work and any modifications or additions
+              to that Work or Derivative Works thereof, that is intentionally
+              submitted to Licensor for inclusion in the Work by the copyright owner
+              or by an individual or Legal Entity authorized to submit on behalf of
+              the copyright owner. For the purposes of this definition, "submitted"
+              means any form of electronic, verbal, or written communication sent
+              to the Licensor or its representatives, including but not limited to
+              communication on electronic mailing lists, source code control systems,
+              and issue tracking systems that are managed by, or on behalf of, the
+              Licensor for the purpose of discussing and improving the Work, but
+              excluding communication that is conspicuously marked or otherwise
+              designated in writing by the copyright owner as "Not a Contribution."
+        
+              "Contributor" shall mean Licensor and any individual or Legal Entity
+              on behalf of whom a Contribution has been received by Licensor and
+              subsequently incorporated within the Work.
+        
+           2. Grant of Copyright License. Subject to the terms and conditions of
+              this License, each Contributor hereby grants to You a perpetual,
+              worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+              copyright license to reproduce, prepare Derivative Works of,
+              publicly display, publicly perform, sublicense, and distribute the
+              Work and such Derivative Works in Source or Object form.
+        
+           3. Grant of Patent License. Subject to the terms and conditions of
+              this License, each Contributor hereby grants to You a perpetual,
+              worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+              (except as stated in this section) patent license to make, have made,
+              use, offer to sell, sell, import, and otherwise transfer the Work,
+              where such license applies only to those patent claims licensable
+              by such Contributor that are necessarily infringed by their
+              Contribution(s) alone or by combination of their Contribution(s)
+              with the Work to which such Contribution(s) was submitted. If You
+              institute patent litigation against any entity (including a
+              cross-claim or counterclaim in a lawsuit) alleging that the Work
+              or a Contribution incorporated within the Work constitutes direct
+              or contributory patent infringement, then any patent licenses
+              granted to You under this License for that Work shall terminate
+              as of the date such litigation is filed.
+        
+           4. Redistribution. You may reproduce and distribute copies of the
+              Work or Derivative Works thereof in any medium, with or without
+              modifications, and in Source or Object form, provided that You
+              meet the following conditions:
+        
+              (a) You must give any other recipients of the Work or
+                  Derivative Works a copy of this License; and
+        
+              (b) You must cause any modified files to carry prominent notices
+                  stating that You changed the files; and
+        
+              (c) You must retain, in the Source form of any Derivative Works
+                  that You distribute, all copyright, patent, trademark, and
+                  attribution notices from the Source form of the Work,
+                  excluding those notices that do not pertain to any part of
+                  the Derivative Works; and
+        
+              (d) If the Work includes a "NOTICE" text file as part of its
+                  distribution, then any Derivative Works that You distribute must
+                  include a readable copy of the attribution notices contained
+                  within such NOTICE file, excluding those notices that do not
+                  pertain to any part of the Derivative Works, in at least one
+                  of the following places: within a NOTICE text file distributed
+                  as part of the Derivative Works; within the Source form or
+                  documentation, if provided along with the Derivative Works; or,
+                  within a display generated by the Derivative Works, if and
+                  wherever such third-party notices normally appear. The contents
+                  of the NOTICE file are for informational purposes only and
+                  do not modify the License. You may add Your own attribution
+                  notices within Derivative Works that You distribute, alongside
+                  or as an addendum to the NOTICE text from the Work, provided
+                  that such additional attribution notices cannot be construed
+                  as modifying the License.
+        
+              You may add Your own copyright statement to Your modifications and
+              may provide additional or different license terms and conditions
+              for use, reproduction, or distribution of Your modifications, or
+              for any such Derivative Works as a whole, provided Your use,
+              reproduction, and distribution of the Work otherwise complies with
+              the conditions stated in this License.
+        
+           5. Submission of Contributions. Unless You explicitly state otherwise,
+              any Contribution intentionally submitted for inclusion in the Work
+              by You to the Licensor shall be under the terms and conditions of
+              this License, without any additional terms or conditions.
+              Notwithstanding the above, nothing herein shall supersede or modify
+              the terms of any separate license agreement you may have executed
+              with Licensor regarding such Contributions.
+        
+           6. Trademarks. This License does not grant permission to use the trade
+              names, trademarks, service marks, or product names of the Licensor,
+              except as required for reasonable and customary use in describing the
+              origin of the Work and reproducing the content of the NOTICE file.
+        
+           7. Disclaimer of Warranty. Unless required by applicable law or
+              agreed to in writing, Licensor provides the Work (and each
+              Contributor provides its Contributions) on an "AS IS" BASIS,
+              WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+              implied, including, without limitation, any warranties or conditions
+              of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+              PARTICULAR PURPOSE. You are solely responsible for determining the
+              appropriateness of using or redistributing the Work and assume any
+              risks associated with Your exercise of permissions under this License.
+        
+           8. Limitation of Liability. In no event and under no legal theory,
+              whether in tort (including negligence), contract, or otherwise,
+              unless required by applicable law (such as deliberate and grossly
+              negligent acts) or agreed to in writing, shall any Contributor be
+              liable to You for damages, including any direct, indirect, special,
+              incidental, or consequential damages of any character arising as a
+              result of this License or out of the use or inability to use the
+              Work (including but not limited to damages for loss of goodwill,
+              work stoppage, computer failure or malfunction, or any and all
+              other commercial damages or losses), even if such Contributor
+              has been advised of the possibility of such damages.
+        
+           9. Accepting Warranty or Additional Liability. While redistributing
+              the Work or Derivative Works thereof, You may choose to offer,
+              and charge a fee for, acceptance of support, warranty, indemnity,
+              or other liability obligations and/or rights consistent with this
+              License. However, in accepting such obligations, You may act only
+              on Your own behalf and on Your sole responsibility, not on behalf
+              of any other Contributor, and only if You agree to indemnify,
+              defend, and hold each Contributor harmless for any liability
+              incurred by, or claims asserted against, such Contributor by reason
+              of your accepting any such warranty or additional liability.
+        
+           END OF TERMS AND CONDITIONS
+        
+           APPENDIX: How to apply the Apache License to your work.
+        
+              To apply the Apache License to your work, attach the following
+              boilerplate notice, with the fields enclosed by brackets "[]"
+              replaced with your own identifying information. (Don't include
+              the brackets!)  The text should be enclosed in the appropriate
+              comment syntax for the file format. We also recommend that a
+              file or class name and description of purpose be included on the
+              same "printed page" as the copyright notice for easier
+              identification within third-party archives.
+        
+           Copyright [yyyy] [name of copyright owner]
+        
+           Licensed under the Apache License, Version 2.0 (the "License");
+           you may not use this file except in compliance with the License.
+           You may obtain a copy of the License at
+        
+               http://www.apache.org/licenses/LICENSE-2.0
+        
+           Unless required by applicable law or agreed to in writing, software
+           distributed under the License is distributed on an "AS IS" BASIS,
+           WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+           See the License for the specific language governing permissions and
+           limitations under the License.
+        
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE.md
+Requires-Dist: torch>=2.1.0
+Requires-Dist: optimum-quanto
+
+**------------------ Memory Management for the GPU Poor by DeepBeepMeep ------------------**
+
+This module contains multiples optimisations so that models such as Flux (and derived), Mochi, CogView, HunyuanVideo, ...  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
+- RAM: minimum 48 GB, recommended 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:
+  from mmgp import offload
+  offload.me(pipe)
+If you don't have enough RAM you may disable RAM pinning but model switching option pinInRAM= False
+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 is corresponds to the main image / video model which usually needs to be quantized (this is done by default)
+
+Becareful, lots of models uses the T5 XXL as a text encoder. However, quite often their corresponding pipeline configuratons points 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)
+
+You are free to use my code as long you give me proper credits. You may contact me on twitter @deepbeepmeep
+
+Credits 
+-------
+Huggingface / accelerate for the hooking examples
+Huggingface / quanto for their very useful quantizer
+gau-nernst for his Pinnig RAM examples

mmgp-1.0.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+mmgp/__init__.py,sha256=N-ewQSpv5s0q9LUWWMY8pf13pLoUVsrFLk4ueOjqSbM,345
+mmgp/_version.py,sha256=DGJ4pj32xs3_DRJhSzQwCiRNnAQrMgo09USYpyMZsKc,411
+mmgp/mmgp.py,sha256=aMVV3jeB0zTq0NCJDphUFeW8_6CTqqWCKpmMdMnEjK8,19596
+mmgp-1.0.0.dist-info/LICENSE.md,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mmgp-1.0.0.dist-info/METADATA,sha256=IPg_YFm2mD3MkGFLRui5IDtlrjjHiqE7_MU96Q4FGQ4,15791
+mmgp-1.0.0.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+mmgp-1.0.0.dist-info/top_level.txt,sha256=nPtltlwG1fODvMY43SuGiMhijSwrtBTrp1E2-xiv49U,5
+mmgp-1.0.0.dist-info/RECORD,,

mmgp-1.0.0.dist-info/WHEEL

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

mmgp-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+mmgp