diffengine.engine.hooks

Submodules

• diffengine.engine.hooks.compile_hook
• diffengine.engine.hooks.controlnet_save_hook
• diffengine.engine.hooks.fast_norm_hook
• diffengine.engine.hooks.ip_adapter_save_hook
• diffengine.engine.hooks.lcm_ema_update_hook
• diffengine.engine.hooks.peft_save_hook
• diffengine.engine.hooks.prior_save_hook
• diffengine.engine.hooks.sd_checkpoint_hook
• diffengine.engine.hooks.t2i_adapter_save_hook
• diffengine.engine.hooks.transformer_checkpoint_hook
• diffengine.engine.hooks.unet_ema_hook
• diffengine.engine.hooks.visualization_hook

Package Contents

Classes

CompileHook	Compile Hook.
ControlNetSaveHook	ControlNet Save Hook.
FastNormHook	Fast Normalization Hook.
IPAdapterSaveHook	IP Adapter Save Hook.
LCMEMAUpdateHook	LCM EMA Update Hook.
PeftSaveHook	Peft Save Hook.
PriorSaveHook	Prior Save Hook.
SDCheckpointHook	Delete 'vae' from checkpoint for efficient save.
T2IAdapterSaveHook	T2I Adapter Save Hook.
TransformerCheckpointHook	Delete 'vae' from checkpoint for efficient save.
UnetEMAHook	Unet EMA Hook.
VisualizationHook	Basic hook that invoke visualizers after train epoch.

class diffengine.engine.hooks.CompileHook(backend='inductor', mode=None, *, compile_main=False)

Bases: mmengine.hooks.Hook

Compile Hook.

Args:

backend (str): The backend to use for compilation.

Defaults to “inductor”.

mode (str): The mode to use for compilation. Defaults to None. compile_main (bool): Whether to compile the main network like unet or

transformer. Defaults to False.

priority = VERY_LOW

before_train(runner)

Compile the model.

Args:

runner (Runner): The runner of the training process.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

Parameters:

• backend (str) –

• mode (str | None) –

• compile_main (bool) –

class diffengine.engine.hooks.ControlNetSaveHook

Bases: mmengine.hooks.Hook

ControlNet Save Hook.

Save ControlNet weights with diffusers format and pick up ControlNet weights from checkpoint.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.FastNormHook(*, fuse_text_encoder_ln=False, fuse_main_ln=True, fuse_gn=False)

Bases: mmengine.hooks.Hook

Fast Normalization Hook.

Replace the normalization layer with a faster one.

Args:

fuse_text_encoder_ln (bool): Whether to fuse the text encoder layer

normalization. Defaults to False.

fuse_main_ln (bool): Whether to replace the layer normalization

in main module like unet or transformer. Defaults to True.

fuse_gn (bool)Whether to replace the group normalization.

Defaults to False.

priority = VERY_LOW

_replace_ln(module, name, device)

Replace the layer normalization with a fused one.

Parameters:

• module (torch.nn.Module) –

• name (str) –

• device (str) –

Return type:

None

_replace_gn(module, name, device)

Replace the layer normalization with a fused one.

Parameters:

• module (torch.nn.Module) –

• name (str) –

• device (str) –

Return type:

None

_replace_gn_forward(module, name)

Replace the group normalization forward with a faster one.

Parameters:

• module (torch.nn.Module) –

• name (str) –

Return type:

None

before_train(runner)

Replace the normalization layer with a faster one.

Args:

runner (Runner): The runner of the training process.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

Parameters:

• fuse_text_encoder_ln (bool) –

• fuse_main_ln (bool) –

• fuse_gn (bool) –

class diffengine.engine.hooks.IPAdapterSaveHook

Bases: mmengine.hooks.Hook

IP Adapter Save Hook.

Save IP-Adapter weights with diffusers format and pick up weights from checkpoint.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.LCMEMAUpdateHook

Bases: mmengine.hooks.hook.Hook

LCM EMA Update Hook.

before_run(runner)

Create an ema copy of the model.

Args:

