Class LoRANetwork

wan/utils/lora_utils.py:158–353

class LoRANetwork(torch.nn.Module):
    TRANSFORMER_TARGET_REPLACE_MODULE = [
        "CogVideoXTransformer3DModel", "WanTransformer3DModel", \
        "Wan2_2Transformer3DModel", "FluxTransformer2DModel", "QwenImageTransformer2DModel", \
        "Wan2_2Transformer3DModel_Animate", "Wan2_2Transformer3DModel_S2V", "FantasyTalkingTransformer3DModel",
    ]
    TEXT_ENCODER_TARGET_REPLACE_MODULE = ["T5LayerSelfAttention", "T5LayerFF", "BertEncoder", "T5SelfAttention", "T5CrossAttention"]
    LORA_PREFIX_TRANSFORMER = "lora_unet"
    LORA_PREFIX_TEXT_ENCODER = "lora_te"
    def __init__(
        self,
        text_encoder: Union[List[T5EncoderModel], T5EncoderModel],
        unet,
        multiplier: float = 1.0,
        lora_dim: int = 4,
        alpha: float = 1,
        dropout: Optional[float] = None,
        module_class: Type[object] = LoRAModule,
        skip_name: str = None,
        target_name: str = None,
        varbose: Optional[bool] = False,
    ) -> None:
        super().__init__()
        self.multiplier = multiplier

        self.lora_dim = lora_dim
        self.alpha = alpha
        self.dropout = dropout

        print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
        print(f"neuron dropout: p={self.dropout}")

        # create module instances
        def create_modules(
            is_unet: bool,
            root_module: torch.nn.Module,
            target_replace_modules: List[torch.nn.Module],
        ) -> List[LoRAModule]:
            prefix = (
                self.LORA_PREFIX_TRANSFORMER
                if is_unet
                else self.LORA_PREFIX_TEXT_ENCODER
            )
            loras = []
            skipped = []
            for name, module in root_module.named_modules():
                if module.__class__.__name__ in target_replace_modules:
                    for child_name, child_module in module.named_modules():
                        is_linear = child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "LoRACompatibleLinear"
                        is_conv2d = child_module.__class__.__name__ == "Conv2d" or child_module.__class__.__name__ == "LoRACompatibleConv"
                        is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
                        
                        if skip_name is not None and skip_name in child_name:
                            continue
                        
                        if target_name is not None:
                            target_name_in = False
                            if isinstance(target_name, str):