Hugging Face's logo

Diffusers documentation

SanaTransformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

SanaTransformer2DModel

Um modelo Diffusion Transformer para dados 2D de SANA: Efficient High-Resolution Image Synthesis with Linear Diffusion Transformers foi introduzido pela NVIDIA e pelo MIT HAN Lab, por Enze Xie, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, Song Han.

O resumo do artigo é:

Apresentamos o Sana, um framework de texto para imagem capaz de gerar imagens com resolução de até 4096×4096 de forma eficiente. O Sana pode sintetizar imagens de alta resolução e alta qualidade, com forte alinhamento entre texto e imagem, a uma velocidade notavelmente rápida, sendo implantável em GPUs de notebook. Os principais designs incluem: (1) Autoencoder de compressão profunda: diferentemente dos AEs tradicionais, que comprimem imagens apenas 8×, treinamos um AE que pode comprimir imagens 32×, reduzindo efetivamente o número de tokens latentes. (2) Linear DiT: substituímos todo o vanilla attention no DiT por linear attention, que é mais eficiente em altas resoluções sem sacrificar a qualidade. (3) Text encoder apenas com decoder: substituímos o T5 por um pequeno LLM moderno apenas com decoder como text encoder e projetamos instruções humanas complexas com in-context learning para aprimorar o alinhamento entre imagem e texto. (4) Treinamento e amostragem eficientes: propomos o Flow-DPM-Solver para reduzir os passos de amostragem, com rotulagem e seleção de legendas eficientes para acelerar a convergência. Como resultado, o Sana-0.6B é muito competitivo em relação a modelos de difusão gigantes modernos (por exemplo, Flux-12B), sendo 20 vezes menor e mais de 100 vezes mais rápido em throughput medido. Além disso, o Sana-0.6B pode ser implantado em uma GPU de notebook de 16GB, levando menos de 1 segundo para gerar uma imagem com resolução de 1024×1024. O Sana permite a criação de conteúdo a baixo custo. O código e o modelo serão publicamente liberados.

O modelo pode ser carregado com o seguinte trecho de código.

from diffusers import SanaTransformer2DModel

transformer = SanaTransformer2DModel.from_pretrained("Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers", subfolder="transformer", dtype=torch.bfloat16)

SanaTransformer2DModel

classe diffusers.SanaTransformer2DModel

< >

( in_channels: int = 32out_channels: int | None = 32num_attention_heads: int = 70attention_head_dim: int = 32num_layers: int = 20num_cross_attention_heads: int | None = 20cross_attention_head_dim: int | None = 112cross_attention_dim: int | None = 2240caption_channels: int = 2304mlp_ratio: float = 2.5dropout: float = 0.0attention_bias: bool = Falsesample_size: int = 32patch_size: int = 1norm_elementwise_affine: bool = Falsenorm_eps: float = 1e-06interpolation_scale: int | None = Noneguidance_embeds: bool = Falseguidance_embeds_scale: float = 0.1qk_norm: str | None = Nonetimestep_scale: float = 1.0 )

Parâmetros

  • in_channels (int, padrão 32) — O número de canais na entrada.
  • out_channels (int, opcional, padrão 32) — O número de canais na saída.
  • num_attention_heads (int, padrão 70) — O número de cabeças a serem usadas para a atenção multi-cabeça.
  • attention_head_dim (int, padrão 32) — O número de canais em cada cabeça.
  • num_layers (int, padrão 20) — O número de blocos de camadas de Transformer a serem usados.
  • num_cross_attention_heads (int, opcional, padrão 20) — O número de cabeças a serem usadas para cross-attention.
  • cross_attention_head_dim (int, opcional, padrão 112) — O número de canais em cada cabeça para a atenção cruzada.
  • cross_attention_dim (int, opcional, padrão 2240) — O número de canais na saída da cross-attention.
  • caption_channels (int, padrão 2304) — O número de canais nos embeddings de legenda.
  • mlp_ratio (float, padrão 2.5) — A razão de expansão a ser usada na camada GLUMBConv.
  • dropout (float, padrão 0.0) — A probabilidade de dropout.
  • attention_bias (bool, padrão False) — Se deve usar viés na camada de atenção.
  • sample_size (int, padrão 32) — O tamanho base do latente de entrada.
  • patch_size (int, padrão 1) — O tamanho dos patches a serem usados na camada de incorporação de patches.
  • norm_elementwise_affine (bool, padrão False) — Se deve usar afinidade elementar na camada de normalização.
  • norm_eps (float, padrão 1e-6) — O valor de epsilon para a camada de normalização.
  • qk_norm (str, opcional, padrão None) — A normalização a ser usada para a query e a key.
  • timestep_scale (float, padrão 1.0) — A escala a ser usada para os timesteps.

Um modelo Transformer 2D introduzido na família de modelos Sana.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorguidance: typing.Optional[torch.Tensor] = Noneencoder_attention_mask: typing.Optional[torch.Tensor] = Noneattention_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonecontrolnet_block_samples: tuple[torch.Tensor] | None = Nonereturn_dict: bool = True )

Parâmetros

  • hidden_states (torch.Tensor de forma (batch_size, in_channels, height, width)) — hidden_states de entrada.
  • encoder_hidden_states (torch.Tensor de forma (batch_size, sequence_len, embed_dims)) — Incorporações condicionais (incorporações calculadas a partir das condições de entrada, como prompts) a serem usadas.
  • timestep (torch.LongTensor) — Usado para indicar o passo de desruido.
  • guidance (torch.Tensor, opcional) — Embedding de escala de orientação.
  • encoder_attention_mask (torch.Tensor, opcional) — Máscara de atenção cruzada aplicada a encoder_hidden_states.
  • attention_mask (torch.Tensor, opcional) — Máscara de autoatenção aplicada a hidden_states.
  • attention_kwargs (dict, opcional) — Um dicionário de kwargs que, se especificado, é passado ao AttentionProcessor conforme definido em self.processor em diffusers.models.attention_processor.
  • controlnet_block_samples (tuple de torch.Tensor, opcional) — Uma lista de tensores que, se especificada, é adicionada aos resíduos dos blocos do transformer.
  • return_dict (bool, opcional, padrão True) — Se deve ou não retornar um ~models.transformer_2d.Transformer2DModelOutput em vez de uma tupla simples.

O método forward do SanaTransformer2DModel.

Transformer2DModelOutput

classe diffusers.models.modeling_outputs.Transformer2DModelOutput

< >

( sample: torch.Tensor )

Parâmetros

  • sample (torch.Tensor de forma (batch_size, num_channels, height, width) ou (batch size, num_vector_embeds - 1, num_latent_pixels) se Transformer2DModel for discreto) — A saída dos estados ocultos condicionada à entrada encoder_hidden_states. Se discreto, retorna distribuições de probabilidade para os pixels latentes sem ruído.

A saída de Transformer2DModel.

Update on GitHub

←RFD4ProteinaTransformer SanaVideoTransformer3DModel→