sketchkit.colorization.manga_ninja.manganinja.models package¶

Submodules¶

sketchkit.colorization.manga_ninja.manganinja.models.attention module¶

class sketchkit.colorization.manga_ninja.manganinja.models.attention.BasicTransformerBlock(dim: int, num_attention_heads: int, attention_head_dim: int, dropout=0.0, cross_attention_dim: int | None = None, activation_fn: str = 'geglu', num_embeds_ada_norm: int | None = None, attention_bias: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, upcast_attention: bool = False, norm_elementwise_affine: bool = True, norm_type: str = 'layer_norm', norm_eps: float = 1e-05, final_dropout: bool = False, attention_type: str = 'default', positional_embeddings: str | None = None, num_positional_embeddings: int | None = None)[source]¶

Bases: Module

A basic Transformer block.

Parameters:

dim (int) – The number of channels in the input and output.
num_attention_heads (int) – The number of heads to use for multi-head attention.
attention_head_dim (int) – The number of channels in each head.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.
cross_attention_dim (int, optional) – The size of the encoder_hidden_states vector for cross attention.
activation_fn (str, optional, defaults to “geglu”) – Activation function to be used in feed-forward.
( (attention_bias) – obj: int, optional): The number of diffusion steps used during training. See Transformer2DModel.
( – obj: bool, optional, defaults to False): Configure if the attentions should contain a bias parameter.
only_cross_attention (bool, optional) – Whether to use only cross-attention layers. In this case two cross attention layers are used.
double_self_attention (bool, optional) – Whether to use two self-attention layers. In this case no cross attention layers are used.
upcast_attention (bool, optional) – Whether to upcast the attention computation to float32. This is useful for mixed precision training.
norm_elementwise_affine (bool, optional, defaults to True) – Whether to use learnable elementwise affine parameters for normalization.
norm_type (str, optional, defaults to “layer_norm”) – The normalization layer to use. Can be “layer_norm”, “ada_norm” or “ada_norm_zero”.
final_dropout (bool optional, defaults to False) – Whether to apply a final dropout after the last feed-forward layer.
attention_type (str, optional, defaults to “default”) – The type of attention to use. Can be “default” or “gated” or “gated-text-image”.
positional_embeddings (str, optional, defaults to None) – The type of positional embeddings to apply to.
num_positional_embeddings (int, optional, defaults to None) – The maximum number of positional embeddings to apply.

forward(hidden_states: FloatTensor, attention_mask: FloatTensor | None = None, encoder_hidden_states: FloatTensor | None = None, encoder_attention_mask: FloatTensor | None = None, timestep: LongTensor | None = None, cross_attention_kwargs: Dict[str, Any] = None, class_labels: LongTensor | None = None) → FloatTensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

set_chunk_feed_forward(chunk_size: int | None, dim: int = 0)[source]¶

class sketchkit.colorization.manga_ninja.manganinja.models.attention.TemporalBasicTransformerBlock(dim: int, num_attention_heads: int, attention_head_dim: int, dropout=0.0, cross_attention_dim: int | None = None, activation_fn: str = 'geglu', num_embeds_ada_norm: int | None = None, attention_bias: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, unet_use_cross_frame_attention=None, unet_use_temporal_attention=None)[source]¶

Bases: Module

