Function remove_context_parallel

src/diffusers/hooks/context_parallel.py:112–126

(module: torch.nn.Module, plan: dict[str, ContextParallelModelPlan])

def remove_context_parallel(module: torch.nn.Module, plan: dict[str, ContextParallelModelPlan]) -> None:
    for module_id, cp_model_plan in plan.items():
        submodule = _get_submodule_by_name(module, module_id)
        if not isinstance(submodule, list):
            submodule = [submodule]

        for m in submodule:
            registry = HookRegistry.check_if_exists_or_initialize(m)
            if isinstance(cp_model_plan, dict):
                hook_name = _CONTEXT_PARALLEL_INPUT_HOOK_TEMPLATE.format(module_id)
            elif isinstance(cp_model_plan, (ContextParallelOutput, list, tuple)):
                hook_name = _CONTEXT_PARALLEL_OUTPUT_HOOK_TEMPLATE.format(module_id)
            else:
                raise ValueError(f"Unsupported context parallel model plan type: {type(cp_model_plan)}")
            registry.remove_hook(hook_name)


class ContextParallelSplitHook(ModelHook):