Method forward

pipelines/modeling_uvit.py:169–347

Args: prompt_embeds (`torch.Tensor` of shape `(batch size, seq_len, text_dim)`): CLIP-embedded text representation. image_latents (`torch.Tensor` of shape `(batch size, latent channels, height, width)`): Latent image representation fro

(
        self,
        prompt_embeds: torch.Tensor,
        image_latents: torch.Tensor,
        image_embeds: torch.Tensor,
        label_latents: torch.Tensor,
        timestep_img: Union[torch.Tensor, float, int],
        timestep_text: Union[torch.Tensor, float, int],
        data_type: Optional[Union[torch.Tensor, float, int]] = 1,
        encoder_hidden_states=None,
        cross_attention_kwargs=None,
    )

        return {"pos_embed"}

    def forward(
        self,
        prompt_embeds: torch.Tensor,
        image_latents: torch.Tensor,
        image_embeds: torch.Tensor,
        label_latents: torch.Tensor,
        timestep_img: Union[torch.Tensor, float, int],
        timestep_text: Union[torch.Tensor, float, int],
        data_type: Optional[Union[torch.Tensor, float, int]] = 1,
        encoder_hidden_states=None,
        cross_attention_kwargs=None,
    ):
        """
        Args:
            prompt_embeds (`torch.Tensor` of shape `(batch size, seq_len, text_dim)`):
                CLIP-embedded text representation.
            image_latents (`torch.Tensor` of shape `(batch size, latent channels, height, width)`):
                Latent image representation from the VAE encoder.
            image_embeds (`torch.Tensor` of shape `(batch size, 1, clip_img_dim)`):
                CLIP-embedded image representation (unsqueezed in the first dimension).
            label_latents (`torch.Tensor` of shape `(batch size, seq_len, label_dim)`):
                Task-specific-model-embedded label representation.
            timestep_text (`torch.long` or `float` or `int`):
                Current denoising step for the label.
            timestep_img (`torch.long` or `float` or `int`):
                Current denoising step for the image.
            data_type: (`torch.int` or `float` or `int`, *optional*, defaults to `1`):
                Only used in UniDiffuser-v1-style models. Can be either `1`, to use weights trained on nonpublic data,
                or `0` otherwise.
            encoder_hidden_states ( `torch.LongTensor` of shape `(batch size, encoder_hidden_states dim)`, *optional*):
                Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
                self-attention.
            cross_attention_kwargs (*optional*):
                Keyword arguments to supply to the cross attention layers, if used.


        Returns:
            `tuple`: Returns relevant parts of the model's noise prediction: the first element of the tuple is tbe VAE
            image embedding, the second element is the CLIP image embedding, and the third element is label embedding.
        """
        batch_size = image_latents.shape[0]

        # 1. Input
        # 1.1. Map inputs to shape (B, N, inner_dim)
        text_hidden_states = self.text_in(self.pre_text(prompt_embeds))
        vae_hidden_states = self.vae_img_in(image_latents)
        clip_hidden_states = self.clip_img_in(image_embeds)
        label_hidden_states = self.vae_label_in(label_latents)

        num_text_tokens, num_img_tokens = text_hidden_states.size(1), vae_hidden_states.size(1)

        # 1.2. Encode image timesteps to single token (B, 1, inner_dim)
        if not torch.is_tensor(timestep_img):
            timestep_img = torch.tensor([timestep_img], dtype=torch.long, device=vae_hidden_states.device)

        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
        timestep_img = timestep_img * torch.ones(batch_size, dtype=timestep_img.dtype, device=timestep_img.device)