forward(hidden_states, encoder_hidden_states=None, timestep=None, attention_mask=None, video_length=None)[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

sketchkit.colorization.manga_ninja.manganinja.models.attention_processor module¶

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.Attention(query_dim: int, cross_attention_dim: int | None = None, heads: int = 8, dim_head: int = 64, dropout: float = 0.0, bias: bool = False, upcast_attention: bool = False, upcast_softmax: bool = False, cross_attention_norm: str | None = None, cross_attention_norm_num_groups: int = 32, added_kv_proj_dim: int | None = None, norm_num_groups: int | None = None, spatial_norm_dim: int | None = None, out_bias: bool = True, scale_qk: bool = True, only_cross_attention: bool = False, eps: float = 1e-05, rescale_output_factor: float = 1.0, residual_connection: bool = False, _from_deprecated_attn_block: bool = False, processor: AttnProcessor | None = None, out_dim: int = None)[source]¶

Bases: Module

A cross attention layer.

Parameters:

query_dim (int) – The number of channels in the query.
cross_attention_dim (int, optional) – The number of channels in the encoder_hidden_states. If not given, defaults to query_dim.
heads (int, optional, defaults to 8) – The number of heads to use for multi-head attention.
dim_head (int, optional, defaults to 64) – The number of channels in each head.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.
bias (bool, optional, defaults to False) – Set to True for the query, key, and value linear layers to contain a bias parameter.
upcast_attention (bool, optional, defaults to False) – Set to True to upcast the attention computation to float32.
upcast_softmax (bool, optional, defaults to False) – Set to True to upcast the softmax computation to float32.
cross_attention_norm (str, optional, defaults to None) – The type of normalization to use for the cross attention. Can be None, layer_norm, or group_norm.
cross_attention_norm_num_groups (int, optional, defaults to 32) – The number of groups to use for the group norm in the cross attention.
added_kv_proj_dim (int, optional, defaults to None) – The number of channels to use for the added key and value projections. If None, no projection is used.
norm_num_groups (int, optional, defaults to None) – The number of groups to use for the group norm in the attention.
spatial_norm_dim (int, optional, defaults to None) – The number of channels to use for the spatial normalization.
out_bias (bool, optional, defaults to True) – Set to True to use a bias in the output linear layer.
scale_qk (bool, optional, defaults to True) – Set to True to scale the query and key by 1 / sqrt(dim_head).
only_cross_attention (bool, optional, defaults to False) – Set to True to only use cross attention and not added_kv_proj_dim. Can only be set to True if added_kv_proj_dim is not None.
eps (float, optional, defaults to 1e-5) – An additional value added to the denominator in group normalization that is used for numerical stability.
rescale_output_factor (float, optional, defaults to 1.0) – A factor to rescale the output by dividing it with this value.
residual_connection (bool, optional, defaults to False) – Set to True to add the residual connection to the output.
_from_deprecated_attn_block (bool, optional, defaults to False) – Set to True if the attention block is loaded from a deprecated state dict.
processor (AttnProcessor, optional, defaults to None) – The attention processor to use. If None, defaults to AttnProcessor2_0 if torch 2.x is used and AttnProcessor otherwise.

batch_to_head_dim(tensor: Tensor) → Tensor[source]¶

Reshape the tensor from [batch_size, seq_len, dim] to [batch_size // heads, seq_len, dim * heads]. heads is the number of heads initialized while constructing the Attention class.

Parameters:: tensor (torch.Tensor) – The tensor to reshape.
Returns:: The reshaped tensor.
Return type:: torch.Tensor

forward(hidden_states: FloatTensor, encoder_hidden_states: FloatTensor | None = None, encoder_hidden_states_v: FloatTensor | None = None, attention_mask: FloatTensor | None = None, **cross_attention_kwargs) → Tensor[source]¶

The forward method of the Attention class.

Parameters:

hidden_states (torch.Tensor) – The hidden states of the query.
encoder_hidden_states (torch.Tensor, optional) – The hidden states of the encoder.
attention_mask (torch.Tensor, optional) – The attention mask to use. If None, no mask is applied.
**cross_attention_kwargs – Additional keyword arguments to pass along to the cross attention.

Returns:

The output of the attention layer.

Return type:

torch.Tensor

fuse_projections(fuse=True)[source]¶

get_attention_scores(query: Tensor, key: Tensor, attention_mask: Tensor = None) → Tensor[source]¶

Compute the attention scores.

Parameters:

query (torch.Tensor) – The query tensor.
key (torch.Tensor) – The key tensor.
attention_mask (torch.Tensor, optional) – The attention mask to use. If None, no mask is applied.

Returns:

The attention probabilities/scores.

Return type:

torch.Tensor

Get the attention processor in use.

Parameters:: return_deprecated_lora (bool, optional, defaults to False) – Set to True to return the deprecated LoRA attention processor.
Returns:: The attention processor in use.
Return type:: “AttentionProcessor”

head_to_batch_dim(tensor: Tensor, out_dim: int = 3) → Tensor[source]¶

Reshape the tensor from [batch_size, seq_len, dim] to [batch_size, seq_len, heads, dim // heads] heads is the number of heads initialized while constructing the Attention class.

Parameters:

tensor (torch.Tensor) – The tensor to reshape.
out_dim (int, optional, defaults to 3) – The output dimension of the tensor. If 3, the tensor is reshaped to [batch_size * heads, seq_len, dim // heads].

Returns:

The reshaped tensor.

Return type:

torch.Tensor

norm_encoder_hidden_states(encoder_hidden_states: Tensor) → Tensor[source]¶

Normalize the encoder hidden states. Requires self.norm_cross to be specified when constructing the Attention class.

Parameters:: encoder_hidden_states (torch.Tensor) – Hidden states of the encoder.
Returns:: The normalized encoder hidden states.
Return type:: torch.Tensor

prepare_attention_mask(attention_mask: Tensor, target_length: int, batch_size: int, out_dim: int = 3) → Tensor[source]¶

Prepare the attention mask for the attention computation.

Parameters:

attention_mask (torch.Tensor) – The attention mask to prepare.
target_length (int) – The target length of the attention mask. This is the length of the attention mask after padding.
batch_size (int) – The batch size, which is used to repeat the attention mask.
out_dim (int, optional, defaults to 3) – The output dimension of the attention mask. Can be either 3 or 4.

Returns:

The prepared attention mask.

Return type:

torch.Tensor

set_attention_slice(slice_size: int) → None[source]¶

Set the slice size for attention computation.

Parameters:: slice_size (int) – The slice size for attention computation.

set_processor(processor: AttnProcessor) → None[source]¶

Set the attention processor to use.

Parameters:: processor (AttnProcessor) – The attention processor to use.

set_use_memory_efficient_attention_xformers(use_memory_efficient_attention_xformers: bool, attention_op: Callable | None = None) → None[source]¶

Set whether to use memory efficient attention from xformers or not.

Parameters:

use_memory_efficient_attention_xformers (bool) – Whether to use memory efficient attention from xformers or not.
attention_op (Callable, optional) – The attention operation to use. Defaults to None which uses the default attention operation from xformers.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.AttnAddedKVProcessor[source]¶

Bases: object

Processor for performing attention-related computations with extra learnable key and value matrices for the text encoder.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.AttnAddedKVProcessor2_0[source]¶

Bases: object

Processor for performing scaled dot-product attention (enabled by default if you’re using PyTorch 2.0), with extra learnable key and value matrices for the text encoder.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.AttnProcessor[source]¶

Bases: object

Default processor for performing attention-related computations.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.AttnProcessor2_0[source]¶

Bases: object

Processor for implementing scaled dot-product attention (enabled by default if you’re using PyTorch 2.0).

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.CustomDiffusionAttnProcessor(train_kv: bool = True, train_q_out: bool = True, hidden_size: int | None = None, cross_attention_dim: int | None = None, out_bias: bool = True, dropout: float = 0.0)[source]¶

Bases: Module

Processor for implementing attention for the Custom Diffusion method.

Parameters:

train_kv (bool, defaults to True) – Whether to newly train the key and value matrices corresponding to the text features.
train_q_out (bool, defaults to True) – Whether to newly train query matrices corresponding to the latent image features.
hidden_size (int, optional, defaults to None) – The hidden size of the attention layer.
cross_attention_dim (int, optional, defaults to None) – The number of channels in the encoder_hidden_states.
out_bias (bool, defaults to True) – Whether to include the bias parameter in train_q_out.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.CustomDiffusionAttnProcessor2_0(train_kv: bool = True, train_q_out: bool = True, hidden_size: int | None = None, cross_attention_dim: int | None = None, out_bias: bool = True, dropout: float = 0.0)[source]¶

Bases: Module

Processor for implementing attention for the Custom Diffusion method using PyTorch 2.0’s memory-efficient scaled dot-product attention.

Parameters:

train_kv (bool, defaults to True) – Whether to newly train the key and value matrices corresponding to the text features.
train_q_out (bool, defaults to True) – Whether to newly train query matrices corresponding to the latent image features.
hidden_size (int, optional, defaults to None) – The hidden size of the attention layer.
cross_attention_dim (int, optional, defaults to None) – The number of channels in the encoder_hidden_states.
out_bias (bool, defaults to True) – Whether to include the bias parameter in train_q_out.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.CustomDiffusionXFormersAttnProcessor(train_kv: bool = True, train_q_out: bool = False, hidden_size: int | None = None, cross_attention_dim: int | None = None, out_bias: bool = True, dropout: float = 0.0, attention_op: Callable | None = None)[source]¶

Bases: Module

Processor for implementing memory efficient attention using xFormers for the Custom Diffusion method.

Args: train_kv (bool, defaults to True):

Whether to newly train the key and value matrices corresponding to the text features.

train_q_out (bool, defaults to True):: Whether to newly train query matrices corresponding to the latent image features.
hidden_size (int, optional, defaults to None):: The hidden size of the attention layer.
cross_attention_dim (int, optional, defaults to None):: The number of channels in the encoder_hidden_states.
out_bias (bool, defaults to True):: Whether to include the bias parameter in train_q_out.
dropout (float, optional, defaults to 0.0):: The dropout probability to use.
attention_op (Callable, optional, defaults to None):: The base [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use as the attention operator. It is recommended to set to None, and allow xFormers to choose the best operator.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.FusedAttnProcessor2_0[source]¶

Bases: object

Processor for implementing scaled dot-product attention (enabled by default if you’re using PyTorch 2.0). It uses fused projection layers. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused.

This API is currently 🧪 experimental in nature and can change in future.

</Tip>

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.IPAdapterAttnProcessor(hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0)[source]¶

Bases: Module

Attention processor for Multiple IP-Adapater.

Parameters:

hidden_size (int) – The hidden size of the attention layer.
cross_attention_dim (int) – The number of channels in the encoder_hidden_states.
num_tokens (int, Tuple[int] or List[int], defaults to (4,)) – The context length of the image features.
scale (float or List[float], defaults to 1.0) – the weight scale of image prompt.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.IPAdapterAttnProcessor2_0(hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0)[source]¶

Bases: Module

Attention processor for IP-Adapater for PyTorch 2.0.

Parameters:

hidden_size (int) – The hidden size of the attention layer.
cross_attention_dim (int) – The number of channels in the encoder_hidden_states.
num_tokens (int, Tuple[int] or List[int], defaults to (4,)) – The context length of the image features.
scale (float or List[float], defaults to 1.0) – the weight scale of image prompt.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.LoRAAttnAddedKVProcessor(hidden_size: int, cross_attention_dim: int | None = None, rank: int = 4, network_alpha: int | None = None)[source]¶

Bases: Module

Processor for implementing the LoRA attention mechanism with extra learnable key and value matrices for the text encoder.

Parameters:

hidden_size (int, optional) – The hidden size of the attention layer.
cross_attention_dim (int, optional, defaults to None) – The number of channels in the encoder_hidden_states.
rank (int, defaults to 4) – The dimension of the LoRA update matrices.
network_alpha (int, optional) – Equivalent to alpha but it’s usage is specific to Kohya (A1111) style LoRAs.
kwargs (dict) – Additional keyword arguments to pass to the LoRALinearLayer layers.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.LoRAAttnProcessor(hidden_size: int, cross_attention_dim: int | None = None, rank: int = 4, network_alpha: int | None = None, **kwargs)[source]¶: Bases: Module

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.LoRAAttnProcessor2_0(hidden_size: int, cross_attention_dim: int | None = None, rank: int = 4, network_alpha: int | None = None, **kwargs)[source]¶: Bases: Module

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.LoRAXFormersAttnProcessor(hidden_size: int, cross_attention_dim: int, rank: int = 4, attention_op: Callable | None = None, network_alpha: int | None = None, **kwargs)[source]¶

Bases: Module

Processor for implementing the LoRA attention mechanism with memory efficient attention using xFormers.

Parameters:

hidden_size (int, optional) – The hidden size of the attention layer.
cross_attention_dim (int, optional) – The number of channels in the encoder_hidden_states.
rank (int, defaults to 4) – The dimension of the LoRA update matrices.
attention_op (Callable, optional, defaults to None) – The base [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use as the attention operator. It is recommended to set to None, and allow xFormers to choose the best operator.
network_alpha (int, optional) – Equivalent to alpha but it’s usage is specific to Kohya (A1111) style LoRAs.
kwargs (dict) – Additional keyword arguments to pass to the LoRALinearLayer layers.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.SlicedAttnAddedKVProcessor(slice_size)[source]¶

Bases: object

Processor for implementing sliced attention with extra learnable key and value matrices for the text encoder.

Parameters:: slice_size (int, optional) – The number of steps to compute attention. Uses as many slices as attention_head_dim // slice_size, and attention_head_dim must be a multiple of the slice_size.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.SlicedAttnProcessor(slice_size: int)[source]¶

Bases: object

Processor for implementing sliced attention.

Parameters:: slice_size (int, optional) – The number of steps to compute attention. Uses as many slices as attention_head_dim // slice_size, and attention_head_dim must be a multiple of the slice_size.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.SpatialNorm(f_channels: int, zq_channels: int)[source]¶

Bases: Module

Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002.

Parameters:

f_channels (int) – The number of channels for input to group normalization layer, and output of the spatial norm layer.
zq_channels (int) – The number of channels for the quantized vector as described in the paper.

forward(f: FloatTensor, zq: FloatTensor) → FloatTensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.XFormersAttnAddedKVProcessor(attention_op: Callable | None = None)[source]¶

Bases: object

Processor for implementing memory efficient attention using xFormers.

Parameters:: attention_op (Callable, optional, defaults to None) – The base [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use as the attention operator. It is recommended to set to None, and allow xFormers to choose the best operator.

class sketchkit.colorization.manga_ninja.manganinja.models.attention_processor.XFormersAttnProcessor(attention_op: Callable | None = None)[source]¶

Bases: object

Processor for implementing memory efficient attention using xFormers.

Parameters:: attention_op (Callable, optional, defaults to None) – The base [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use as the attention operator. It is recommended to set to None, and allow xFormers to choose the best operator.

sketchkit.colorization.manga_ninja.manganinja.models.mutual_self_attention_multi_scale module¶

class sketchkit.colorization.manga_ninja.manganinja.models.mutual_self_attention_multi_scale.ReferenceAttentionControl(unet, mode='write', do_classifier_free_guidance=False, attention_auto_machine_weight=inf, gn_auto_machine_weight=1.0, style_fidelity=1.0, reference_attn=True, reference_adain=False, fusion_blocks='midup', batch_size=1)[source]¶

Bases: object

clear()[source]¶

register_reference_hooks(mode, do_classifier_free_guidance, attention_auto_machine_weight, gn_auto_machine_weight, style_fidelity, reference_attn, reference_adain, dtype=torch.float16, batch_size=1, num_images_per_prompt=1, device=device(type='cpu'), fusion_blocks='midup')[source]¶

update(writer, point_embedding_ref=None, point_embedding_main=None, dtype=torch.float16)[source]¶

sketchkit.colorization.manga_ninja.manganinja.models.mutual_self_attention_multi_scale.filter_matrices_by_size(matrix_list, reference_matrix)[source]¶

sketchkit.colorization.manga_ninja.manganinja.models.mutual_self_attention_multi_scale.torch_dfs(model: Module)[source]¶

sketchkit.colorization.manga_ninja.manganinja.models.refunet_2d_condition module¶

class sketchkit.colorization.manga_ninja.manganinja.models.refunet_2d_condition.RefUNet2DConditionModel(sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, down_block_types: Tuple[str] = ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D'), mid_block_type: str | None = 'UNetMidBlock2DCrossAttn', up_block_types: Tuple[str] = ('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D'), only_cross_attention: bool | Tuple[bool] = False, block_out_channels: Tuple[int] = (320, 640, 1280, 1280), layers_per_block: int | Tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, dropout: float = 0.0, act_fn: str = 'silu', norm_num_groups: int | None = 32, norm_eps: float = 1e-05, cross_attention_dim: int | Tuple[int] = 1280, transformer_layers_per_block: int | Tuple[int] | Tuple[Tuple] = 1, reverse_transformer_layers_per_block: Tuple[Tuple[int]] | None = None, encoder_hid_dim: int | None = None, encoder_hid_dim_type: str | None = None, attention_head_dim: int | Tuple[int] = 8, num_attention_heads: int | Tuple[int] | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, class_embed_type: str | None = None, addition_embed_type: str | None = None, addition_time_embed_dim: int | None = None, num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = 'default', resnet_skip_time_act: bool = False, resnet_out_scale_factor: int = 1.0, time_embedding_type: str = 'positional', time_embedding_dim: int | None = None, time_embedding_act_fn: str | None = None, timestep_post_act: str | None = None, time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, projection_class_embeddings_input_dim: int | None = None, attention_type: str = 'default', class_embeddings_concat: bool = False, mid_block_only_cross_attention: bool | None = None, cross_attention_norm: str | None = None, addition_embed_type_num_heads=64)[source]¶

Bases: ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin

A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output.

This model inherits from [ModelMixin]. Check the superclass documentation for it’s generic methods implemented for all models (such as downloading or saving).

Parameters:

sample_size (int or Tuple[int, int], optional, defaults to None) – Height and width of input/output sample.
in_channels (int, optional, defaults to 4) – Number of channels in the input sample.
out_channels (int, optional, defaults to 4) – Number of channels in the output.
center_input_sample (bool, optional, defaults to False) – Whether to center the input sample.
flip_sin_to_cos (bool, optional, defaults to False) – Whether to flip the sin to cos in the time embedding.
freq_shift (int, optional, defaults to 0) – The frequency shift to apply to the time embedding.
down_block_types (Tuple[str], optional, defaults to (“CrossAttnDownBlock2D”, “CrossAttnDownBlock2D”, “CrossAttnDownBlock2D”, “DownBlock2D”)) – The tuple of downsample blocks to use.
mid_block_type (str, optional, defaults to “UNetMidBlock2DCrossAttn”) – Block type for middle of UNet, it can be one of UNetMidBlock2DCrossAttn, UNetMidBlock2D, or UNetMidBlock2DSimpleCrossAttn. If None, the mid block layer is skipped.
up_block_types (Tuple[str], optional, defaults to (“UpBlock2D”, “CrossAttnUpBlock2D”, “CrossAttnUpBlock2D”, “CrossAttnUpBlock2D”)) – The tuple of upsample blocks to use.
only_cross_attention (bool or Tuple[bool], optional, default to False) – Whether to include self-attention in the basic transformer blocks, see [~models.attention.BasicTransformerBlock].
block_out_channels (Tuple[int], optional, defaults to (320, 640, 1280, 1280)) – The tuple of output channels for each block.
layers_per_block (int, optional, defaults to 2) – The number of layers per block.
downsample_padding (int, optional, defaults to 1) – The padding to use for the downsampling convolution.
mid_block_scale_factor (float, optional, defaults to 1.0) – The scale factor to use for the mid block.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.
act_fn (str, optional, defaults to “silu”) – The activation function to use.
norm_num_groups (int, optional, defaults to 32) – The number of groups to use for the normalization. If None, normalization and activation layers is skipped in post-processing.
norm_eps (float, optional, defaults to 1e-5) – The epsilon to use for the normalization.
cross_attention_dim (int or Tuple[int], optional, defaults to 1280) – The dimension of the cross attention features.
transformer_layers_per_block – The number of transformer blocks of type [~models.attention.BasicTransformerBlock]. Only relevant for [~models.unet_2d_blocks.CrossAttnDownBlock2D], [~models.unet_2d_blocks.CrossAttnUpBlock2D], [~models.unet_2d_blocks.UNetMidBlock2DCrossAttn].

_set_gradient_checkpointing(module, value=False)[source]¶

_supports_gradient_checkpointing = True¶

property attn_processors: Dict[str, AttnProcessor | CustomDiffusionAttnProcessor | AttnAddedKVProcessor | AttnAddedKVProcessor2_0 | JointAttnProcessor2_0 | PAGJointAttnProcessor2_0 | PAGCFGJointAttnProcessor2_0 | FusedJointAttnProcessor2_0 | AllegroAttnProcessor2_0 | AuraFlowAttnProcessor2_0 | FusedAuraFlowAttnProcessor2_0 | FluxAttnProcessor2_0 | FluxAttnProcessor2_0_NPU | FusedFluxAttnProcessor2_0 | FusedFluxAttnProcessor2_0_NPU | CogVideoXAttnProcessor2_0 | FusedCogVideoXAttnProcessor2_0 | XFormersAttnAddedKVProcessor | XFormersAttnProcessor | XLAFlashAttnProcessor2_0 | AttnProcessorNPU | AttnProcessor2_0 | MochiVaeAttnProcessor2_0 | MochiAttnProcessor2_0 | StableAudioAttnProcessor2_0 | HunyuanAttnProcessor2_0 | FusedHunyuanAttnProcessor2_0 | PAGHunyuanAttnProcessor2_0 | PAGCFGHunyuanAttnProcessor2_0 | LuminaAttnProcessor2_0 | FusedAttnProcessor2_0 | CustomDiffusionXFormersAttnProcessor | CustomDiffusionAttnProcessor2_0 | SlicedAttnProcessor | SlicedAttnAddedKVProcessor | SanaLinearAttnProcessor2_0 | PAGCFGSanaLinearAttnProcessor2_0 | PAGIdentitySanaLinearAttnProcessor2_0 | SanaMultiscaleLinearAttention | SanaMultiscaleAttnProcessor2_0 | SanaMultiscaleAttentionProjection | IPAdapterAttnProcessor | IPAdapterAttnProcessor2_0 | IPAdapterXFormersAttnProcessor | SD3IPAdapterJointAttnProcessor2_0 | PAGIdentitySelfAttnProcessor2_0 | PAGCFGIdentitySelfAttnProcessor2_0 | LoRAAttnProcessor | LoRAAttnProcessor2_0 | LoRAXFormersAttnProcessor | LoRAAttnAddedKVProcessor]¶: Returns: dict of attention processors: A dictionary containing all attention processors used in the model with indexed by its weight name.

disable_freeu()[source]¶: Disables the FreeU mechanism.

enable_freeu(s1, s2, b1, b2)[source]¶

Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.

The suffixes after the scaling factors represent the stage blocks where they are being applied.

Please refer to the [official repository](https://github.com/ChenyangSi/FreeU) for combinations of values that are known to work well for different pipelines such as Stable Diffusion v1, v2, and Stable Diffusion XL.

Parameters:

s1 (float) – Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to mitigate the “oversmoothing effect” in the enhanced denoising process.
s2 (float) – Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to mitigate the “oversmoothing effect” in the enhanced denoising process.
b1 (float) – Scaling factor for stage 1 to amplify the contributions of backbone features.
b2 (float) – Scaling factor for stage 2 to amplify the contributions of backbone features.

forward(sample: FloatTensor, timestep: Tensor | float | int, encoder_hidden_states: Tensor, class_labels: Tensor | None = None, timestep_cond: Tensor | None = None, attention_mask: Tensor | None = None, cross_attention_kwargs: Dict[str, Any] | None = None, added_cond_kwargs: Dict[str, Tensor] | None = None, down_block_additional_residuals: Tuple[Tensor] | None = None, mid_block_additional_residual: Tensor | None = None, down_intrablock_additional_residuals: Tuple[Tensor] | None = None, encoder_attention_mask: Tensor | None = None, return_dict: bool = True) → UNet2DConditionOutput | Tuple[source]¶

The [UNet2DConditionModel] forward method.

Parameters:

sample (torch.FloatTensor) – The noisy input tensor with the following shape (batch, channel, height, width).
timestep (torch.FloatTensor or float or int) – The number of timesteps to denoise an input.
encoder_hidden_states (torch.FloatTensor) – The encoder hidden states with shape (batch, sequence_length, feature_dim).
class_labels (torch.Tensor, optional, defaults to None) – Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
timestep_cond – (torch.Tensor, optional, defaults to None): Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed through the self.time_embedding layer to obtain the timestep embeddings.
attention_mask (torch.Tensor, optional, defaults to None) – An attention mask of shape (batch, key_tokens) is applied to encoder_hidden_states. If 1 the mask is kept, otherwise if 0 it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to “discard” tokens.
cross_attention_kwargs (dict, optional) – A kwargs dictionary that if specified is passed along to the AttentionProcessor as defined under self.processor in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
added_cond_kwargs – (dict, optional): A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks.
down_block_additional_residuals (tuple of torch.Tensor, optional) – (tuple of torch.Tensor, optional): A tuple of tensors that if specified are added to the residuals of down unet blocks.
mid_block_additional_residual (torch.Tensor, optional) – (torch.Tensor, optional): A tensor that if specified is added to the residual of the middle unet block.
encoder_attention_mask (torch.Tensor) – A cross-attention mask of shape (batch, sequence_length) is applied to encoder_hidden_states. If True the mask is kept, otherwise if False it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to “discard” tokens.
return_dict (bool, optional, defaults to True) – Whether or not to return a [~models.unet_2d_condition.UNet2DConditionOutput] instead of a plain tuple.
cross_attention_kwargs – A kwargs dictionary that if specified is passed along to the [AttnProcessor].
added_cond_kwargs – (dict, optional): A kwargs dictionary containin additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks.
down_block_additional_residuals – additional residuals to be added to UNet long skip connections from down blocks to up blocks for example from ControlNet side model(s)
mid_block_additional_residual – additional residual to be added to UNet mid block output, for example from ControlNet side model
down_intrablock_additional_residuals (tuple of torch.Tensor, optional) – additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)

Returns:

If return_dict is True, an [~models.unet_2d_condition.UNet2DConditionOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Return type:

[~models.unet_2d_condition.UNet2DConditionOutput] or tuple

set_attention_slice(slice_size)[source]¶

Enable sliced attention computation.

When this option is enabled, the attention module splits the input tensor in slices to compute attention in several steps. This is useful for saving some memory in exchange for a small decrease in speed.

Parameters:: slice_size (str or int or list(int), optional, defaults to “auto”) – When “auto”, input to the attention heads is halved, so attention is computed in two steps. If “max”, maximum amount of memory is saved by running only one slice at a time. If a number is provided, uses as many slices as attention_head_dim // slice_size. In this case, attention_head_dim must be a multiple of slice_size.

Sets the attention processor to use to compute attention.

Parameters:

processor (dict of AttentionProcessor or only AttentionProcessor) –

The instantiated processor class or a dictionary of processor classes that will be set as the processor for all Attention layers.

If processor is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.

set_default_attn_processor()[source]¶: Disables custom attention processors and sets the default attention implementation.

class sketchkit.colorization.manga_ninja.manganinja.models.refunet_2d_condition.UNet2DConditionOutput(sample: FloatTensor = None, ref_features: Tuple[FloatTensor] = None)[source]¶

Bases: BaseOutput

The output of [UNet2DConditionModel].

Parameters:: sample (torch.FloatTensor of shape (batch_size, num_channels, height, width)) – The hidden states output conditioned on encoder_hidden_states input. Output of last layer of model.

ref_features: Tuple[FloatTensor] = None¶

sample: FloatTensor = None¶

sketchkit.colorization.manga_ninja.manganinja.models.transformer_2d module¶

class sketchkit.colorization.manga_ninja.manganinja.models.transformer_2d.Transformer2DModel(num_attention_heads: int = 16, attention_head_dim: int = 88, in_channels: int | None = None, out_channels: int | None = None, num_layers: int = 1, dropout: float = 0.0, norm_num_groups: int = 32, cross_attention_dim: int | None = None, attention_bias: bool = False, sample_size: int | None = None, num_vector_embeds: int | None = None, patch_size: int | None = None, activation_fn: str = 'geglu', num_embeds_ada_norm: int | None = None, use_linear_projection: bool = False, only_cross_attention: bool = False, double_self_attention: bool = False, upcast_attention: bool = False, norm_type: str = 'layer_norm', norm_elementwise_affine: bool = True, norm_eps: float = 1e-05, attention_type: str = 'default', caption_channels: int = None)[source]¶

Bases: ModelMixin, ConfigMixin

A 2D Transformer model for image-like data.

Parameters:

num_attention_heads (int, optional, defaults to 16) – The number of heads to use for multi-head attention.
attention_head_dim (int, optional, defaults to 88) – The number of channels in each head.
in_channels (int, optional) – The number of channels in the input and output (specify if the input is continuous).
num_layers (int, optional, defaults to 1) – The number of layers of Transformer blocks to use.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.
cross_attention_dim (int, optional) – The number of encoder_hidden_states dimensions to use.
sample_size (int, optional) – The width of the latent images (specify if the input is discrete). This is fixed during training since it is used to learn a number of position embeddings.
num_vector_embeds (int, optional) – The number of classes of the vector embeddings of the latent pixels (specify if the input is discrete). Includes the class for the masked latent pixel.
activation_fn (str, optional, defaults to “geglu”) – Activation function to use in feed-forward.
num_embeds_ada_norm (int, optional) –
The number of diffusion steps used during training. Pass if at least one of the norm_layers is AdaLayerNorm. This is fixed during training since it is used to learn a number of embeddings that are added to the hidden states.

During inference, you can denoise for up to but not more steps than num_embeds_ada_norm.
attention_bias (bool, optional) – Configure if the TransformerBlocks attention should contain a bias parameter.

_set_gradient_checkpointing(module, value=False)[source]¶

_supports_gradient_checkpointing = True¶

forward(hidden_states: Tensor, encoder_hidden_states: Tensor | None = None, timestep: LongTensor | None = None, added_cond_kwargs: Dict[str, Tensor] = None, class_labels: LongTensor | None = None, cross_attention_kwargs: Dict[str, Any] = None, attention_mask: Tensor | None = None, encoder_attention_mask: Tensor | None = None, return_dict: bool = True)[source]¶

The [Transformer2DModel] forward method.

Parameters:

hidden_states (torch.LongTensor of shape (batch size, num latent pixels) if discrete, torch.FloatTensor of shape (batch size, channel, height, width) if continuous) – Input hidden_states.
encoder_hidden_states (torch.FloatTensor of shape (batch size, sequence len, embed dims), optional) – Conditional embeddings for cross attention layer. If not given, cross-attention defaults to self-attention.
timestep (torch.LongTensor, optional) – Used to indicate denoising step. Optional timestep to be applied as an embedding in AdaLayerNorm.
class_labels (torch.LongTensor of shape (batch size, num classes), optional) – Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in AdaLayerZeroNorm.
cross_attention_kwargs (Dict[str, Any], optional) – A kwargs dictionary that if specified is passed along to the AttentionProcessor as defined under self.processor in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
attention_mask (torch.Tensor, optional) – An attention mask of shape (batch, key_tokens) is applied to encoder_hidden_states. If 1 the mask is kept, otherwise if 0 it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to “discard” tokens.
encoder_attention_mask (torch.Tensor, optional) –
Cross-attention mask applied to encoder_hidden_states. Two formats supported:
- Mask (batch, sequence_length) True = keep, False = discard.
- Bias (batch, 1, sequence_length) 0 = keep, -10000 = discard.
If ndim == 2: will be interpreted as a mask, then converted into a bias consistent with the format above. This bias will be added to the cross-attention scores.
return_dict (bool, optional, defaults to True) – Whether or not to return a [~models.unet_2d_condition.UNet2DConditionOutput] instead of a plain tuple.

Returns:

If return_dict is True, an [~models.transformer_2d.Transformer2DModelOutput] is returned, otherwise a tuple where the first element is the sample tensor.

class sketchkit.colorization.manga_ninja.manganinja.models.transformer_2d.Transformer2DModelOutput(sample: FloatTensor, ref_feature: FloatTensor)[source]¶

Bases: BaseOutput

The output of [Transformer2DModel].

Parameters:: sample (torch.FloatTensor of shape (batch_size, num_channels, height, width) or (batch size, num_vector_embeds - 1, num_latent_pixels) if [Transformer2DModel] is discrete) – The hidden states output conditioned on the encoder_hidden_states input. If discrete, returns probability distributions for the unnoised latent pixels.

ref_feature: FloatTensor¶

sample: FloatTensor¶

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks module¶

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.AutoencoderTinyBlock(in_channels: int, out_channels: int, act_fn: str)[source]¶

Bases: Module

Tiny Autoencoder block used in [AutoencoderTiny]. It is a mini residual module consisting of plain conv + ReLU blocks.

Parameters:

in_channels (int) – The number of input channels.
out_channels (int) – The number of output channels.
act_fn (str) – ` The activation function to use. Supported values are “swish”, “mish”, “gelu”, and “relu”.

Returns:

A tensor with the same shape as the input tensor, but with the number of channels equal to out_channels.

Return type:

torch.FloatTensor

forward(x: FloatTensor) → FloatTensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.CrossAttnDownBlock2D(in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | Tuple[int] = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, cross_attention_dim: int = 1280, output_scale_factor: float = 1.0, downsample_padding: int = 1, add_downsample: bool = True, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = 'default')[source]¶

Bases: Module

forward(hidden_states: FloatTensor, temb: FloatTensor | None = None, encoder_hidden_states: FloatTensor | None = None, attention_mask: FloatTensor | None = None, cross_attention_kwargs: Dict[str, Any] | None = None, encoder_attention_mask: FloatTensor | None = None, additional_residuals: FloatTensor | None = None) → Tuple[FloatTensor, Tuple[FloatTensor, ...]][source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.CrossAttnUpBlock2D(in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | Tuple[int] = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, cross_attention_dim: int = 1280, output_scale_factor: float = 1.0, add_upsample: bool = True, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, attention_type: str = 'default')[source]¶

Bases: Module

forward(hidden_states: FloatTensor, res_hidden_states_tuple: Tuple[FloatTensor, ...], temb: FloatTensor | None = None, encoder_hidden_states: FloatTensor | None = None, cross_attention_kwargs: Dict[str, Any] | None = None, upsample_size: int | None = None, attention_mask: FloatTensor | None = None, encoder_attention_mask: FloatTensor | None = None) → FloatTensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.DownBlock2D(in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_downsample: bool = True, downsample_padding: int = 1)[source]¶

Bases: Module

forward(hidden_states: FloatTensor, temb: FloatTensor | None = None, scale: float = 1.0) → Tuple[FloatTensor, Tuple[FloatTensor, ...]][source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.UNetMidBlock2D(in_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, attn_groups: int | None = None, resnet_pre_norm: bool = True, add_attention: bool = True, attention_head_dim: int = 1, output_scale_factor: float = 1.0)[source]¶

Bases: Module

A 2D UNet mid-block [UNetMidBlock2D] with multiple residual blocks and optional attention blocks.

Parameters:

in_channels (int) – The number of input channels.
temb_channels (int) – The number of temporal embedding channels.
dropout (float, optional, defaults to 0.0) – The dropout rate.
num_layers (int, optional, defaults to 1) – The number of residual blocks.
resnet_eps (float, optional, 1e-6) – The epsilon value for the resnet blocks.
resnet_time_scale_shift (str, optional, defaults to default) – The type of normalization to apply to the time embeddings. This can help to improve the performance of the model on tasks with long-range temporal dependencies.
resnet_act_fn (str, optional, defaults to swish) – The activation function for the resnet blocks.
resnet_groups (int, optional, defaults to 32) – The number of groups to use in the group normalization layers of the resnet blocks.
attn_groups (Optional[int], optional, defaults to None) – The number of groups for the attention blocks.
resnet_pre_norm (bool, optional, defaults to True) – Whether to use pre-normalization for the resnet blocks.
add_attention (bool, optional, defaults to True) – Whether to add attention blocks.
attention_head_dim (int, optional, defaults to 1) – Dimension of a single attention head. The number of attention heads is determined based on this value and the number of input channels.
output_scale_factor (float, optional, defaults to 1.0) – The output scale factor.

Returns:

The output of the last residual block, which is a tensor of shape (batch_size, in_channels, height, width).

Return type:

torch.FloatTensor

forward(hidden_states: FloatTensor, temb: FloatTensor | None = None) → FloatTensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.UNetMidBlock2DCrossAttn(in_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, transformer_layers_per_block: int | Tuple[int] = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, num_attention_heads: int = 1, output_scale_factor: float = 1.0, cross_attention_dim: int = 1280, dual_cross_attention: bool = False, use_linear_projection: bool = False, upcast_attention: bool = False, attention_type: str = 'default')[source]¶

Bases: Module

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.UpBlock2D(in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: int | None = None, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, output_scale_factor: float = 1.0, add_upsample: bool = True)[source]¶

Bases: Module

forward(hidden_states: FloatTensor, res_hidden_states_tuple: Tuple[FloatTensor, ...], temb: FloatTensor | None = None, upsample_size: int | None = None, scale: float = 1.0) → FloatTensor[source]¶

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.get_down_block(down_block_type: str, num_layers: int, in_channels: int, out_channels: int, temb_channels: int, add_downsample: bool, resnet_eps: float, resnet_act_fn: str, transformer_layers_per_block: int = 1, num_attention_heads: int | None = None, resnet_groups: int | None = None, cross_attention_dim: int | None = None, downsample_padding: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, resnet_time_scale_shift: str = 'default', attention_type: str = 'default', resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, cross_attention_norm: str | None = None, attention_head_dim: int | None = None, downsample_type: str | None = None, dropout: float = 0.0)[source]¶

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.get_mid_block(mid_block_type: str, temb_channels: int, in_channels: int, resnet_eps: float, resnet_act_fn: str, resnet_groups: int, output_scale_factor: float = 1.0, transformer_layers_per_block: int = 1, num_attention_heads: int | None = None, cross_attention_dim: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, mid_block_only_cross_attention: bool = False, upcast_attention: bool = False, resnet_time_scale_shift: str = 'default', attention_type: str = 'default', resnet_skip_time_act: bool = False, cross_attention_norm: str | None = None, attention_head_dim: int | None = 1, dropout: float = 0.0)[source]¶

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks.get_up_block(up_block_type: str, num_layers: int, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, add_upsample: bool, resnet_eps: float, resnet_act_fn: str, resolution_idx: int | None = None, transformer_layers_per_block: int = 1, num_attention_heads: int | None = None, resnet_groups: int | None = None, cross_attention_dim: int | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, only_cross_attention: bool = False, upcast_attention: bool = False, resnet_time_scale_shift: str = 'default', attention_type: str = 'default', resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, cross_attention_norm: str | None = None, attention_head_dim: int | None = None, upsample_type: str | None = None, dropout: float = 0.0) → Module[source]¶

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_condition module¶

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_condition.UNet2DConditionModel(sample_size: int | None = None, in_channels: int = 4, out_channels: int = 4, center_input_sample: bool = False, flip_sin_to_cos: bool = True, freq_shift: int = 0, down_block_types: Tuple[str] = ('CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'CrossAttnDownBlock2D', 'DownBlock2D'), mid_block_type: str | None = 'UNetMidBlock2DCrossAttn', up_block_types: Tuple[str] = ('UpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D', 'CrossAttnUpBlock2D'), only_cross_attention: bool | Tuple[bool] = False, block_out_channels: Tuple[int] = (320, 640, 1280, 1280), layers_per_block: int | Tuple[int] = 2, downsample_padding: int = 1, mid_block_scale_factor: float = 1, dropout: float = 0.0, act_fn: str = 'silu', norm_num_groups: int | None = 32, norm_eps: float = 1e-05, cross_attention_dim: int | Tuple[int] = 1280, transformer_layers_per_block: int | Tuple[int] | Tuple[Tuple] = 1, reverse_transformer_layers_per_block: Tuple[Tuple[int]] | None = None, encoder_hid_dim: int | None = None, encoder_hid_dim_type: str | None = None, attention_head_dim: int | Tuple[int] = 8, num_attention_heads: int | Tuple[int] | None = None, dual_cross_attention: bool = False, use_linear_projection: bool = False, class_embed_type: str | None = None, addition_embed_type: str | None = None, addition_time_embed_dim: int | None = None, num_class_embeds: int | None = None, upcast_attention: bool = False, resnet_time_scale_shift: str = 'default', resnet_skip_time_act: bool = False, resnet_out_scale_factor: float = 1.0, time_embedding_type: str = 'positional', time_embedding_dim: int | None = None, time_embedding_act_fn: str | None = None, timestep_post_act: str | None = None, time_cond_proj_dim: int | None = None, conv_in_kernel: int = 3, conv_out_kernel: int = 3, projection_class_embeddings_input_dim: int | None = None, attention_type: str = 'default', class_embeddings_concat: bool = False, mid_block_only_cross_attention: bool | None = None, cross_attention_norm: str | None = None, addition_embed_type_num_heads: int = 64)[source]¶

Bases: ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin

A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample shaped output.

This model inherits from [ModelMixin]. Check the superclass documentation for it’s generic methods implemented for all models (such as downloading or saving).

Parameters:

sample_size (int or Tuple[int, int], optional, defaults to None) – Height and width of input/output sample.
in_channels (int, optional, defaults to 4) – Number of channels in the input sample.
out_channels (int, optional, defaults to 4) – Number of channels in the output.
center_input_sample (bool, optional, defaults to False) – Whether to center the input sample.
flip_sin_to_cos (bool, optional, defaults to True) – Whether to flip the sin to cos in the time embedding.
freq_shift (int, optional, defaults to 0) – The frequency shift to apply to the time embedding.
down_block_types (Tuple[str], optional, defaults to (“CrossAttnDownBlock2D”, “CrossAttnDownBlock2D”, “CrossAttnDownBlock2D”, “DownBlock2D”)) – The tuple of downsample blocks to use.
mid_block_type (str, optional, defaults to “UNetMidBlock2DCrossAttn”) – Block type for middle of UNet, it can be one of UNetMidBlock2DCrossAttn, UNetMidBlock2D, or UNetMidBlock2DSimpleCrossAttn. If None, the mid block layer is skipped.
up_block_types (Tuple[str], optional, defaults to (“UpBlock2D”, “CrossAttnUpBlock2D”, “CrossAttnUpBlock2D”, “CrossAttnUpBlock2D”)) – The tuple of upsample blocks to use.
only_cross_attention (bool or Tuple[bool], optional, default to False) – Whether to include self-attention in the basic transformer blocks, see [~models.attention.BasicTransformerBlock].
block_out_channels (Tuple[int], optional, defaults to (320, 640, 1280, 1280)) – The tuple of output channels for each block.
layers_per_block (int, optional, defaults to 2) – The number of layers per block.
downsample_padding (int, optional, defaults to 1) – The padding to use for the downsampling convolution.
mid_block_scale_factor (float, optional, defaults to 1.0) – The scale factor to use for the mid block.
dropout (float, optional, defaults to 0.0) – The dropout probability to use.
act_fn (str, optional, defaults to “silu”) – The activation function to use.
norm_num_groups (int, optional, defaults to 32) – The number of groups to use for the normalization. If None, normalization and activation layers is skipped in post-processing.
norm_eps (float, optional, defaults to 1e-5) – The epsilon to use for the normalization.
cross_attention_dim (int or Tuple[int], optional, defaults to 1280) – The dimension of the cross attention features.
transformer_layers_per_block (int, Tuple[int], or Tuple[Tuple] , optional, defaults to 1) – The number of transformer blocks of type [~models.attention.BasicTransformerBlock]. Only relevant for [~models.unet_2d_blocks.CrossAttnDownBlock2D], [~models.unet_2d_blocks.CrossAttnUpBlock2D], [~models.unet_2d_blocks.UNetMidBlock2DCrossAttn].
reverse_transformer_layers_per_block – (Tuple[Tuple], optional, defaults to None): The number of transformer blocks of type [~models.attention.BasicTransformerBlock], in the upsampling blocks of the U-Net. Only relevant if transformer_layers_per_block is of type Tuple[Tuple] and for [~models.unet_2d_blocks.CrossAttnDownBlock2D], [~models.unet_2d_blocks.CrossAttnUpBlock2D], [~models.unet_2d_blocks.UNetMidBlock2DCrossAttn].
encoder_hid_dim (int, optional, defaults to None) – If encoder_hid_dim_type is defined, encoder_hidden_states will be projected from encoder_hid_dim dimension to cross_attention_dim.
encoder_hid_dim_type (str, optional, defaults to None) – If given, the encoder_hidden_states and potentially other embeddings are down-projected to text embeddings of dimension cross_attention according to encoder_hid_dim_type.
attention_head_dim (int, optional, defaults to 8) – The dimension of the attention heads.
num_attention_heads (int, optional) – The number of attention heads. If not defined, defaults to attention_head_dim
resnet_time_scale_shift (str, optional, defaults to “default”) – Time scale shift config for ResNet blocks (see [~models.resnet.ResnetBlock2D]). Choose from default or scale_shift.
class_embed_type (str, optional, defaults to None) – The type of class embedding to use which is ultimately summed with the time embeddings. Choose from None, “timestep”, “identity”, “projection”, or “simple_projection”.
addition_embed_type (str, optional, defaults to None) – Configures an optional embedding which will be summed with the time embeddings. Choose from None or “text”. “text” will use the TextTimeEmbedding layer.
addition_time_embed_dim – (int, optional, defaults to None): Dimension for the timestep embeddings.
num_class_embeds (int, optional, defaults to None) – Input dimension of the learnable embedding matrix to be projected to time_embed_dim, when performing class conditioning with class_embed_type equal to None.
time_embedding_type (str, optional, defaults to positional) – The type of position embedding to use for timesteps. Choose from positional or fourier.
time_embedding_dim (int, optional, defaults to None) – An optional override for the dimension of the projected time embedding.
time_embedding_act_fn (str, optional, defaults to None) – Optional activation function to use only once on the time embeddings before they are passed to the rest of the UNet. Choose from silu, mish, gelu, and swish.
timestep_post_act (str, optional, defaults to None) – The second activation function to use in timestep embedding. Choose from silu, mish and gelu.
time_cond_proj_dim (int, optional, defaults to None) – The dimension of cond_proj layer in the timestep embedding.
conv_in_kernel (int, optional, default to 3) – The kernel size of conv_in layer.
conv_out_kernel (int, optional, default to 3) – The kernel size of conv_out layer.
projection_class_embeddings_input_dim (int, optional) – The dimension of the class_labels input when class_embed_type=”projection”. Required when class_embed_type=”projection”.
class_embeddings_concat (bool, optional, defaults to False) – Whether to concatenate the time embeddings with the class embeddings.
mid_block_only_cross_attention (bool, optional, defaults to None) – Whether to use cross attention with the mid block when using the UNetMidBlock2DSimpleCrossAttn. If only_cross_attention is given as a single boolean and mid_block_only_cross_attention is None, the only_cross_attention value is used as the value for mid_block_only_cross_attention. Default to False otherwise.

_check_config(down_block_types: Tuple[str], up_block_types: Tuple[str], only_cross_attention: bool | Tuple[bool], block_out_channels: Tuple[int], layers_per_block: int | Tuple[int], cross_attention_dim: int | Tuple[int], transformer_layers_per_block: int | Tuple[int] | Tuple[Tuple[int]], reverse_transformer_layers_per_block: bool, attention_head_dim: int, num_attention_heads: int | Tuple[int] | None)[source]¶

_set_add_embedding(addition_embed_type: str, addition_embed_type_num_heads: int, addition_time_embed_dim: int | None, flip_sin_to_cos: bool, freq_shift: float, cross_attention_dim: int | None, encoder_hid_dim: int | None, projection_class_embeddings_input_dim: int | None, time_embed_dim: int)[source]¶

_set_class_embedding(class_embed_type: str | None, act_fn: str, num_class_embeds: int | None, projection_class_embeddings_input_dim: int | None, time_embed_dim: int, timestep_input_dim: int)[source]¶

_set_encoder_hid_proj(encoder_hid_dim_type: str | None, cross_attention_dim: int | Tuple[int], encoder_hid_dim: int | None)[source]¶

_set_gradient_checkpointing(module, value=False)[source]¶

_set_pos_net_if_use_gligen(attention_type: str, cross_attention_dim: int)[source]¶

_set_time_proj(time_embedding_type: str, block_out_channels: int, flip_sin_to_cos: bool, freq_shift: float, time_embedding_dim: int) → Tuple[int, int][source]¶

_supports_gradient_checkpointing = True¶

property attn_processors: Dict[str, AttnProcessor | CustomDiffusionAttnProcessor | AttnAddedKVProcessor | AttnAddedKVProcessor2_0 | JointAttnProcessor2_0 | PAGJointAttnProcessor2_0 | PAGCFGJointAttnProcessor2_0 | FusedJointAttnProcessor2_0 | AllegroAttnProcessor2_0 | AuraFlowAttnProcessor2_0 | FusedAuraFlowAttnProcessor2_0 | FluxAttnProcessor2_0 | FluxAttnProcessor2_0_NPU | FusedFluxAttnProcessor2_0 | FusedFluxAttnProcessor2_0_NPU | CogVideoXAttnProcessor2_0 | FusedCogVideoXAttnProcessor2_0 | XFormersAttnAddedKVProcessor | XFormersAttnProcessor | XLAFlashAttnProcessor2_0 | AttnProcessorNPU | AttnProcessor2_0 | MochiVaeAttnProcessor2_0 | MochiAttnProcessor2_0 | StableAudioAttnProcessor2_0 | HunyuanAttnProcessor2_0 | FusedHunyuanAttnProcessor2_0 | PAGHunyuanAttnProcessor2_0 | PAGCFGHunyuanAttnProcessor2_0 | LuminaAttnProcessor2_0 | FusedAttnProcessor2_0 | CustomDiffusionXFormersAttnProcessor | CustomDiffusionAttnProcessor2_0 | SlicedAttnProcessor | SlicedAttnAddedKVProcessor | SanaLinearAttnProcessor2_0 | PAGCFGSanaLinearAttnProcessor2_0 | PAGIdentitySanaLinearAttnProcessor2_0 | SanaMultiscaleLinearAttention | SanaMultiscaleAttnProcessor2_0 | SanaMultiscaleAttentionProjection | IPAdapterAttnProcessor | IPAdapterAttnProcessor2_0 | IPAdapterXFormersAttnProcessor | SD3IPAdapterJointAttnProcessor2_0 | PAGIdentitySelfAttnProcessor2_0 | PAGCFGIdentitySelfAttnProcessor2_0 | LoRAAttnProcessor | LoRAAttnProcessor2_0 | LoRAXFormersAttnProcessor | LoRAAttnAddedKVProcessor]¶: Returns: dict of attention processors: A dictionary containing all attention processors used in the model with indexed by its weight name.

disable_freeu()[source]¶: Disables the FreeU mechanism.

enable_freeu(s1: float, s2: float, b1: float, b2: float)[source]¶

Enables the FreeU mechanism from https://arxiv.org/abs/2309.11497.

The suffixes after the scaling factors represent the stage blocks where they are being applied.

Parameters:

s1 (float) – Scaling factor for stage 1 to attenuate the contributions of the skip features. This is done to mitigate the “oversmoothing effect” in the enhanced denoising process.
s2 (float) – Scaling factor for stage 2 to attenuate the contributions of the skip features. This is done to mitigate the “oversmoothing effect” in the enhanced denoising process.
b1 (float) – Scaling factor for stage 1 to amplify the contributions of backbone features.
b2 (float) – Scaling factor for stage 2 to amplify the contributions of backbone features.

The [UNet2DConditionModel] forward method.

Parameters:

sample (torch.FloatTensor) – The noisy input tensor with the following shape (batch, channel, height, width).
timestep (torch.FloatTensor or float or int) – The number of timesteps to denoise an input.
encoder_hidden_states (torch.FloatTensor) – The encoder hidden states with shape (batch, sequence_length, feature_dim).
class_labels (torch.Tensor, optional, defaults to None) – Optional class labels for conditioning. Their embeddings will be summed with the timestep embeddings.
timestep_cond – (torch.Tensor, optional, defaults to None): Conditional embeddings for timestep. If provided, the embeddings will be summed with the samples passed through the self.time_embedding layer to obtain the timestep embeddings.
attention_mask (torch.Tensor, optional, defaults to None) – An attention mask of shape (batch, key_tokens) is applied to encoder_hidden_states. If 1 the mask is kept, otherwise if 0 it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to “discard” tokens.
cross_attention_kwargs (dict, optional) – A kwargs dictionary that if specified is passed along to the AttentionProcessor as defined under self.processor in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
added_cond_kwargs – (dict, optional): A kwargs dictionary containing additional embeddings that if specified are added to the embeddings that are passed along to the UNet blocks.
down_block_additional_residuals – (tuple of torch.Tensor, optional): A tuple of tensors that if specified are added to the residuals of down unet blocks.
mid_block_additional_residual – (torch.Tensor, optional): A tensor that if specified is added to the residual of the middle unet block.
down_intrablock_additional_residuals (tuple of torch.Tensor, optional) – additional residuals to be added within UNet down blocks, for example from T2I-Adapter side model(s)
encoder_attention_mask (torch.Tensor) – A cross-attention mask of shape (batch, sequence_length) is applied to encoder_hidden_states. If True the mask is kept, otherwise if False it is discarded. Mask will be converted into a bias, which adds large negative values to the attention scores corresponding to “discard” tokens.
return_dict (bool, optional, defaults to True) – Whether or not to return a [~models.unets.unet_2d_condition.UNet2DConditionOutput] instead of a plain tuple.

Returns:

If return_dict is True, an [~models.unets.unet_2d_condition.UNet2DConditionOutput] is returned, otherwise a tuple is returned where the first element is the sample tensor.

Return type:

[~models.unets.unet_2d_condition.UNet2DConditionOutput] or tuple

fuse_qkv_projections()[source]¶

Enables fused QKV projections. For self-attention modules, all projection matrices (i.e., query, key, value) are fused. For cross-attention modules, key and value projection matrices are fused.

This API is 🧪 experimental.

</Tip>

get_aug_embed(emb: Tensor, encoder_hidden_states: Tensor, added_cond_kwargs: Dict[str, Any]) → Tensor | None[source]¶

get_class_embed(sample: Tensor, class_labels: Tensor | None) → Tensor | None[source]¶

get_time_embed(sample: Tensor, timestep: Tensor | float | int) → Tensor | None[source]¶

process_encoder_hidden_states(encoder_hidden_states: Tensor, added_cond_kwargs: Dict[str, Any]) → Tensor[source]¶

set_attention_slice(slice_size: str | int | List[int] = 'auto')[source]¶

Enable sliced attention computation.

Parameters:: slice_size (str or int or list(int), optional, defaults to “auto”) – When “auto”, input to the attention heads is halved, so attention is computed in two steps. If “max”, maximum amount of memory is saved by running only one slice at a time. If a number is provided, uses as many slices as attention_head_dim // slice_size. In this case, attention_head_dim must be a multiple of slice_size.

Sets the attention processor to use to compute attention.

Parameters:

processor (dict of AttentionProcessor or only AttentionProcessor) –

The instantiated processor class or a dictionary of processor classes that will be set as the processor for all Attention layers.

If processor is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.

set_default_attn_processor()[source]¶: Disables custom attention processors and sets the default attention implementation.

unfuse_qkv_projections()[source]¶

Disables the fused QKV projection if enabled.

This API is 🧪 experimental.

</Tip>

unload_lora()[source]¶: Unloads LoRA weights.

class sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_condition.UNet2DConditionOutput(sample: FloatTensor = None)[source]¶

Bases: BaseOutput

The output of [UNet2DConditionModel].

Parameters:: sample (torch.FloatTensor of shape (batch_size, num_channels, height, width)) – The hidden states output conditioned on encoder_hidden_states input. Output of last layer of model.

sample: FloatTensor = None¶

sketchkit.colorization.manga_ninja.manganinja.models package¶

Submodules¶

sketchkit.colorization.manga_ninja.manganinja.models.attention module¶

sketchkit.colorization.manga_ninja.manganinja.models.attention_processor module¶

sketchkit.colorization.manga_ninja.manganinja.models.mutual_self_attention_multi_scale module¶

sketchkit.colorization.manga_ninja.manganinja.models.refunet_2d_condition module¶

sketchkit.colorization.manga_ninja.manganinja.models.transformer_2d module¶

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_blocks module¶

sketchkit.colorization.manga_ninja.manganinja.models.unet_2d_condition module¶

Module contents¶