Update get_block_names func (#1047)

Author: mengniwang95

auto_round/utils/model.py

@@ -631,18 +631,13 @@ def _get_vlm_block_names(model, quant_vision=False):
         block_names = []
         target_modules = []
         vision_blocks_tuple = ("vision", "visual", "image", "img")
-        last_block_name = ""
-        for n, m in model.named_modules():
-            if hasattr(type(m), "__name__") and "ModuleList" in type(m).__name__:
-                if quant_vision or all(key not in n.lower() for key in (vision_blocks_tuple)):
-                    if last_block_name and last_block_name in n:
-                        continue
-                    target_modules.append((n, m))
-                    last_block_name = n
+        target_modules = _search_block("", model)
+
         for i, target_m in enumerate(target_modules):
-            block_names.append([])
-            for n, m in target_m[1].named_children():
-                block_names[i].append(target_m[0] + "." + n)
+            if quant_vision or all(key not in target_m[0].lower() for key in (vision_blocks_tuple)):
+                block_names.append([])
+                for n, m in target_m[1].named_children():
+                    block_names[-1].append(target_m[0] + "." + n)
         return block_names
 
     if quant_vision or not is_pure_text_model(model):

test/test_cuda/test_get_block_name.py

@@ -30,6 +30,7 @@ def tearDownClass(self):
         shutil.rmtree("runs", ignore_errors=True)
 
     def check_block_names(self, block_names, prefixs=[], n_layers=[]):
+        assert len(block_names) == len(prefixs) == len(n_layers)
         for i, block_name in enumerate(block_names):
             prefix = prefixs[i]
             n_layer = n_layers[i]
@@ -196,6 +197,9 @@ def test_flux(self):
         self.check_block_names(block_names, ["transformer_blocks", "single_transformer_blocks"], [19, 38])
         self.assertTrue(any(["context_embedder" not in n for n in block_names]))
 
+        block_names = get_block_names(model, quant_vision=True)
+        self.check_block_names(block_names, ["transformer_blocks", "single_transformer_blocks"], [19, 38])
+
 
 if __name__ == "__main__":
     unittest.main()