Support GLM-Image model quantizaiton

auto_round/autoround.py

@@ -25,6 +25,7 @@
     LLMCompressor,
     MLLMCompressor,
 )
+from auto_round.compressors.diffusion.hybrid import HybridCompressor, is_hybrid_diffusion_model
 from auto_round.logger import deprecated, logger
 from auto_round.schemes import QuantizationScheme
 from auto_round.utils import is_diffusion_model, is_mllm_model
@@ -162,7 +163,19 @@ def __new__(
 
         model_cls = []
 
-        if (extra_config and not extra_config.mllm_config.is_default()) or is_mllm_model(model, platform=platform):
+        has_multimodal_assets = kwargs.get("processor") is not None or kwargs.get("image_processor") is not None
+
+        if is_hybrid_diffusion_model(model):
+            logger.info("using Hybrid AR+Diffusion mode for hybrid model.")
+            model_cls.append(HybridCompressor)
+            if extra_config:
+                extra_config.mllm_config = None
+                extra_config.diffusion_config = None
+        elif (
+            (extra_config and not extra_config.mllm_config.is_default())
+            or has_multimodal_assets
+            or is_mllm_model(model, platform=platform)
+        ):
             logger.info("using MLLM mode for multimodal model.")
             model_cls.append(MLLMCompressor)
             if extra_config:

auto_round/compressors/__init__.py

@@ -17,6 +17,7 @@
 from auto_round.compressors.base import LLMCompressor
 from auto_round.compressors.mllm.compressor import MLLMCompressor
 from auto_round.compressors.diffusion.compressor import DiffusionCompressor
+from auto_round.compressors.diffusion.hybrid import HybridCompressor
 from auto_round.compressors.config import (
     DiffusionExtraConfig,
     ExtraConfig,

auto_round/compressors/diffusion/compressor.py

@@ -33,6 +33,8 @@
     extract_block_names_to_str,
     find_matching_blocks,
     get_block_names,
+    merge_block_output_keys,
+    wrap_block_forward_positional_to_kwargs,
 )
 
 pipeline_utils = LazyImport("diffusers.pipelines.pipeline_utils")
@@ -168,6 +170,9 @@ def _update_inputs(self, inputs: dict, q_inputs: dict) -> tuple[dict, dict]:
             q_inputs = {k: q_inputs.pop(k, None) for k in input_id_str}
         return inputs, q_inputs
 
+    def _get_block_forward_func(self, name):
+        return wrap_block_forward_positional_to_kwargs(super()._get_block_forward_func(name))
+
     def _split_inputs(self, inputs: dict, first_input_name: str) -> tuple[dict, dict]:
         input_id_str = [key for key in inputs.keys() if "hidden_state" in key]
         input_ids = {k: inputs.pop(k, None) for k in input_id_str}
@@ -201,7 +206,7 @@ def _get_current_q_output(
         )
         if isinstance(current_input_ids, dict):
             hidden_states = current_input_ids.pop("hidden_states")
-            current_input_others.update(current_input_ids)
+            merge_block_output_keys(block, current_input_others, current_input_ids)
             current_input_ids = hidden_states
         output_q = block_forward(block, current_input_ids, current_input_others, self.amp, self.amp_dtype, device, idx)
         return output_q.to(cache_device)
@@ -247,7 +252,7 @@ def _get_block_outputs(
             )
             if isinstance(tmp_input_ids, dict):
                 hidden_states = tmp_input_ids.pop("hidden_states")
-                tmp_input_others.update(tmp_input_ids)
+                merge_block_output_keys(block, tmp_input_others, tmp_input_ids)
                 tmp_input_ids = hidden_states
 
             tmp_output = block_forward(block, tmp_input_ids, tmp_input_others, self.amp, self.amp_dtype, device, None)