diffengine.engine.hooks.visualization_hook

Module Contents

Classes

VisualizationHook

Basic hook that invoke visualizers after train epoch.
class diffengine.engine.hooks.visualization_hook.VisualizationHook(prompt, interval=1, height=None, width=None, *, by_epoch=True, **kwargs)[source]

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[source]
before_train(runner)[source]

Before train hook.

Parameters:

runner (mmengine.runner.Runner) –

Return type:

None

after_train_iter(runner, batch_idx, data_batch=None, outputs=None)[source]

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)[source]

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) –