fix: recognize Modular Diffusers pipelines as diffusion models

auto_round/utils/model.py

@@ -867,8 +867,17 @@ def diffusion_load_model(
         return pipe, pipe.model
 
     pipelines = LazyImport("diffusers.pipelines")
+    modular_pipeline = LazyImport("diffusers.modular_pipelines.modular_pipeline")
     if isinstance(pretrained_model_name_or_path, str):
         model_index = os.path.join(pretrained_model_name_or_path, "model_index.json")
+        if not os.path.exists(model_index) and os.path.exists(
+            os.path.join(pretrained_model_name_or_path, MODULAR_PIPELINE_INDEX_NAME)
+        ):
+            raise NotImplementedError(
+                f"{pretrained_model_name_or_path} is a Modular Diffusers pipeline "
+                f"({MODULAR_PIPELINE_INDEX_NAME}), which auto_round cannot assemble from a path yet. "
+                "Build the ModularPipeline yourself and pass the pipeline object as `model` instead."
+            )
         with open(model_index, "r", encoding="utf-8") as file:
             config = json.load(file)
 
@@ -886,13 +895,20 @@ def diffusion_load_model(
         )
         pipe_config = pipe.load_config(pretrained_model_name_or_path)
 
-    elif isinstance(pretrained_model_name_or_path, pipelines.pipeline_utils.DiffusionPipeline):
+    elif isinstance(
+        pretrained_model_name_or_path,
+        (pipelines.pipeline_utils.DiffusionPipeline, modular_pipeline.ModularPipeline),
+    ):
         pipe = pretrained_model_name_or_path
-        pipe_config = pipe.load_config(pipe.config["_name_or_path"])
+        # a pipeline assembled in-process, as Modular Diffusers ones typically are,
+        # has no _name_or_path to reload the on-disk index from
+        name_or_path = pipe.config.get("_name_or_path", None)
+        pipe_config = pipe.load_config(name_or_path) if name_or_path is not None else {}
 
     else:
         raise ValueError(
-            f"Only support str or DiffusionPipeline class for model, but get {type(pretrained_model_name_or_path)}"
+            f"Only support str, DiffusionPipeline or ModularPipeline class for model, "
+            f"but get {type(pretrained_model_name_or_path)}"
         )
 
     # add missing key
@@ -1104,7 +1120,10 @@ def is_mllm_model(model_or_path: Union[str, torch.nn.Module], platform: str = No
     # For dummy model, model_path could be "".
     # Only try to download if the path looks like a HF repo id (not a local filesystem path).
     # Skip download for absolute paths or relative paths that contain current/parent dir markers.
-    _is_local_path = os.path.isabs(model_path) or model_path.startswith("./") or model_path.startswith("../")
+    # model_path is None for a model or pipeline built in-process, which has no name or path
+    _is_local_path = isinstance(model_path, str) and (
+        os.path.isabs(model_path) or model_path.startswith("./") or model_path.startswith("../")
+    )
     if model_path and not os.path.isdir(model_path) and not _is_local_path:
         model_path = download_or_get_path(model_path, platform=platform)
 
@@ -1147,6 +1166,38 @@ def is_gguf_model(model_path: Union[str, torch.nn.Module]) -> bool:
     return is_gguf_file
 
 
+## ModularPipeline.config_name, spelled out to keep the diffusers import lazy
+MODULAR_PIPELINE_INDEX_NAME = "modular_model_index.json"
+
+
+def _find_pipeline_index_file(model_dir_or_repo: str) -> Optional[str]:
+    """Return the pipeline index file of a diffusers directory or repo, if it has one.
+
+    Standard pipelines ship ``model_index.json``, Modular Diffusers pipelines ship
+    ``modular_model_index.json`` instead.
+    """
+    index_names = ("model_index.json", MODULAR_PIPELINE_INDEX_NAME)
+
+    if os.path.isdir(model_dir_or_repo):
+        for name in index_names:
+            index_file = os.path.join(model_dir_or_repo, name)
+            if os.path.exists(index_file):
+                check_diffusers_installed()
+                return index_file
+        return None
+
+    from huggingface_hub import hf_hub_download
+
+    for name in index_names:
+        try:
+            index_file = hf_hub_download(model_dir_or_repo, name)
+            check_diffusers_installed()
+            return index_file
+        except Exception as e:
+            logger.debug(f"No {name} found in {model_dir_or_repo}: {e}")
+    return None
+
+
 def is_diffusion_model(model_or_path: Union[str, object], trust_remote_code: bool = True) -> bool:
     from auto_round.utils.common import LazyImport
 
@@ -1171,25 +1222,12 @@ def is_diffusion_model(model_or_path: Union[str, object], trust_remote_code: boo
             logger.warning(
                 f"Failed to load config for {model_or_path}, trying to check model_index.json for diffusion pipeline."
             )
-        index_file = None
-        if not os.path.isdir(model_or_path):
-            try:
-                from huggingface_hub import hf_hub_download
-
-                index_file = hf_hub_download(model_or_path, "model_index.json")
-                check_diffusers_installed()
-            except Exception as e:
-                print(e)
-                index_file = None
-
-        elif os.path.exists(os.path.join(model_or_path, "model_index.json")):
-            check_diffusers_installed()
-            index_file = os.path.join(model_or_path, "model_index.json")
-        return index_file is not None
+        return _find_pipeline_index_file(model_or_path) is not None
     elif not isinstance(model_or_path, torch.nn.Module):
         check_diffusers_installed()
         pipeline_utils = LazyImport("diffusers.pipelines.pipeline_utils")
-        return isinstance(model_or_path, pipeline_utils.DiffusionPipeline)
+        modular_pipeline = LazyImport("diffusers.modular_pipelines.modular_pipeline")
+        return isinstance(model_or_path, (pipeline_utils.DiffusionPipeline, modular_pipeline.ModularPipeline))
     else:
         return False
 

test/test_cpu/models/test_diffusion.py

@@ -70,3 +70,42 @@ def test_flux(setup_flux):
 #     all_inputs = autoround.cache_inter_data(["transformer_blocks.0"], 2)
 #     assert len(all_inputs["transformer_blocks.0"]["hidden_states"]) == 4
 #     shutil.rmtree(output_dir, ignore_errors=True)
+
+
+def _build_empty_modular_pipeline():
+    """A ModularPipeline with no components, so nothing is downloaded or loaded."""
+    from diffusers.modular_pipelines import SequentialPipelineBlocks
+
+    class EmptyBlocks(SequentialPipelineBlocks):
+        block_classes = []
+        block_names = []
+
+    return EmptyBlocks().init_pipeline()
+
+
+def test_modular_pipeline_is_detected_as_diffusion():
+    """A ModularPipeline is not a DiffusionPipeline, but it is still a diffusion model."""
+    pytest.importorskip("diffusers.modular_pipelines")
+
+    from auto_round.utils.model import detect_model_type, is_diffusion_model, is_mllm_model
+
+    pipe = _build_empty_modular_pipeline()
+
+    # is_mllm_model is consulted first and has no path to inspect for an in-process pipeline
+    assert is_mllm_model(pipe) is False
+    assert is_diffusion_model(pipe) is True
+    assert detect_model_type(pipe) == "diffusion"
+
+
+def test_modular_model_index_dir_is_detected_as_diffusion(tmp_path):
+    """Modular Diffusers ships modular_model_index.json instead of model_index.json."""
+    pytest.importorskip("diffusers.modular_pipelines")
+
+    from auto_round.utils.model import diffusion_load_model, is_diffusion_model
+
+    (tmp_path / "modular_model_index.json").write_text("{}", encoding="utf-8")
+
+    assert is_diffusion_model(str(tmp_path)) is True
+    # assembling a modular pipeline from a path is not supported, but it must say so
+    with pytest.raises(NotImplementedError, match="Modular Diffusers"):
+        diffusion_load_model(str(tmp_path))