mmagic.models.utils.model_utils

Module Contents

Functions

default_init_weights(module[, scale])	Initialize network weights.
make_layer(block, num_blocks, **kwarg)	Make layers by stacking the same blocks.
get_module_device(module)	Get the device of a module.
set_requires_grad(nets[, requires_grad])	Set requires_grad for all the networks.
generation_init_weights(module[, init_type, init_gain])	Default initialization of network weights for image generation.
get_valid_noise_size(→ Optional[int])	Get the value of noise_size from input, generator and check the
get_valid_num_batches(→ int)	Try get the valid batch size from inputs.
build_module(→ Any)	Build module from config or return the module itself.
xformers_is_enable(→ bool)	Check whether xformers is installed.
set_xformers(→ torch.nn.Module)	Set xformers' efficient Attention for attention modules.
set_tomesd(model[, ratio, max_downsample, sx, sy, ...])	Patches a stable diffusion model with ToMe. Apply this to the highest
remove_tomesd(model)	Removes a patch from a ToMe Diffusion module if it was already patched.

mmagic.models.utils.model_utils.default_init_weights(module, scale=1)

Initialize network weights.

参数

• modules (nn.Module) – Modules to be initialized.

• scale (float) – Scale initialized weights, especially for residual blocks. Default: 1.

mmagic.models.utils.model_utils.make_layer(block, num_blocks, **kwarg)

Make layers by stacking the same blocks.

参数

• block (nn.module) – nn.module class for basic block.

• num_blocks (int) – number of blocks.

返回

Stacked blocks in nn.Sequential.

返回类型

nn.Sequential

mmagic.models.utils.model_utils.get_module_device(module)

Get the device of a module.

参数

module (nn.Module) – A module contains the parameters.

返回

The device of the module.

返回类型

torch.device

mmagic.models.utils.model_utils.set_requires_grad(nets, requires_grad=False)

Set requires_grad for all the networks.

参数

• nets (nn.Module | list[nn.Module]) – A list of networks or a single network.

• requires_grad (bool) – Whether the networks require gradients or not

mmagic.models.utils.model_utils.generation_init_weights(module, init_type='normal', init_gain=0.02)

Default initialization of network weights for image generation.

By default, we use normal init, but xavier and kaiming might work better for some applications.

参数

• module (nn.Module) – Module to be initialized.

• init_type (str) – The name of an initialization method: normal | xavier | kaiming | orthogonal. Default: ‘normal’.

• init_gain (float) – Scaling factor for normal, xavier and orthogonal. Default: 0.02.

mmagic.models.utils.model_utils.get_valid_noise_size(noise_size: Optional[int], generator: Union[Dict, torch.nn.Module]) → Optional[int]

Get the value of noise_size from input, generator and check the consistency of these values. If no conflict is found, return that value.

参数

• noise_size (Optional[int]) – noise_size passed to BaseGAN_refactor’s initialize function.

• generator (ModelType) – The config or the model of generator.

返回

The noise size feed to generator.

返回类型

int | None

mmagic.models.utils.model_utils.get_valid_num_batches(batch_inputs: Optional[mmagic.utils.typing.ForwardInputs] = None, data_samples: List[mmagic.structures.DataSample] = None) → int

Try get the valid batch size from inputs.

• If some values in batch_inputs are Tensor and ‘num_batches’ is in batch_inputs, we check whether the value of ‘num_batches’ and the the length of first dimension of all tensors are same. If the values are not same, AssertionError will be raised. If all values are the same, return the value.

• If no values in batch_inputs is Tensor, ‘num_batches’ must be contained in batch_inputs. And this value will be returned.

• If some values are Tensor and ‘num_batches’ is not contained in batch_inputs, we check whether all tensor have the same length on the first dimension. If the length are not same, AssertionError will be raised. If all length are the same, return the length as batch size.

• If batch_inputs is a Tensor, directly return the length of the first dimension as batch size.

参数

batch_inputs (ForwardInputs) – Inputs passed to forward().

返回

The batch size of samples to generate.

返回类型

int

mmagic.models.utils.model_utils.build_module(module: Union[dict, torch.nn.Module], builder: mmengine.registry.Registry, *args, **kwargs) → Any

Build module from config or return the module itself.

参数

• module (Union[dict, nn.Module]) – The module to build.

• builder (Registry) – The registry to build module.

• *args – Arguments passed to build function.

• **kwargs –

  Arguments passed to build function.

返回

The built module.

返回类型

Any

mmagic.models.utils.model_utils.xformers_is_enable(verbose: bool = False) → bool

Check whether xformers is installed. :param verbose: Whether to print the log. :type verbose: bool

返回

Whether xformers is installed.

返回类型

bool

mmagic.models.utils.model_utils.set_xformers(module: torch.nn.Module, prefix: str = '') → torch.nn.Module

Set xformers’ efficient Attention for attention modules.

参数

• module (nn.Module) – The module to set xformers.

• prefix (str) – The prefix of the module name.

返回

The module with xformers’ efficient Attention.

返回类型

nn.Module

mmagic.models.utils.model_utils.set_tomesd(model: torch.nn.Module, ratio: float = 0.5, max_downsample: int = 1, sx: int = 2, sy: int = 2, use_rand: bool = True, merge_attn: bool = True, merge_crossattn: bool = False, merge_mlp: bool = False)

Patches a stable diffusion model with ToMe. Apply this to the highest level stable diffusion object.

Refer to: https://github.com/dbolya/tomesd/blob/main/tomesd/patch.py#L173 # noqa

参数

• model (torch.nn.Module) – A top level Stable Diffusion module to patch in place.

• ratio (float) – The ratio of tokens to merge. I.e., 0.4 would reduce the total number of tokens by 40%.The maximum value for this is 1-(1/(sx * sy)). By default, the max ratio is 0.75 (usually <= 0.5 is recommended). Higher values result in more speed-up, but with more visual quality loss.

• max_downsample (int) – Apply ToMe to layers with at most this amount of downsampling. E.g., 1 only applies to layers with no downsampling, while 8 applies to all layers. Should be chosen from [1, 2, 4, or 8]. 1 and 2 are recommended.

• sx (int, int) – The stride for computing dst sets. A higher stride means you can merge more tokens, default setting of (2, 2) works well in most cases. sx and sy do not need to divide image size.

• sy (int, int) – The stride for computing dst sets. A higher stride means you can merge more tokens, default setting of (2, 2) works well in most cases. sx and sy do not need to divide image size.

• use_rand (bool) – Whether or not to allow random perturbations when computing dst sets. By default: True, but if you’re having weird artifacts you can try turning this off.

• merge_attn (bool) – Whether or not to merge tokens for attention (recommended).

• merge_crossattn (bool) – Whether or not to merge tokens for cross attention (not recommended).

• merge_mlp (bool) – Whether or not to merge tokens for the mlp layers (particular not recommended).

返回

Model patched by ToMe.

返回类型

model (torch.nn.Module)

mmagic.models.utils.model_utils.remove_tomesd(model: torch.nn.Module)

Removes a patch from a ToMe Diffusion module if it was already patched.

Refer to: https://github.com/dbolya/tomesd/blob/main/tomesd/patch.py#L251 # noqa