mmgp 3.4.2__py3-none-any.whl → 3.4.3__py3-none-any.whl

mmgp/offload.py

@@ -1,4 +1,4 @@
-# ------------------ Memory Management 3.4.2 for the GPU Poor by DeepBeepMeep (mmgp)------------------
+# ------------------ Memory Management 3.4.3 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
@@ -619,7 +619,7 @@ def _welcome():
     if welcome_displayed:
          return 
     welcome_displayed = True
-    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.4.1) by DeepBeepMeep ************{ENDC}{UNBOLD}")
+    print(f"{BOLD}{HEADER}************ Memory Management for the GPU Poor (mmgp 3.4.3) by DeepBeepMeep ************{ENDC}{UNBOLD}")
 
 def change_dtype(model, new_dtype, exclude_buffers = False):
     for submodule_name, submodule in model.named_modules():  
@@ -1749,7 +1749,7 @@ class offload:
     @torch.compiler.disable()
     def gpu_unload_blocks(self, model_id, blocks_name):
         # cl = clock.start()
-        if blocks_name != None:
+        if blocks_name != None and blocks_name == self.loaded_blocks[model_id]:
             self.loaded_blocks[model_id] = None 
 
 
@@ -1805,7 +1805,13 @@ class offload:
 
             loaded_block = self.loaded_blocks[model_id]
             if loaded_block != None:
-                self.gpu_unload_blocks(model_id, loaded_block)      
+                self.gpu_unload_blocks(model_id, loaded_block)
+                entry_name = model_id + "/" + loaded_block
+                next_blocks_entry = self.next_blocks_names[entry_name] if entry_name in self.next_blocks_names else None
+                if next_blocks_entry != None:
+                    pos = next_blocks_entry.rfind("/")
+                    torch.cuda.synchronize()
+                    self.gpu_unload_blocks(model_id, next_blocks_entry[pos+1:])      
                 self.loaded_blocks[model_id] = None  
  
         self.active_models = []
@@ -2297,6 +2303,8 @@ def all(pipe_or_dict_of_modules, pinnedMemory = False, pinnedPEFTLora = False, p
                         if model_dtype== None:
                             model_dtype = dtype
                         else:
+                            if model_dtype != dtype:
+                                pass
                             assert model_dtype == dtype
                     current_model_size +=  torch.numel(p.data) * p.data.element_size()
                 current_model._dtype = model_dtype

{mmgp-3.4.2.dist-info → mmgp-3.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mmgp
-Version: 3.4.2
+Version: 3.4.3
 Summary: Memory Management for the GPU Poor
 Author-email: deepbeepmeep <deepbeepmeep@yahoo.com>
 License:                     GNU GENERAL PUBLIC LICENSE
@@ -17,7 +17,7 @@ Dynamic: license-file
 
 
 <p align="center">
-  <H2>Memory Management 3.4.1 for the GPU Poor by DeepBeepMeep</H2>	
+  <H2>Memory Management 3.4.3 for the GPU Poor by DeepBeepMeep</H2>	
 </p>
 
 

mmgp-3.4.3.dist-info/RECORD

@@ -0,0 +1,9 @@
+__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mmgp/__init__.py,sha256=A9qBwyQMd1M7vshSTOBnFGP1MQvS2hXmTcTCMUcmyzE,509
+mmgp/offload.py,sha256=fR6ACUxT4rVIbfxM8p-bkKCbNYwZi6061yyfjzCEnlM,114769
+mmgp/safetensors2.py,sha256=GFbyNrQpvtV1RTen6_iUl47LC3jtM0PMtB9dsERTElY,18355
+mmgp-3.4.3.dist-info/licenses/LICENSE.md,sha256=HjzvY2grdtdduZclbZ46B2M-XpT4MDCxFub5ZwTWq2g,93
+mmgp-3.4.3.dist-info/METADATA,sha256=PXvefWGCj1dt5XmXvP5NXoq6W7wuDaZaCkuyp_i_2tY,16350
+mmgp-3.4.3.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+mmgp-3.4.3.dist-info/top_level.txt,sha256=waGaepj2qVfnS2yAOkaMu4r9mJaVjGbEi6AwOUogU_U,14
+mmgp-3.4.3.dist-info/RECORD,,

{mmgp-3.4.2.dist-info → mmgp-3.4.3.dist-info}/WHEEL

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

mmgp-3.4.2.dist-info/RECORD

@@ -1,9 +0,0 @@
-__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mmgp/__init__.py,sha256=A9qBwyQMd1M7vshSTOBnFGP1MQvS2hXmTcTCMUcmyzE,509
-mmgp/offload.py,sha256=WG2yJEBTZ8BBdutJN5_eUUby_eGEUky68pwNH1p4-PU,114217
-mmgp/safetensors2.py,sha256=GFbyNrQpvtV1RTen6_iUl47LC3jtM0PMtB9dsERTElY,18355
-mmgp-3.4.2.dist-info/licenses/LICENSE.md,sha256=HjzvY2grdtdduZclbZ46B2M-XpT4MDCxFub5ZwTWq2g,93
-mmgp-3.4.2.dist-info/METADATA,sha256=THPFQEc2FzUzAcm8WpW_CuO0IPJThvisNMeq-oJO2CQ,16350
-mmgp-3.4.2.dist-info/WHEEL,sha256=GHB6lJx2juba1wDgXDNlMTyM13ckjBMKf-OnwgKOCtA,91
-mmgp-3.4.2.dist-info/top_level.txt,sha256=waGaepj2qVfnS2yAOkaMu4r9mJaVjGbEi6AwOUogU_U,14
-mmgp-3.4.2.dist-info/RECORD,,