diff --git a/auto_round/utils/model.py b/auto_round/utils/model.py index d84b18b64..f01bb6949 100644 --- a/auto_round/utils/model.py +++ b/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 diff --git a/test/test_cpu/models/test_diffusion.py b/test/test_cpu/models/test_diffusion.py index 3d047cb72..e32a1c9d9 100644 --- a/test/test_cpu/models/test_diffusion.py +++ b/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))