Shortcuts

mmedit.models.editors.guided_diffusion

Package Contents

Classes

AblatedDiffusionModel

Guided diffusion Model.

EncoderUNetModel

The half UNet model with attention and timestep embedding.

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

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

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

_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: mmedit.utils.typing.ForwardInputs, data_samples: Optional[list] = None, mode: Optional[str] = None) List[mmedit.structures.EditDataSample][source]

_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[EditDataSample]

val_step(data: dict) mmedit.utils.typing.SampleList[source]

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 detials in Metrics and Evaluator.

Returns

Generated image or image dict.

Return type

SampleList

test_step(data: dict) mmedit.utils.typing.SampleList[source]

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

Parameters

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

Returns

Generated image or image dict.

Return type

List[EditDataSample]

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

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

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

Bases: torch.nn.Module

The half UNet model with attention and timestep embedding.

For usage, see UNet.

convert_to_fp16()[source]

Convert the torso of the model to float16.

convert_to_fp32()[source]

Convert the torso of the model to float32.

forward(x, timesteps)[source]

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.

Read the Docs v: latest
Versions
master
latest
stable
zyh-re-docs
zyh-doc-notfound-extend
zyh-api-rendering
Downloads
pdf
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.