fix(models): keep pipelines whose class already serves the requested task

get_diffusers_task reports image-to-image for any class that diffusers
registers under both the text-to-image and image-to-image tables, so every
text-to-image generation on such a model (Klein, FLUX.2, Kontext,
Qwen-Image-Edit, LongCat) went through AutoPipeline.from_pipe and came back
as a new object. The prompt cache is keyed on the pipeline identity, so the
text encoder ran again on every image. set_diffuser_pipe now returns the
pipeline unchanged when its class is registered for the requested task.
This commit is contained in:
CalamitousFelicitousness
2026-09-06 16:00:03 +01:00
parent fa7243b239
commit c389ce92d0
+13
View File
@@ -1147,6 +1147,17 @@ def get_diffusers_task(pipe: diffusers.DiffusionPipeline) -> DiffusersTaskType:
return DiffusersTaskType.TEXT_2_IMAGE
def pipe_serves_task(pipe: diffusers.DiffusionPipeline, task_type: DiffusersTaskType) -> bool:
"""True when the pipeline class is registered for the task in the diffusers auto-pipeline tables."""
mappings = {
DiffusersTaskType.TEXT_2_IMAGE: diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING,
DiffusersTaskType.IMAGE_2_IMAGE: diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING,
DiffusersTaskType.INPAINTING: diffusers.pipelines.auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING,
}
mapping = mappings.get(task_type)
return mapping is not None and pipe.__class__ in mapping.values()
def switch_pipe(cls: type[diffusers.DiffusionPipeline] | str, pipeline: diffusers.DiffusionPipeline | None = None, force = False, args: dict | None = None):
"""
args:
@@ -1354,6 +1365,8 @@ def set_diffuser_pipe(pipe, new_pipe_type):
return pipe
if get_diffusers_task(pipe) == new_pipe_type:
return pipe
if pipe_serves_task(pipe, new_pipe_type): # a class registered for several tasks classifies as one of them
return pipe
if get_diffusers_task(pipe) == DiffusersTaskType.MODULAR:
return pipe