mmagic.models.editors.guided_diffusion

Package Contents

Classes

AblatedDiffusionModel	Guided diffusion Model.
EncoderUNetModel	The half UNet model with attention and timestep embedding.

class mmagic.models.editors.guided_diffusion.AblatedDiffusionModel(data_preprocessor, unet, diffusion_scheduler, use_fp16=False, classifier=None, classifier_scale=1.0, rgb2bgr=False, pretrained_cfgs=None)

Bases: mmengine.model.BaseModel

Guided diffusion Model.

Parameters

• data_preprocessor (dict, optional) – The pre-process config of BaseDataPreprocessor.

• unet (ModelType) – Config of denoising Unet.

• diffusion_scheduler (ModelType) – Config of diffusion_scheduler scheduler.

• use_fp16 (bool) – Whether to use fp16 for unet model. Defaults to False.

• classifier (ModelType) – Config of classifier. Defaults to None.

• pretrained_cfgs (dict) – Path Config for pretrained weights. Usually this is a dict contains module name and the corresponding ckpt path.Defaults to None.

property device

Get current device of the model.

Returns

The current device of the model.

Return type

torch.device

load_pretrained_models(pretrained_cfgs)

_summary_

Parameters

pretrained_cfgs (_type_) – _description_

infer(scheduler_kwargs=None, init_image=None, batch_size=1, num_inference_steps=1000, labels=None, classifier_scale=0.0, show_progress=False)

_summary_

Parameters

• init_image (_type_, optional) – _description_. Defaults to None.

• batch_size (int, optional) – _description_. Defaults to 1.

• num_inference_steps (int, optional) – _description_. Defaults to 1000.

• labels (_type_, optional) – _description_. Defaults to None.

• show_progress (bool, optional) – _description_. Defaults to False.

Returns

_description_

Return type

_type_

forward(inputs: mmagic.utils.typing.ForwardInputs, data_samples: Optional[list] = None, mode: Optional[str] = None) → List[mmagic.structures.DataSample]

_summary_

Parameters

• inputs (ForwardInputs) – _description_

• data_samples (Optional[list], optional) – _description_. Defaults to None.

• mode (Optional[str], optional) – _description_. Defaults to None.

Returns

_description_

Return type

List[DataSample]

val_step(data: dict) → mmagic.utils.typing.SampleList

Gets the generated image of given data.

Calls self.data_preprocessor(data) and self(inputs, data_sample, mode=None) in order. Return the generated results which will be passed to evaluator.

Parameters

data (dict) – Data sampled from metric specific sampler. More details in Metrics and Evaluator.

Returns

Generated image or image dict.

Return type

SampleList

test_step(data: dict) → mmagic.utils.typing.SampleList

Gets the generated image of given data. Same as val_step().

Parameters

data (dict) – Data sampled from metric specific sampler. More details in Metrics and Evaluator.

Returns

Generated image or image dict.

Return type

List[DataSample]

train_step(data: dict, optim_wrapper: mmengine.optim.OptimWrapperDict)

_summary_

Parameters

• data (dict) – _description_

• optim_wrapper (OptimWrapperDict) – _description_

Returns

_description_

Return type

_type_

get_module(model: torch.nn.Module, module_name: str) → torch.nn.Module

Get an inner module from model.

Since we will wrapper DDP for some model, we have to judge whether the module can be indexed directly.

Parameters

• model (nn.Module) – This model may wrapped with DDP or not.

• module_name (str) – The name of specific module.

Returns

Returned sub module.

Return type

nn.Module

class mmagic.models.editors.guided_diffusion.EncoderUNetModel(image_size, in_channels, model_channels, out_channels, num_res_blocks, attention_resolutions, dropout=0, channel_mult=(1, 2, 4, 8), conv_resample=True, dims=2, use_checkpoint=False, use_fp16=False, num_heads=1, num_head_channels=- 1, num_heads_upsample=- 1, use_scale_shift_norm=False, resblock_updown=False, use_new_attention_order=False, pool='adaptive')

Bases: torch.nn.Module

The half UNet model with attention and timestep embedding.

For usage, see UNet.

convert_to_fp16()

Convert the torso of the model to float16.

convert_to_fp32()

Convert the torso of the model to float32.

forward(x, timesteps)

Apply the model to an input batch.

Parameters

• x – an [N x C x …] Tensor of inputs.

• timesteps – a 1-D batch of timesteps.

Returns

an [N x K] Tensor of outputs.