diffengine.models.editors.amused

Submodules

Package Contents

Classes

AMUSEd

aMUSEd.

AMUSEdPreprocessor

AMUSEdPreprocessor.
class diffengine.models.editors.amused.AMUSEd(tokenizer, text_encoder, vae, transformer, model='amused/amused-512', loss=None, transformer_lora_config=None, text_encoder_lora_config=None, prior_loss_weight=1.0, data_preprocessor=None, vae_batch_size=8, *, finetune_text_encoder=False, gradient_checkpointing=False, enable_xformers=False)[source]

Bases: mmengine.model.BaseModel

aMUSEd.

Args:

tokenizer (dict): Config of tokenizer. text_encoder (dict): Config of text encoder. vae (dict): Config of vae. transformer (dict): Config of transformer. model (str): pretrained model name.

Defaults to “amused/amused-512”.

loss (dict): Config of loss. Defaults to

dict(type='L2Loss', loss_weight=1.0).

transformer_lora_config (dict, optional): The LoRA config dict for

Transformer. example. dict(type=”LoRA”, r=4). type is chosen from LoRA, LoHa, LoKr. Other config are same as the config of PEFT. https://github.com/huggingface/peft Defaults to None.

text_encoder_lora_config (dict, optional): The LoRA config dict for

Text Encoder. example. dict(type=”LoRA”, r=4). type is chosen from LoRA, LoHa, LoKr. Other config are same as the config of PEFT. https://github.com/huggingface/peft Defaults to None.

prior_loss_weight (float): The weight of prior preservation loss.

It works when training dreambooth with class images.

data_preprocessor (dict, optional): The pre-process config of

SDDataPreprocessor.

vae_batch_size (int): The batch size of vae. Defaults to 8. finetune_text_encoder (bool, optional): Whether to fine-tune text

encoder. Defaults to False.

gradient_checkpointing (bool): Whether or not to use gradient

checkpointing to save memory at the expense of slower backward pass. Defaults to False.

enable_xformers (bool): Whether or not to enable memory efficient

attention. Defaults to False.

property device: torch.device

Get device information.

Returns:

torch.device

Return type:

device.

set_lora()[source]

Set LORA for model.

Return type:

None

prepare_model()[source]

Prepare model for training.

Disable gradient for some models.

Return type:

None

set_xformers()[source]

Set xformers for model.

Return type:

None

infer(prompt, negative_prompt=None, height=None, width=None, num_inference_steps=12, output_type='pil', **kwargs)[source]

Inference function.

Args:
prompt (List[str]):

The prompt or prompts to guide the image generation.

negative_prompt (Optional[str]):

The prompt or prompts to guide the image generation. Defaults to None.

height (int, optional):

The height in pixels of the generated image. Defaults to None.

width (int, optional):

The width in pixels of the generated image. Defaults to None.

num_inference_steps (int): Number of inference steps.

Defaults to 12.

output_type (str): The output format of the generate image.

Choose between ‘pil’ and ‘latent’. Defaults to ‘pil’.

**kwargs: Other arguments.

Parameters:

  • prompt (list[str]) –

  • negative_prompt (str | None) –

  • height (int | None) –

  • width (int | None) –

  • num_inference_steps (int) –

  • output_type (str) –

Return type:

list[numpy.ndarray]

val_step(data)[source]

Val step.

Parameters:

data (Union[tuple, dict, list]) –

Return type:

list

test_step(data)[source]

Test step.

Parameters:

data (Union[tuple, dict, list]) –

Return type:

list

_forward_vae(img, num_batches)[source]

Forward vae.

Parameters:

  • img (torch.Tensor) –

  • num_batches (int) –

Return type:

torch.Tensor

forward(inputs, data_samples=None, mode='loss')[source]

Forward function.

Args:

inputs (dict): The input dict. data_samples (Optional[list], optional): The data samples.

Defaults to None.

mode (str, optional): The mode. Defaults to “loss”.

Returns:

dict: The loss dict.

Parameters:

  • inputs (dict) –

  • data_samples (Optional[list]) –

  • mode (str) –

Return type:

dict

Parameters:

  • tokenizer (dict) –

  • text_encoder (dict) –

  • vae (dict) –

  • transformer (dict) –

  • model (str) –

  • loss (dict | None) –

  • transformer_lora_config (dict | None) –

  • text_encoder_lora_config (dict | None) –

  • prior_loss_weight (float) –

  • data_preprocessor (dict | torch.nn.Module | None) –

  • vae_batch_size (int) –

  • finetune_text_encoder (bool) –

  • gradient_checkpointing (bool) –

  • enable_xformers (bool) –

class diffengine.models.editors.amused.AMUSEdPreprocessor(non_blocking=False)[source]

Bases: mmengine.model.base_model.data_preprocessor.BaseDataPreprocessor

AMUSEdPreprocessor.

Parameters:

non_blocking (Optional[bool]) –

forward(data, training=False)[source]

Preprocesses the data into the model input format.

After the data pre-processing of cast_data(), forward will stack the input tensor list to a batch tensor at the first dimension.

Args:

data (dict): Data returned by dataloader training (bool): Whether to enable training time augmentation.

Returns:

dict or list: Data in the same format as the model input.

Parameters:

  • data (dict) –

  • training (bool) –

Return type:

dict | list