Function lora_forward

omini/pipeline/flux_omini.py:144–184

Apply a single, explicitly-selected LoRA adapter to ``module`` at scale 1.0. This is a fast, allocation-free replacement for the previous ``specify_lora`` context manager, which mutated ``module.scaling`` on every call. Semantics are preserved exactly: * If ``module`` is not a

(module, x: torch.Tensor, adapter)

def lora_forward(module, x: torch.Tensor, adapter) -> torch.Tensor:
    """
    Apply a single, explicitly-selected LoRA adapter to ``module`` at scale 1.0.

    This is a fast, allocation-free replacement for the previous
    ``specify_lora`` context manager, which mutated ``module.scaling`` on every
    call. Semantics are preserved exactly:

    * If ``module`` is not a PEFT ``BaseTunerLayer`` it is called directly.
    * The base (non-LoRA) projection is always applied.
    * When ``adapter`` is not None and is an *active* adapter with weights on
      this module, its LoRA delta is added with scale hardcoded to ``1.0``
      (matching ``specify_lora`` setting ``scaling[adapter] = 1``). All other
      adapters contribute nothing (matching ``scaling[other] = 0``).

    LoRA dropout is applied via the adapter's own dropout module, which is an
    Identity (or ``p == 0``) at inference/eval, so inference stays bit-identical
    while training matches PEFT's semantics.
    """
    if not isinstance(module, BaseTunerLayer):
        return module(x)

    result = module.base_layer(x)

    # No adapter requested, adapters disabled, or weights already merged into
    # the base layer -> nothing more to add (mirrors PEFT's forward).
    if adapter is None or module.disable_adapters or module.merged:
        return result

    # Only an *active* adapter with LoRA weights on this module contributes,
    # exactly as in the original specify_lora + PEFT forward path.
    if adapter not in module.active_adapters or adapter not in module.lora_A:
        return result

    torch_result_dtype = result.dtype
    lora_A = module.lora_A[adapter]
    lora_B = module.lora_B[adapter]
    dropout = module.lora_dropout[adapter]  # Identity at eval / when p == 0
    x = x.to(lora_A.weight.dtype)
    result = result + lora_B(lora_A(dropout(x)))  # scale hardcoded to 1.0
    return result.to(torch_result_dtype)


def _adanorm_zero_forward(norm, x, emb, adapter):