mmagic.models.archs.attention_injection

Module Contents

Classes

AttentionInjection

Wrapper for stable diffusion unet.

Functions

torch_dfs(model)

Attributes

AttentionStatus

mmagic.models.archs.attention_injection.AttentionStatus

mmagic.models.archs.attention_injection.torch_dfs(model: torch.nn.Module)

class mmagic.models.archs.attention_injection.AttentionInjection(module: torch.nn.Module, injection_weight=5)

Bases: torch.nn.Module

Wrapper for stable diffusion unet.

参数

module (nn.Module) – The module to be wrapped.

forward(x: torch.Tensor, t, encoder_hidden_states=None, down_block_additional_residuals=None, mid_block_additional_residual=None, ref_x=None) → torch.Tensor

Forward and add LoRA mapping.

参数

x (Tensor) – The input tensor.

返回

The output tensor.

返回类型

Tensor