runner (Runner): The runner of the training process.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

after_train_iter(runner, batch_idx, data_batch=None, outputs=None)

Update ema parameter.

Args:

runner (Runner): The runner of the training process. batch_idx (int): The index of the current batch in the train loop. data_batch (Sequence[dict], optional): Data from dataloader.

Defaults to None.

outputs (dict, optional): Outputs from model. Defaults to None.

Parameters:

• runner (mmengine.runner.Runner) –

• batch_idx (int) –

• data_batch (mmengine.hooks.hook.DATA_BATCH) –

• outputs (dict | None) –

Return type:

None

class diffengine.engine.hooks.PeftSaveHook

Bases: mmengine.hooks.Hook

Peft Save Hook.

Save LoRA weights with diffusers format and pick up LoRA weights from checkpoint.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.PriorSaveHook

Bases: mmengine.hooks.Hook

Prior Save Hook.

Save Prior weights with diffusers format and pick up Prior weights from checkpoint.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.SDCheckpointHook

Bases: mmengine.hooks.Hook

Delete ‘vae’ from checkpoint for efficient save.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.T2IAdapterSaveHook

Bases: mmengine.hooks.Hook

T2I Adapter Save Hook.

Save T2I-Adapter weights with diffusers format and pick up weights from checkpoint.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.TransformerCheckpointHook

Bases: mmengine.hooks.Hook

Delete ‘vae’ from checkpoint for efficient save.

priority = VERY_LOW

before_save_checkpoint(runner, checkpoint)

Before save checkpoint hook.

Args:

runner (Runner): The runner of the training, validation or testing

process.

checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.UnetEMAHook(ema_type='ExponentialMovingAverage', strict_load=False, begin_iter=0, begin_epoch=0, **kwargs)

Bases: mmengine.hooks.ema_hook.EMAHook

Unet EMA Hook.

Parameters:

• ema_type (str) –

• strict_load (bool) –

• begin_iter (int) –

• begin_epoch (int) –

before_run(runner)

Create an ema copy of the model.

Args:

runner (Runner): The runner of the training process.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

_swap_ema_state_dict(checkpoint)

Swap the state dict values of model with ema_model.

Parameters:

checkpoint (dict) –

Return type:

None

after_load_checkpoint(runner, checkpoint)

Resume ema parameters from checkpoint.

Args:

runner (Runner): The runner of the testing process. checkpoint (dict): Model’s checkpoint.

Parameters:

• runner (mmengine.runner.Runner) –

• checkpoint (dict) –

Return type:

None

class diffengine.engine.hooks.VisualizationHook(prompt, interval=1, height=None, width=None, *, by_epoch=True, **kwargs)

Bases: mmengine.hooks.Hook

Basic hook that invoke visualizers after train epoch.

Args:

prompt (List[str]):

The prompt or prompts to guide the image generation.

interval (int): Visualization interval (every k iterations).

Defaults to 1.

by_epoch (bool): Whether to visualize by epoch. Defaults to True. height (int, optional, defaults to

self.unet.config.sample_size * self.vae_scale_factor): The height in pixels of the generated image.

width (int, optional, defaults to

self.unet.config.sample_size * self.vae_scale_factor): The width in pixels of the generated image.

priority = NORMAL

before_train(runner)

Before train hook.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

after_train_iter(runner, batch_idx, data_batch=None, outputs=None)

After train iter hook.

Args:

runner (Runner): The runner of the training process. batch_idx (int): The index of the current batch. data_batch (DATA_BATCH, optional): The current data batch. outputs (dict, optional): The outputs of the current batch.

Parameters:

• runner (mmengine.runner.Runner) –

• batch_idx (int) –

• data_batch (mmengine.hooks.hook.DATA_BATCH) –

• outputs (Optional[dict]) –

Return type:

None

after_train_epoch(runner)

After train epoch hook.

Args:

runner (Runner): The runner of the training process.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

Parameters:

• prompt (list[str]) –

• interval (int) –

• height (int | None) –

• width (int | None) –

• by_epoch (bool) –