diff --git a/src/diffusers/pipelines/pipeline_loading_utils.py b/src/diffusers/pipelines/pipeline_loading_utils.py index 69bce1a1c533..e13d6dbea4a6 100644 --- a/src/diffusers/pipelines/pipeline_loading_utils.py +++ b/src/diffusers/pipelines/pipeline_loading_utils.py @@ -936,7 +936,11 @@ def _fetch_class_library_tuple(module): pipeline_dir = module_path_items[-2] if len(module_path_items) > 2 else None path = not_compiled_module.__module__.split(".") - is_pipeline_module = pipeline_dir in path and hasattr(pipelines, pipeline_dir) + # A same-named folder in another library (e.g. `transformers.models.diffusion_gemma` vs + # `diffusers.pipelines.diffusion_gemma`) must not count as a pipeline module. + is_pipeline_module = ( + path[0] == diffusers_module.__name__ and pipeline_dir in path and hasattr(pipelines, pipeline_dir) + ) # if library is not in LOADABLE_CLASSES, then it is a custom module. # Or if it's a pipeline module, then the module is inside the pipeline diff --git a/tests/pipelines/test_pipeline_utils.py b/tests/pipelines/test_pipeline_utils.py index 6bf79ae5dde3..521dc8802d73 100644 --- a/tests/pipelines/test_pipeline_utils.py +++ b/tests/pipelines/test_pipeline_utils.py @@ -1111,3 +1111,29 @@ def test_push_to_hub_library_name(self): # Reset repo delete_repo(repo_id, token=TOKEN) + + +class TestFetchClassLibraryTuple: + def test_diffusers_model(self): + from diffusers import UNet2DConditionModel + from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple + + assert _fetch_class_library_tuple(UNet2DConditionModel) == ("diffusers", "UNet2DConditionModel") + + def test_pipeline_module_class(self): + from diffusers.pipelines.deepfloyd_if import IFWatermarker + from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple + + assert _fetch_class_library_tuple(IFWatermarker) == ("deepfloyd_if", "IFWatermarker") + + def test_other_library_class_shadowing_pipeline_dir(self): + from diffusers.pipelines.pipeline_loading_utils import _fetch_class_library_tuple + + # A transformers class whose model folder shares its name with a diffusers pipeline folder + # (e.g. `transformers.models.diffusion_gemma` vs `diffusers.pipelines.diffusion_gemma`) must + # resolve to its own library, not to the pipeline folder. + class FakeModel: + pass + + FakeModel.__module__ = "transformers.models.diffusion_gemma.modeling_diffusion_gemma" + assert _fetch_class_library_tuple(FakeModel) == ("transformers", "FakeModel")